WebsocketDemo

2022-03-24

项目里有想法想加个聊天功能，因此稍微写个websocket小demo熟悉一下大致的过程

后端代码

使用go语言的gin框架"github.com/gin-gonic/gin"和第三方websocket库"github.com/gorilla/websocket"编写，代码如下

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"net/http"
)

/* connect为全局的map,用于存储连接,实际开发中可换成redis缓存 */
var connect = make(map[string]*websocket.Conn)
/* 测试环境没有数据库,因此用data这个map模拟数据库,即处理当对方不在线时的数据缓存 */
var data = make(map[string][]byte)

/* 创建一个websocket的Upgrader实例,具体参数可以看api文档,ReadBufferSize和WriteBufferSize若不填值则默认4096,CheckOrigin函数用于验证请求域,此处直接返回true,真实开发时需验证请求域,若合法则返回true */
var up = websocket.Upgrader{
	ReadBufferSize: 1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

/* gin框架的handlerFunc */
func ws(c *gin.Context)  {
  /* 模拟数据,因ws默认发起GET请求,因此从GET请求携带的参数中获得所需参数 */
  /* 此处的from和to分别模拟用户本人和需要通信的另一名用户,也可以将to参数设置为某个会话群的id,以实现多人聊天而非单对单聊天 */
	from := c.Query("from")
	to := c.Query("to")

  /* 将普通的http请求Upgrade成websocket请求,因为websocket请求之前也需要先进行握手,websocket和http均承载于TCP连接之上 */
	webs, err := up.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		println(err)
		return
	}

  /* 将当前连接存入全局map,以便会话另一端用户发起会话时找到当前连接 */
	key := fmt.Sprintf("from:%s,to:%s", from, to)
	connect[key] = webs
	/* 另一端会话的key */
	rkey := fmt.Sprintf("from:%s,to:%s", to, from)
	/* go routine执行函数 */
	go func() {
    /* 定义当连接即将断开前的处理函数,关闭websocket,删除map中的当前连接 */
		defer func() {
			webs.Close()
			println(fmt.Sprintf("链接%s断开", key))
			delete(connect, key)
		}()
		/* for循环持续监听websocket */
		for {
      /* 如果data中存在另一端发送的数据(即真实开发中,数据库内部存有对当前用户发送的数据),则在连接建立后立即将数据发送给当前用户,并删除数据或标记数据为已发送 */
			if v, ok := data[rkey]; ok {
				err = webs.WriteMessage(1, v)
				if err != nil {
					println(err)
					return
				}
				delete(data, rkey)
			}
      /* 监听用户发送来的数据,可以使用ReadJSON()和WriteJSON()来接收和发送复杂数据 */
			t, ms, err := webs.ReadMessage()
			if err != nil {
				println(err)
				return
			}

			println(fmt.Sprintf("type:%d, message:%s", t,ms))
			/* 数据回显,也可仅发送'服务器成功接收信息'的通知 */
			err = webs.WriteMessage(t, ms)
			if err != nil {
				println(err)
				return
			}
			/* 如果另一端的连接存在,说明双方均在线,向另一位用户发送当前用户所发送的信息;否则,将数据写入数据库并标记为未发送 */
			if webs2, ok := connect[rkey]; ok {
				err = webs2.WriteMessage(t, ms)
				if err != nil {
					println(err)
					return
				}
			} else {
				data[key] = ms
			}
		}
	}()
}

func init()  {
	r := gin.Default()
  /* 正常监听GET请求即可,在handlerFunc内部会将GET请求提升为websocket */
	r.GET("/ws", ws)
	r.Run(":8000")
}

前端

由于项目是小程序，所以demo写的也是小程序，普通Javascript大同小异

Page({
  data: {
    input: "",
    from: "",
    to: "",
    message: [],
  },
  /* 用于存储websocket实例,由于该实例不涉及页面交互,出于性能考虑,不放入data中 */
  ws: null,

  onLoad() {
    const that = this;
    /* 这里使用设备类型来假性设定from和to,实际开发时有用户唯一id可用 */
    wx.getSystemInfo({
      success (res) {
        console.log(res.platform)
        if(res.platform === "devtools") {
          that.setData({
            from: "devtools",
            to: "phone"
          })
        } else {
          that.setData({
            from: "phone",
            to: "devtools"
          })
        }
      }
    })
		
    /* 微信小程序的封装,相当于new Websocket(url),这里使用了同一台路由器下的两个设备(开发者工具和微信小程序)进行模拟调试 */
    this.ws = wx.connectSocket({
      url: 'ws://192.168.1.2:8000/ws?from=' + this.data.from + '&to=' + this.data.to
    })
		/* 添加接受信息的回调函数,将数据放入数组进行渲染 */
    this.ws.onMessage((res) => {
      console.log(res);
      let message = this.data.message;
      message.push(res.data);
      this.setData({
        message,
      })
    })
  },
	/* 在页面的Unload事件中关闭websocket连接 */
  onUnload() {
    this.ws.close();
  },
	/* 发送数据事件,input+button,按下button时触发此函数,input的输入数据双向绑定到this.data.input上 */
  send() {
    console.log(this.data.input);
    this.ws.send({
      data: this.data.input
    });
    this.setData({
      input: "",
    })
  }
})

写在最后

由于只是个小demo，并没有连接池、连接控制、发送心跳包判断是否在线等一些重要的处理，写真实项目时会注意的😃