利用 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)
   }
  }
 }
}