利用 Go 语言编写一个简单的 WebSocket 推送服务

2020-01-28 13:04:04刘景俊

这里说一下 Upgrader *websocket.Upgrader,这是 gorilla/websocket 包的对象,它用来升级 HTTP 请求。

如果一个结构体参数过多,通常不建议直接初始化,而是使用它提供的 New 方法。这里是:


// NewServer creates a new Server.
func NewServer(addr string) *Server {
 return &Server{
  Addr:  addr,
  WSPath: serverDefaultWSPath,
  PushPath: serverDefaultPushPath,
 }
}

这也是 Go 语言对外提供初始化方法的一种常见用法。

然后 Server 使用 ListenAndServe 方法启动并监听端口,与 http 包的使用类似:


// ListenAndServe listens on the TCP network address and handle websocket
// request.
func (s *Server) ListenAndServe() error {
 b := &binder{
  userID2EventConnMap: make(map[string]*[]eventConn),
  connID2UserIDMap: make(map[string]string),
 }
 // websocket request handler
 wh := websocketHandler{
  upgrader: defaultUpgrader,
  binder: b,
 }
 if s.Upgrader != nil {
  wh.upgrader = s.Upgrader
 }
 if s.AuthToken != nil {
  wh.calcUserIDFunc = s.AuthToken
 }
 s.wh = &wh
 http.Handle(s.WSPath, s.wh)
 // push request handler
 ph := pushHandler{
  binder: b,
 }
 if s.PushAuth != nil {
  ph.authFunc = s.PushAuth
 }
 s.ph = &ph
 http.Handle(s.PushPath, s.ph)
 return http.ListenAndServe(s.Addr, nil)
}

这里我们生成了两个 Handler,分别为 websocketHandler 和 pushHandler。websocketHandler 负责与浏览器建立连接并传输数据,而 pushHandler 则处理推送端的请求。可以看到,这里两个 Handler 都封装了一个 binder 对象。这个 binder 用于维护 token <-> userID <-> Conn 的关系:


// binder is defined to store the relation of userID and eventConn
type binder struct {
 mu sync.RWMutex
 // map stores key: userID and value of related slice of eventConn
 userID2EventConnMap map[string]*[]eventConn

 // map stores key: connID and value: userID
 connID2UserIDMap map[string]string
}

websocketHandler

具体看一下 websocketHandler 的实现。


// websocketHandler defines to handle websocket upgrade request.
type websocketHandler struct {
 // upgrader is used to upgrade request.
 upgrader *websocket.Upgrader

 // binder stores relations about websocket connection and userID.
 binder *binder

 // calcUserIDFunc defines to calculate userID by token. The userID will
 // be equal to token if this function is nil.
 calcUserIDFunc func(token string) (userID string, ok bool)
}

很简单的结构。websocketHandler 实现了 http.Handler 接口:


// First try to upgrade connection to websocket. If success, connection will
// be kept until client send close message or server drop them.
func (wh *websocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 wsConn, err := wh.upgrader.Upgrade(w, r, nil)
 if err != nil {
  return
 }
 defer wsConn.Close()
 // handle Websocket request
 conn := NewConn(wsConn)
 conn.AfterReadFunc = func(messageType int, r io.Reader) {
  var rm RegisterMessage
  decoder := json.NewDecoder(r)
  if err := decoder.Decode(&rm); err != nil {
   return
  }
  // calculate userID by token
  userID := rm.Token
  if wh.calcUserIDFunc != nil {
   uID, ok := wh.calcUserIDFunc(rm.Token)
   if !ok {
    return
   }
   userID = uID
  }
  // bind
  wh.binder.Bind(userID, rm.Event, conn)
 }
 conn.BeforeCloseFunc = func() {
  // unbind
  wh.binder.Unbind(conn)
 }
 conn.Listen()
}