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