第四章节：WebSocket实现用户一对一聊天

前言

上一章节中我们实现了在WebSocket Server端管理用户和WebSocket连接的绑定和释放，本章节我们尝试实现两个client端之间的消息收发。

消息实体

首先我们需要定义一个统一的消息模型

在models.go中添加Message结构体

type Message struct {
    ID         int64
    Type       string
    FromUserID int64
    ToUserID   int64
    Data       interface{}
}

消息发送服务

在conn.go中定义接口Sender

type Sender interface {
    Send(msg model.Message)
}

编写服务实现LocalSender

type LocalSender struct {
}

func (sender *LocalSender) Send(msg model.Message) error {
    // 获取对方的WebSocket连接
    conn := GetConnManager().FindConn(msg.ToUserID)
    if conn == nil {
        return fmt.Errorf("%v does not online", msg.ToUserID)
    }
    // 发送JSON消息
    return conn.WriteJSON(msg)
}

WebSocket Server端

修改server.go

package main

import (
    "log"
    "net/http"

    "aoki.com/go-im/src/conn"
    "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
    },
}

var ConnManager *conn.ConnManager = conn.GetConnManager()
var Sender = conn.LocalSender{}

func ws(c *gin.Context) {
    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
    }
    ConnManager.AddConn(user.ID, conn)
    log.Printf("%s connected... \n", user.Name)
    defer OnDisconnect(*user)
    for {
        v := &model.Message{}
        err = conn.ReadJSON(v)
        if err != nil {
            log.Println("Read Json Message Failed", err)
        }
        err = Sender.Send(*v)
        if err != nil {
            log.Println("Send Message Failed", err)
        }
    }
}

func OnDisconnect(user model.User) {
    log.Printf("%v disconnect...\n", user.Name)
    ConnManager.DelConn(user.ID)
}

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

WebSocket Client端

修改client.go

package main

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

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

func main() {
    uri := "ws://localhost:8848/ws"
    // 从命令行获取登录用户ID和Name
    args := os.Args[1:]
    userID, er := strconv.ParseInt(args[0], 10, 64)
    if er != nil {
        log.Fatal("Parse UserId Failed:", er)
    }
    u := model.User{
        ID:   userID,
        Name: args[1],
    }
    header := http.Header{
        "X-TOKEN": []string{u.GenToken()},
    }
    conn, _, err := websocket.DefaultDialer.Dial(uri, header)
    if err != nil {
        log.Fatal("dial failed:", err)
    }
    defer conn.Close()

    go func() {
        for {
            // 读取JSON消息
            m := &model.Message{}
            err = conn.ReadJSON(m)
            if err != nil {
                log.Println("Read WS message failed:", err)
                return
            }
            log.Printf("Received Message %v From %v \n", m.Data, m.FromUserID)
        }
    }()

    scanner := bufio.NewScanner(os.Stdin)
    for scanner.Scan() {
        // 从控制台读取内容，前一段为收信人的用户ID，后一段为消息内容
        str := strings.Split(scanner.Text(), " ")
        ToUserID, er := strconv.ParseInt(str[0], 10, 64)
        if er != nil {
            log.Println("cannot recognize ToUserId ", err)
            break
        }
        m := model.Message{
            FromUserID: u.ID,
            ToUserID:   ToUserID,
            Type:       "Text",
            Data:       str[1],
        }
        if err := conn.WriteJSON(m); err != nil {
            log.Println("Send WS message failed:", err)
        }
    }
    if err := scanner.Err(); err != nil {
        log.Fatal(err)
    }
}

模拟交互

在Terminal中启动server.go, 然后新开3个Terminal，分别启动3个client进程

go run src/main/client.go 1001 AOKI
go run src/main/client.go 1002 Tony
go run src/main/client.go 1003 Frank

模拟3个用户登录IM服务，然后分别使用3个Terminal向其他用户发送消息，并在Terminal中确认收到的消息

发送方和接收方消息都能对上，消息都能正常接收、发送

小结

本章节我们实现了使用WebSocket让用户间能够正常收发消息，但是还存在一些问题，比如当对方不在线时，消息将会丢失，这是不合常理的，正常应该等对方上线时会立即收到这条未读消息，所以在下一章节我们探讨将消息持久化并做已读未读状态标记。