第二章节: 识别WebSocket客户端用户

前言

上一章节中我们创建了一个最简单的WebSocket服务端，但是在服务端的视角下这些客户端连接都是一样的，当用户A向用户B发送消息时，实际对应服务端的逻辑是FindConnB然后往ConnB里面WriteMessage，所以第一步就需要将User和Conn一一绑定然后保存起来，那么第一步就是当有新连接时，识别连接的用户。因为WebSocket实际也是Http协议，所以和Http服务一样通过HttpHeader中携带的Token来解析出当前连接的用户.

用户实体

在go-im的src目录下创建model文件夹，在model文件夹下创建models.go

go 复制代码

package model

import (
    "encoding/base64"
    "encoding/json"
    "net/http"
)

type User struct {
    ID   int64
    Name string
}

// Resolve User From HttpHeader -> X-TOKEN
func ResolveUser(request *http.Request) *User {
    // 从HttpHeader中获取X-TOKEN
    token := request.Header.Get("X-TOKEN")
    if len(token) == 0 {
        return nil
    }
    // 先进行Base64解密
    decoded_token, err := base64.StdEncoding.DecodeString(token)
    if err != nil {
        return nil
    }
    u := &User{}
    // 然后JSON反序列化为User对象
    err = json.Unmarshal(decoded_token, u)
    if err != nil {
        return nil
    }
    return u
}

// 将用户序列化为X-Token
func (u User) GenToken() string {
    jsonBytes, err := json.Marshal(u)
    if err != nil {
        return ""
    }
    return base64.StdEncoding.EncodeToString(jsonBytes)
}

改造Server

修改server.go代码，识别连接中的用户, 如果未解析出用户，将断开连接

go 复制代码

package main

import (
    "log"
    "net/http"

    "aoki.com/go-im/src/model"
    "github.com/gin-gonic/gin"
    "github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
    // Allow Cross Origin
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

// 处理WS
func ws(c *gin.Context) {
    // 升级为WS协议 
    conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
    if err != nil {
        panic(err)
    }
    defer conn.Close()
    user := model.ResolveUser(c.Request)
    if user == nil {
        log.Println("ResolveUser failed...")
        return
    }
    log.Printf("%s connected... \n", user.Name)
    // 循环读取数据
    for {
        mt, message, err := conn.ReadMessage()
        if err != nil {
            log.Printf("Read Message Failed... \n", err)
            break
        }
        // 读到什么往WS客户端发什么
        log.Printf("Received Message %s... \n", message)
        err = conn.WriteMessage(mt, message)
        if err != nil {
            log.Printf("Write Message Failed... \n", err)
            break
        }
    }
}

// 启动Websocket server
func main() {
    server := gin.Default()
    server.GET("/ws", ws)
    server.Run("localhost:8848")
}

改造Client

此时Server端做了处理，如果没有从HttpHeader中解析出用户，将会主动断开连接，所以此时客户端就必须携带对应的令牌才能连接

修改client.go

go 复制代码

package main

import (
	"bufio"
	"fmt"
	"log"
	"net/http"
	"os"

	"aoki.com/go-im/src/model"
	"github.com/gorilla/websocket"
)

func main() {
    uri := "ws://localhost:8848/ws"
    u := model.User{
        ID:   1001,
        Name: "AOKI",
    }
    header := http.Header{
        "X-TOKEN": []string{u.GenToken()},
    }
    conn, _, err := websocket.DefaultDialer.Dial(uri, header)
    if err != nil {
        log.Fatal("dial:", err)
    }
    defer conn.Close()

    go func() {
        for {
            _, message, err := conn.ReadMessage()
            if err != nil {
                fmt.Println("Read WS message failed:", err)
                return
            }
            log.Println("Received: ", string(message))
        }
    }()

    scanner := bufio.NewScanner(os.Stdin)
    for scanner.Scan() {
        line := scanner.Text()
        if err := conn.WriteMessage(websocket.TextMessage, []byte(line)); err != nil {
            log.Println("write failed:", err)
        }
    }
    if err := scanner.Err(); err != nil {
        log.Fatal(err)
    }
}

启动交互

此时我们分别在两个Terminal启动server和client，可以看到server端的日志

表示通过HttpHeader的Token识别用户连接方式没有问题

下一章节我们在Server端创建连接管理器，将用户和其对应的连接进行绑定，并管理所有的客户端连接