76 lines
1.5 KiB
Go
76 lines
1.5 KiB
Go
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
|
|
}
|
|
}
|
|
}
|