这里说一下 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()
}









