// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved. // Use of this source code is governed by a BSD-style // license that can be found in the LICENSE file. package websockets import ( "net/http" "time" "github.com/axllent/mailpit/config" "github.com/axllent/mailpit/internal/logger" "github.com/gorilla/websocket" ) 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{' '} // MessageHub global MessageHub *Hub ) var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: func(r *http.Request) bool { return true }, // allow multi-domain EnableCompression: true, // experimental compression } // 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 } // 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)) _ = c.conn.WriteMessage(websocket.PingMessage, []byte{}) } } } // ServeWs handles websocket requests from the peer. func ServeWs(hub *Hub, w http.ResponseWriter, r *http.Request) { if config.UIAuthFile != "" { if config.UIAuthFile != "" { user, pass, ok := r.BasicAuth() if !ok { basicAuthResponse(w) return } if !config.UIAuth.Match(user, pass) { basicAuthResponse(w) return } } } conn, err := upgrader.Upgrade(w, r, nil) if err != nil { logger.Log().Error(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() } // BasicAuthResponse returns an basic auth response to the browser func basicAuthResponse(w http.ResponseWriter) { w.Header().Set("WWW-Authenticate", `Basic realm="Login"`) w.WriteHeader(http.StatusUnauthorized) _, _ = w.Write([]byte("Unauthorised.\n")) }