dwarf-go/dwarfctl/internal/transport/client.go

package transport

import (
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	pb "github.com/antitbone/dwarfctl/proto"
	"google.golang.org/protobuf/proto"
)

const (
	WsMajorVersion = 2
	WsMinorVersion = 3
	WsPort         = 9900
)

// MsgType mirrors WsMessageType: request=0, response=1, notification=2, reply=3.
type MsgType uint32

const (
	MsgRequest      MsgType = 0
	MsgResponse     MsgType = 1
	MsgNotification MsgType = 2
	MsgReply        MsgType = 3
)

// Client is a WebSocket client for the DWARF telescope control plane.
type Client struct {
	conn      *websocket.Conn
	clientID  string
	deviceID  uint32
	mu        sync.Mutex
	pending   map[uint32]chan *pb.WsPacket
	notifyChs []chan *pb.WsPacket
	done      chan struct{}
	Debug     bool
}

// NewClient creates a client with the given client_id and device_id.
func NewClient(clientID string, deviceID uint32) *Client {
	return &Client{
		clientID: clientID,
		deviceID: deviceID,
		pending:  make(map[uint32]chan *pb.WsPacket),
		done:     make(chan struct{}),
	}
}

// Connect opens the WebSocket to the telescope.
func (c *Client) Connect(ip string) error {
	url := fmt.Sprintf("ws://%s:%d/?client_id=%s", ip, WsPort, c.clientID)
	dialer := websocket.Dialer{
		HandshakeTimeout: 10 * time.Second,
	}
	conn, _, err := dialer.Dial(url, nil)
	if err != nil {
		return fmt.Errorf("ws dial %s: %w", url, err)
	}
	c.conn = conn
	// Keep the connection alive — the telescope closes idle connections
	// after ~30s. The PongHandler also resets the read deadline on each pong.
	conn.SetPingHandler(func(string) error {
		conn.WriteControl(websocket.PongMessage, nil, time.Now().Add(5*time.Second))
		return nil
	})
	go c.readLoop()
	go c.heartbeatLoop()
	return nil
}

// heartbeatLoop sends periodic ping frames to keep the WebSocket alive.
// The telescope drops idle connections after ~30s.
func (c *Client) heartbeatLoop() {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-c.done:
			return
		case <-ticker.C:
			c.mu.Lock()
			if c.conn == nil {
				c.mu.Unlock()
				return
			}
			err := c.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(5*time.Second))
			c.mu.Unlock()
			if err != nil && c.Debug {
				log.Printf("[WS] heartbeat ping error: %v", err)
			}
		}
	}
}

// Close shuts down the connection.
func (c *Client) Close() error {
	close(c.done)
	if c.conn != nil {
		return c.conn.Close()
	}
	return nil
}

// IsConnected returns true if the WebSocket is open.
func (c *Client) IsConnected() bool {
	return c.conn != nil
}

// readLoop continuously reads WsPacket frames and dispatches them.
func (c *Client) readLoop() {
	for {
		select {
		case <-c.done:
			return
		default:
		}
		_, data, err := c.conn.ReadMessage()
		if err != nil {
			if c.Debug {
				log.Printf("[WS] read error: %v", err)
			}
			return
		}
		pkt := &pb.WsPacket{}
		if err := proto.Unmarshal(data, pkt); err != nil {
			if c.Debug {
				log.Printf("[WS] unmarshal error: %v (raw %d bytes)", err, len(data))
			}
			continue
		}
		if c.Debug {
			log.Printf("[WS] recv cmd=%d module=%d type=%d data=%d bytes",
				pkt.GetCmd(), pkt.GetModuleId(), pkt.GetType(), len(pkt.GetData()))
		}
		// Dispatch based on type. Many telescopes send responses with
		// type=0 (default proto3 value) instead of type=1. So we try
		// response dispatch for type 0, 1, and 3.
		switch MsgType(pkt.GetType()) {
		case MsgNotification:
			c.dispatchNotification(pkt)
			// Also try response dispatch — some notifications double as acks
			c.dispatchResponse(pkt)
		default:
			// Covers type=0 (unset), type=1 (response), type=3 (reply)
			c.dispatchResponse(pkt)
			c.dispatchNotification(pkt)
		}
	}
}

// dispatchResponse delivers a response to the pending request waiter.
func (c *Client) dispatchResponse(pkt *pb.WsPacket) {
	c.mu.Lock()
	ch, ok := c.pending[pkt.GetCmd()]
	if ok {
		delete(c.pending, pkt.GetCmd())
	}
	c.mu.Unlock()
	if ok {
		ch <- pkt
	}
}

// dispatchNotification fans out to all subscribers.
func (c *Client) dispatchNotification(pkt *pb.WsPacket) {
	c.mu.Lock()
	chs := make([]chan *pb.WsPacket, len(c.notifyChs))
	copy(chs, c.notifyChs)
	c.mu.Unlock()
	for _, ch := range chs {
		select {
		case ch <- pkt:
		default:
		}
	}
}

// SubscribeNotifications returns a channel for receiving NOTIFY packets.
func (c *Client) SubscribeNotifications() chan *pb.WsPacket {
	ch := make(chan *pb.WsPacket, 64)
	c.mu.Lock()
	c.notifyChs = append(c.notifyChs, ch)
	c.mu.Unlock()
	return ch
}

// Send sends a raw command with the given module_id and cmd, carrying
// the serialized inner proto message as data. It returns the response packet.
func (c *Client) Send(moduleID uint32, cmd uint32, data []byte, timeout time.Duration) (*pb.WsPacket, error) {
	pkt := &pb.WsPacket{
		MajorVersion: WsMajorVersion,
		MinorVersion: WsMinorVersion,
		DeviceId:     c.deviceID,
		ModuleId:     moduleID,
		Cmd:          cmd,
		Type:         uint32(MsgRequest),
		Data:         data,
		ClientId:     c.clientID,
	}
	raw, err := proto.Marshal(pkt)
	if err != nil {
		return nil, fmt.Errorf("marshal WsPacket: %w", err)
	}

	respCh := make(chan *pb.WsPacket, 1)
	c.mu.Lock()
	c.pending[cmd] = respCh
	c.mu.Unlock()

	if err := c.conn.WriteMessage(websocket.BinaryMessage, raw); err != nil {
		c.mu.Lock()
		delete(c.pending, cmd)
		c.mu.Unlock()
		return nil, fmt.Errorf("ws write: %w", err)
	}

	select {
	case resp := <-respCh:
		return resp, nil
	case <-time.After(timeout):
		c.mu.Lock()
		delete(c.pending, cmd)
		c.mu.Unlock()
		return nil, fmt.Errorf("timeout waiting for response to cmd %d", cmd)
	case <-c.done:
		return nil, fmt.Errorf("connection closed")
	}
}

// SendNotify sends a command without waiting for a response (fire-and-forget).
func (c *Client) SendNotify(moduleID uint32, cmd uint32, data []byte) error {
	pkt := &pb.WsPacket{
		MajorVersion: WsMajorVersion,
		MinorVersion: WsMinorVersion,
		DeviceId:     c.deviceID,
		ModuleId:     moduleID,
		Cmd:          cmd,
		Type:         uint32(MsgRequest),
		Data:         data,
		ClientId:     c.clientID,
	}
	raw, err := proto.Marshal(pkt)
	if err != nil {
		return err
	}
	return c.conn.WriteMessage(websocket.BinaryMessage, raw)
}