// Copyright 2013 Gary Burd. All rights reserved. // Use of this source code is governed by a BSD-style // license that can be found in the LICENSE file. package main import ( "github.com/gorilla/websocket" "log" "net/http" "time" ) 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 ) // connection is an middleman between the websocket connection and the hub. type connection struct { // The websocket connection. ws *websocket.Conn // Buffered channel of outbound messages. send chan []byte } // readPump pumps messages from the websocket connection to the hub. func (c *connection) readPump() { defer func() { h.unregister <- c c.ws.Close() }() c.ws.SetReadLimit(maxMessageSize) c.ws.SetReadDeadline(time.Now().Add(pongWait)) c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil }) for { _, message, err := c.ws.ReadMessage() if err != nil { break } h.broadcast <- message } } // write writes a message with the given message type and payload. func (c *connection) write(mt int, payload []byte) error { c.ws.SetWriteDeadline(time.Now().Add(writeWait)) return c.ws.WriteMessage(mt, payload) } // writePump pumps messages from the hub to the websocket connection. func (c *connection) writePump() { ticker := time.NewTicker(pingPeriod) defer func() { ticker.Stop() c.ws.Close() }() for { select { case message, ok := <-c.send: if !ok { c.write(websocket.CloseMessage, []byte{}) return } if err := c.write(websocket.TextMessage, message); err != nil { return } case <-ticker.C: if err := c.write(websocket.PingMessage, []byte{}); err != nil { return } } } } // serverWs handles webocket requests from the peer. func serveWs(w http.ResponseWriter, r *http.Request) { if r.Method != "GET" { http.Error(w, "Method not allowed", 405) return } if r.Header.Get("Origin") != "http://"+r.Host { http.Error(w, "Origin not allowed", 403) return } ws, err := websocket.Upgrade(w, r, nil, 1024, 1024) if _, ok := err.(websocket.HandshakeError); ok { http.Error(w, "Not a websocket handshake", 400) return } else if err != nil { log.Println(err) return } c := &connection{send: make(chan []byte, 256), ws: ws} h.register <- c go c.writePump() c.readPump() }