feat: 基本的websocket echo服务

interval/ws/handler.go

@@ -0,0 +1,75 @@
+package ws
+
+import (
+	"context"
+	"io"
+	"log"
+	"net/http"
+	"time"
+
+	"github.com/coder/websocket"
+)
+
+func Handler(ctx context.Context, h *Hub) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
+			OriginPatterns: []string{"*"},
+		})
+
+		log.Println("New WebSocket connection from", r.RemoteAddr, "at", time.Now().Format(time.RFC3339))
+
+		if err != nil {
+			log.Println("WebSocket accept error:", err)
+			return
+		}
+
+		c := NewClient(r.RemoteAddr, conn, ctx)
+		log.Println("Client", r.RemoteAddr, "connected.")
+		h.RegisterClient(c)
+
+		go echo(c, h)
+		go heartbeat(c)
+	}
+}
+
+func echo(c *Client, h *Hub) {
+	defer func() {
+		if c.Conn != nil {
+			log.Println("Closing WebSocket connection")
+			h.UnregisterClient(c)
+			c.Cancel()
+			_ = c.Conn.Close(websocket.StatusNormalClosure, "echo finished")
+		}
+	}()
+
+	for {
+		typ, r, err := c.Conn.Reader(c.Ctx)
+
+		if err != nil {
+			if websocket.CloseStatus(err) == websocket.StatusNormalClosure {
+				log.Println("WebSocket connection closed normally")
+			} else {
+				log.Println("WebSocket reader error:", err)
+			}
+
+			return
+		}
+
+		w, err := c.Conn.Writer(c.Ctx, typ)
+		if err != nil {
+			log.Println("WebSocket writer error:", err)
+			return
+		}
+
+		_, err = io.Copy(w, r)
+		if err != nil {
+			log.Println("WebSocket copy error:", err)
+			return
+		}
+
+		if err = w.Close(); err != nil {
+			log.Println("WebSocket writer close error:", err)
+			return
+		}
+	}
+}