// readPump pumps messages from the websocket connection to the hub. func (c *connection) readPump() { defer func() { h.unregister <- c c.ws.Close() }() for { // Use deadline to detect dead or stuck clients. c.ws.SetReadDeadline(time.Now().Add(readWait)) op, r, err := c.ws.NextReader() if err != nil { return } if op == websocket.OpBinary { // unexpected return } if op != websocket.OpText { // ignore pongs and other control messages. continue } lr := io.LimitedReader{R: r, N: maxMessageSize + 1} message, err := ioutil.ReadAll(&lr) if err != nil { return } if lr.N <= 0 { // Message is larger than max allowed message size. c.ws.WriteControl(websocket.OpClose, websocket.FormatCloseMessage(websocket.CloseMessageTooBig, ""), time.Now().Add(time.Second)) return } h.broadcast <- message } }
// readPump pumps messages from the websocket connection to the hub. func (c *connection) readPump() { defer c.ws.Close() for { c.ws.SetReadDeadline(time.Now().Add(readWait)) op, r, err := c.ws.NextReader() if err != nil { return } if op != websocket.OpText { continue } lr := io.LimitedReader{R: r, N: maxMessageSize + 1} message, err := ioutil.ReadAll(&lr) if err != nil { return } if lr.N <= 0 { c.ws.WriteControl(websocket.OpClose, websocket.FormatCloseMessage(websocket.CloseMessageTooBig, ""), time.Now().Add(time.Second)) return } h.broadcast <- message } }
// echoReadAll echoes messages from the client by reading the entire message // with ioutil.ReadAll. func echoReadAll(w http.ResponseWriter, r *http.Request, writeMessage bool) { conn, err := websocket.Upgrade(w, r.Header, nil, 4096, 4096) if err != nil { log.Println("Upgrade:", err) http.Error(w, "Bad request", 400) return } defer conn.Close() for { op, r, err := conn.NextReader() if err != nil { if err != io.EOF { log.Println("NextReader:", err) } return } if op == websocket.OpPong { continue } if op == websocket.OpText { r = &validator{r: r} } b, err := ioutil.ReadAll(r) if err != nil { if err == errInvalidUTF8 { conn.WriteControl(websocket.OpClose, websocket.FormatCloseMessage(websocket.CloseInvalidFramePayloadData, ""), time.Time{}) } log.Println("ReadAll:", err) return } if writeMessage { err = conn.WriteMessage(op, b) if err != nil { log.Println("WriteMessage:", err) } } else { w, err := conn.NextWriter(op) if err != nil { log.Println("NextWriter:", err) return } if _, err := w.Write(b); err != nil { log.Println("Writer:", err) return } if err := w.Close(); err != nil { log.Println("Close:", err) return } } } }
// echoCopy echoes messages from the client using io.Copy. func echoCopy(w http.ResponseWriter, r *http.Request, writerOnly bool) { conn, err := websocket.Upgrade(w, r.Header, nil, 4096, 4096) if err != nil { log.Println("Upgrade:", err) http.Error(w, "Bad request", 400) return } defer conn.Close() for { op, r, err := conn.NextReader() if err != nil { if err != io.EOF { log.Println("NextReader:", err) } return } if op == websocket.OpPong { continue } if op == websocket.OpText { r = &validator{r: r} } w, err := conn.NextWriter(op) if err != nil { log.Println("NextWriter:", err) return } if op == websocket.OpText { r = &validator{r: r} } if writerOnly { _, err = io.Copy(struct{ io.Writer }{w}, r) } else { _, err = io.Copy(w, r) } if err != nil { if err == errInvalidUTF8 { conn.WriteControl(websocket.OpClose, websocket.FormatCloseMessage(websocket.CloseInvalidFramePayloadData, ""), time.Time{}) } log.Println("Copy:", err) return } err = w.Close() if err != nil { log.Println("Close:", err) return } } }
// serverWs handles webocket requests from the client. func serveSocket(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.Header, "", 1024, 1024) if err != nil { http.Error(w, err.Error(), 400) log.Println(err) return } roomid, err := strconv.Atoi(r.URL.RawQuery) if err != nil { ws.WriteControl(websocket.OpClose, websocket.FormatCloseMessage(4000, "Not a valid room id"), time.Now().Add(writeWait)) ws.Close() return } room := roomserver.GetRoom(roomid) if room == nil { ws.WriteControl(websocket.OpClose, websocket.FormatCloseMessage(4004, "Room not found"), time.Now().Add(writeWait)) ws.Close() return } c := &connection{send: make(chan []byte), receive: make(chan []byte), ws: ws, room: room} room.Register <- c go c.writePump() c.readPump() }
// Follow handles FOLLOW requests to retrieve the contents of a file and a // real-time stream of data that is appended to the file. func (h Handler) Follow(w http.ResponseWriter, r *http.Request) { path := h.resolve(r.URL.Path) h.logf("FOLLOW %s", path) // If this file isn't currently being written to, we don't need to update to // a WebSocket; we can just return the static file. if !h.isWriting(path) { h.serveFile(w, r) return } c := make(chan []byte) h.addFollower(path, r, c) defer h.removeFollower(path, r) // TODO(sqs): race conditions galore f, err := os.Open(path) if err != nil { http.Error(w, "failed to open file: "+err.Error(), http.StatusInternalServerError) return } defer f.Close() // Open WebSocket. ws, err := websocket.Upgrade(w, r.Header, nil, readBufSize, writeBufSize) if err != nil { if _, ok := err.(websocket.HandshakeError); ok { // Serve file via HTTP (not WebSocket). h.serveFile(w, r) return } h.logf("failed to upgrade to WebSocket: %s", err) return } defer ws.Close() // Send persisted file contents. for { sw, err := ws.NextWriter(websocket.OpText) if err != nil { h.logf("NextWriter for file failed: %s", err) return } n, err := io.Copy(sw, f) if err != nil { h.logf("File write to WebSocket failed: %s", err) sw.Close() return } err = sw.Close() if err != nil { h.logf("Failed to close WebSocket file writer: %s", err) return } // Finished reading file. if n == 0 { break } } // Follow new writes to file. var lastPing time.Time for { tick := time.NewTicker(50 * time.Millisecond) select { case <-tick.C: if !h.isWriting(path) { goto done } if time.Since(lastPing) > followKeepaliveInterval { ws.WriteMessage(websocket.OpPing, []byte{}) lastPing = time.Now() } case data := <-c: sw, err := ws.NextWriter(websocket.OpText) if err != nil { h.logf("NextWriter failed: %s", err) return } _, err = sw.Write(data) if err != nil { h.logf("Write to WebSocket failed: %s", err) sw.Close() return } err = sw.Close() if err != nil { h.logf("Failed to close WebSocket writer: %s", err) return } } } done: err = ws.WriteControl(websocket.OpClose, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Time{}) if err != nil { h.logf("Failed to close WebSocket: %s", err) return } err = f.Close() if err != nil { h.logf("Failed to close destination file: %s", err) return } }