124 lines
2.3 KiB
Go
124 lines
2.3 KiB
Go
package ws
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log"
|
|
|
|
"git.jinshen.cn/remilia/push-server/interval/protocol"
|
|
"github.com/coder/websocket"
|
|
"github.com/coder/websocket/wsjson"
|
|
)
|
|
|
|
// Client represents a connected client in the hub.
|
|
type Client struct {
|
|
ID string
|
|
Conn *websocket.Conn
|
|
SendChan chan []byte
|
|
Hub *Hub
|
|
|
|
Ctx context.Context
|
|
Cancel context.CancelFunc
|
|
|
|
inited bool
|
|
}
|
|
|
|
func NewClient(id string, conn *websocket.Conn, hub *Hub, parentCtx context.Context) *Client {
|
|
ctx, cancel := context.WithCancel(parentCtx)
|
|
|
|
return &Client{
|
|
ID: id,
|
|
Conn: conn,
|
|
SendChan: make(chan []byte, 32),
|
|
Hub: hub,
|
|
|
|
Ctx: ctx,
|
|
Cancel: cancel,
|
|
|
|
inited: false,
|
|
}
|
|
}
|
|
func (c *Client) ReadLoop() {
|
|
defer c.Close()
|
|
|
|
for {
|
|
var msg protocol.ControlMessage
|
|
|
|
err := wsjson.Read(c.Ctx, c.Conn, &msg)
|
|
if err != nil {
|
|
if websocket.CloseStatus(err) == websocket.StatusNormalClosure {
|
|
log.Println("WebSocket closed normally:", err)
|
|
} else {
|
|
log.Println("WebSocket read error:", err)
|
|
}
|
|
return
|
|
}
|
|
|
|
if err := msg.Validate(); err != nil {
|
|
_ = c.Conn.Close(websocket.StatusPolicyViolation, "invalid message")
|
|
return
|
|
}
|
|
|
|
if err := c.handleControlMessage(msg); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Client) WriteLoop() {
|
|
defer c.Close()
|
|
|
|
for {
|
|
select {
|
|
case <-c.Ctx.Done():
|
|
return
|
|
case msg, ok := <-c.SendChan:
|
|
if !ok {
|
|
return
|
|
}
|
|
err := wsjson.Write(c.Ctx, c.Conn, msg)
|
|
if err != nil {
|
|
log.Println("WebSocket write error:", err)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Client) handleControlMessage(msg protocol.ControlMessage) error {
|
|
switch msg.Type {
|
|
case protocol.MsgInit:
|
|
if c.inited {
|
|
return errors.New("already initialized")
|
|
}
|
|
|
|
log.Printf("Client %s initializing with topics: %v", c.ID, msg.Topics)
|
|
c.Hub.RegisterClient(c)
|
|
|
|
for _, t := range msg.Topics {
|
|
c.Hub.Subscribe(protocol.Subscription{
|
|
ClientID: c.ID,
|
|
Topic: protocol.Topic(t),
|
|
})
|
|
}
|
|
|
|
c.inited = true
|
|
|
|
return nil
|
|
default:
|
|
if !c.inited {
|
|
return errors.New("client not initialized")
|
|
}
|
|
log.Println("Unknown control message type:", msg.Type)
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (c *Client) Close() {
|
|
c.Cancel()
|
|
c.Hub.UnregisterClient(c)
|
|
if c.Conn != nil {
|
|
_ = c.Conn.Close(websocket.StatusNormalClosure, "client closed")
|
|
}
|
|
}
|