用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

Conn (此处指 wserver.Conn) 为 websocket.Conn 的包装。

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

            }

        }

    }

}