package http import ( "net/http" "strings" "sync" "time" "github.com/gorilla/websocket" "go.uber.org/zap" ) var wsCon = websocket.Upgrader{} const ( // wsMaxMessageSize caps a single inbound websocket message so one large // frame cannot exhaust process memory. wsMaxMessageSize = 1 << 20 // 1 MiB // wsReadTimeout is how long the server waits for the next client message or // pong before closing an idle connection. wsReadTimeout = 60 * time.Second // wsWriteTimeout bounds a single write so a slow reader cannot block forever. wsWriteTimeout = 10 * time.Second // wsPingInterval is how often the server pings the client to keep the // connection alive and detect dead peers; it must be shorter than wsReadTimeout. wsPingInterval = 30 * time.Second ) // EchoWS godoc // @Summary Echo over websockets // @Description echos content via websockets // @Tags HTTP API // @Accept json // @Produce json // @Router /ws/echo [post] // @Success 202 {object} http.MapResponse // Test: go run ./cmd/podcli/* ws localhost:9898/ws/echo func (s *Server) echoWsHandler(w http.ResponseWriter, r *http.Request) { c, err := wsCon.Upgrade(w, r, nil) if err != nil { s.logger.Warn("websocket upgrade error", zap.Error(err)) return } // Bound per-message size and idle time; refresh the read deadline whenever // the client responds to a ping so live connections stay open. c.SetReadLimit(wsMaxMessageSize) _ = c.SetReadDeadline(time.Now().Add(wsReadTimeout)) c.SetPongHandler(func(string) error { return c.SetReadDeadline(time.Now().Add(wsReadTimeout)) }) var wg sync.WaitGroup wg.Add(1) defer c.Close() done := make(chan struct{}) defer close(done) in := make(chan interface{}) go s.writeWs(c, in) go s.sendHostWs(c, in, done, &wg) go func() { defer close(in) wg.Wait() }() for { _, message, err := c.ReadMessage() if err != nil { if !strings.Contains(err.Error(), "close") { s.logger.Warn("websocket read error", zap.Error(err)) } break } var response = struct { Time time.Time `json:"ts"` Message string `json:"msg"` }{ Time: time.Now(), Message: string(message), } in <- response } } func (s *Server) sendHostWs(ws *websocket.Conn, in chan interface{}, done chan struct{}, wg *sync.WaitGroup) { ticker := time.NewTicker(5 * time.Second) defer ticker.Stop() for { select { case <-ticker.C: var status = struct { Time time.Time `json:"ts"` Host string `json:"server"` }{ Time: time.Now(), Host: s.config.Hostname, } in <- status case <-done: s.logger.Debug("websocket exit") wg.Done() return } } } func (s *Server) writeWs(ws *websocket.Conn, in chan interface{}) { ping := time.NewTicker(wsPingInterval) defer ping.Stop() for { select { case msg := <-in: _ = ws.SetWriteDeadline(time.Now().Add(wsWriteTimeout)) if err := ws.WriteJSON(msg); err != nil { if !strings.Contains(err.Error(), "close") { s.logger.Warn("websocket write error", zap.Error(err)) } return } case <-ping.C: _ = ws.SetWriteDeadline(time.Now().Add(wsWriteTimeout)) if err := ws.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } }