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

// Server defines parameters for running websocket server.
type Server struct {
 // Address for server to listen on
 Addr string

 // Path for websocket request, default "/ws".
 WSPath string

 // Path for push message, default "/push".
 PushPath string

 // Upgrader is for upgrade connection to websocket connection using
 // "github.com/gorilla/websocket".
 //
 // If Upgrader is nil, default upgrader will be used. Default upgrader is
 // set ReadBufferSize and WriteBufferSize to 1024, and CheckOrigin always
 // returns true.
 Upgrader *websocket.Upgrader

 // Check token if it's valid and return userID. If token is valid, userID
 // must be returned and ok should be true. Otherwise ok should be false.
 AuthToken func(token string) (userID string, ok bool)

 // Authorize push request. Message will be sent if it returns true,
 // otherwise the request will be discarded. Default nil and push request
 // will always be accepted.
 PushAuth func(r *http.Request) bool

 wh *websocketHandler
 ph *pushHandler
}

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

// 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)
}

// 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 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)
}

// 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()
}

// Authorize if needed. Then decode the request and push message to each
// realted websocket connection.
func (s *pushHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 if r.Method != http.MethodPost {
  w.WriteHeader(http.StatusMethodNotAllowed)
  return
 }
 // authorize
 if s.authFunc != nil {
  if ok := s.authFunc(r); !ok {
   w.WriteHeader(http.StatusUnauthorized)
   return
  }
 }
 // read request
 var pm PushMessage
 decoder := json.NewDecoder(r.Body)
 if err := decoder.Decode(pm); err != nil {
  w.WriteHeader(http.StatusBadRequest)
  w.Write([]byte(ErrRequestIllegal.Error()))
  return
 }
 // validate the data
 if pm.UserID == "" || pm.Event == "" || pm.Message == "" {
  w.WriteHeader(http.StatusBadRequest)
  w.Write([]byte(ErrRequestIllegal.Error()))
  return
 }
 cnt, err := s.push(pm.UserID, pm.Event, pm.Message)
 if err != nil {
  w.WriteHeader(http.StatusInternalServerError)
  w.Write([]byte(err.Error()))
  return
 }
 result := strings.NewReader(fmt.Sprintf("message sent to %d clients", cnt))
 io.Copy(w, result)
}

// Conn wraps websocket.Conn with Conn. It defines to listen and read
// data from Conn.
type Conn struct {
 Conn *websocket.Conn
 AfterReadFunc func(messageType int, r io.Reader)
 BeforeCloseFunc func()

 once sync.Once
 id  string
 stopCh chan struct{}
}

// Listen listens for receive data from websocket connection. It blocks
// until websocket connection is closed.
func (c *Conn) Listen() {
 c.Conn.SetCloseHandler(func(code int, text string) error {
  if c.BeforeCloseFunc != nil {
   c.BeforeCloseFunc()
  }
  if err := c.Close(); err != nil {
   log.Println(err)
  }
  message := websocket.FormatCloseMessage(code, "")
  c.Conn.WriteControl(websocket.CloseMessage, message, time.Now().Add(time.Second))
  return nil
 })
 // Keeps reading from Conn util get error.
ReadLoop:
 for {
  select {
  case -c.stopCh:
   break ReadLoop
  default:
   messageType, r, err := c.Conn.NextReader()
   if err != nil {
    // TODO: handle read error maybe
    break ReadLoop
   }
   if c.AfterReadFunc != nil {
    c.AfterReadFunc(messageType, r)
   }
  }
 }
}