package main import ( "encoding/json" "log" "net/http" "strings" "time" "github.com/gorilla/websocket" "golang.org/x/time/rate" ) // code adapted from the gorilla websocket chat demo code const ( // Time allowed to write a message to the peer. writeWait = 10 * time.Second // Time allowed to read the next pong message from the peer. pongWait = 60 * time.Second // Send pings to peer with this period. Must be less than pongWait. pingPeriod = (pongWait * 9) / 10 // Maximum message size allowed from peer. maxMessageSize = 512 ) var ( newline = []byte{'\n'} space = []byte{' '} ) var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, } // Client is a middleman between the websocket connection and the hub. type Client struct { hub *Hub // The websocket connection. conn *websocket.Conn // Buffered channel of outbound messages. send chan []byte username string auth bool limiter rate.Limiter } type Auth struct { Type string `json:"type"` Username string `json:"username"` Key string `json:"key"` } // readPump pumps messages from the websocket connection to the hub. // // The application runs readPump in a per-connection goroutine. The application // ensures that there is at most one reader on a connection by executing all // reads from this goroutine. func (c *Client) readPump() { defer func() { c.hub.unregister <- c c.conn.Close() }() c.conn.SetReadLimit(maxMessageSize) c.conn.SetReadDeadline(time.Now().Add(pongWait)) c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil }) c.limiter = *rate.NewLimiter(1, 5) for { _, message, err := c.conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { log.Printf("error: %v", err) } break } if !c.limiter.Allow() { c.send <- []byte("{\"type\":\"chaterror\",\"error\":\"You are sending messages too quickly\"}") continue } var v Auth err = json.Unmarshal(message, &v) if err == nil { // check for authorization messages if v.Type == "auth" { // bind the WS client to the username (useful) c.username = strings.ToLower(v.Username) // if auth key matches session, they are authorized to bet/chat etc if DBGetUserByUsername(c.username).Session == v.Key { c.auth = true } // if the user has bet (and then reloaded the page) send them their state if c.hub.allUsers[c.username] != "" { c.send <- []byte("{\"type\":\"state\",\"value\":true,\"bet\":\"" + c.hub.allUsers[c.username] + "\"}") } continue } } // only authorized clients are allowed to push to the Hub if c.auth { c.hub.broadcast <- message } } } // writePump pumps messages from the hub to the websocket connection. // // A goroutine running writePump is started for each connection. The // application ensures that there is at most one writer to a connection by // executing all writes from this goroutine. func (c *Client) writePump() { ticker := time.NewTicker(pingPeriod) defer func() { ticker.Stop() c.conn.Close() }() for { select { case message, ok := <-c.send: c.conn.SetWriteDeadline(time.Now().Add(writeWait)) if !ok { // The hub closed the channel. c.conn.WriteMessage(websocket.CloseMessage, []byte{}) return } w, err := c.conn.NextWriter(websocket.TextMessage) if err != nil { return } w.Write(message) // Add queued chat messages to the current websocket message. n := len(c.send) for i := 0; i < n; i++ { w.Write(newline) w.Write(<-c.send) } if err := w.Close(); err != nil { return } case <-ticker.C: c.conn.SetWriteDeadline(time.Now().Add(writeWait)) if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } } // serveWs handles websocket requests from the peer. func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) { upgrader.CheckOrigin = func(r *http.Request) bool { return true } conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Println(err) return } client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)} client.hub.register <- client // Allow collection of memory referenced by the caller by doing all work in // new goroutines. go client.writePump() go client.readPump() }