first commit

2026-02-18 22:01:54 +01:00
commit 5464e553b3
35 changed files with 5432 additions and 0 deletions
@@ -0,0 +1,138 @@
+package relay
+
+import (
+	"encoding/json"
+	"log"
+	"net/http"
+	"time"
+
+	"github.com/duffy/usb-server/internal/protocol"
+	"github.com/gorilla/websocket"
+)
+
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  64 * 1024,
+	WriteBufferSize: 64 * 1024,
+	CheckOrigin: func(r *http.Request) bool {
+		return true // relay accepts all origins
+	},
+}
+
+// Server is the WebSocket relay server
+type Server struct {
+	hub  *Hub
+	addr string
+}
+
+// NewServer creates a new relay server
+func NewServer(addr string) *Server {
+	return &Server{
+		hub:  NewHub(),
+		addr: addr,
+	}
+}
+
+// Run starts the relay server
+func (s *Server) Run() error {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/ws", s.handleWebSocket)
+	mux.HandleFunc("/health", s.handleHealth)
+
+	log.Printf("[relay] starting on %s", s.addr)
+	return http.ListenAndServe(s.addr, mux)
+}
+
+func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusOK)
+	json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
+}
+
+func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Printf("[relay] upgrade error: %v", err)
+		return
+	}
+	defer conn.Close()
+
+	// Set read limits and deadlines
+	conn.SetReadLimit(1024 * 1024) // 1MB max message
+	conn.SetReadDeadline(time.Now().Add(60 * time.Second))
+	conn.SetPongHandler(func(string) error {
+		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
+		return nil
+	})
+
+	// Wait for registration message
+	_, msgData, err := conn.ReadMessage()
+	if err != nil {
+		log.Printf("[relay] read error during registration: %v", err)
+		return
+	}
+
+	var reg protocol.Register
+	if err := json.Unmarshal(msgData, &reg); err != nil || reg.Type != protocol.MsgRegister {
+		log.Printf("[relay] invalid registration message")
+		conn.WriteJSON(&protocol.ErrorMsg{Type: protocol.MsgError, Message: "invalid registration"})
+		return
+	}
+
+	if reg.Hash == "" || reg.ClientID == "" || (reg.Mode != protocol.ModeShare && reg.Mode != protocol.ModeUse) {
+		conn.WriteJSON(&protocol.ErrorMsg{Type: protocol.MsgError, Message: "missing required fields"})
+		return
+	}
+
+	client := &Client{
+		ID:   reg.ClientID,
+		Hash: reg.Hash,
+		Mode: reg.Mode,
+		Name: reg.Name,
+		Conn: conn,
+		Send: make(chan []byte, 256),
+	}
+
+	s.hub.Register(client)
+	defer s.hub.Unregister(client)
+
+	// Start ping ticker
+	done := make(chan struct{})
+	go func() {
+		ticker := time.NewTicker(30 * time.Second)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ticker.C:
+				client.mu.Lock()
+				err := conn.WriteMessage(websocket.PingMessage, nil)
+				client.mu.Unlock()
+				if err != nil {
+					return
+				}
+			case <-done:
+				return
+			}
+		}
+	}()
+	defer close(done)
+
+	// Read loop
+	for {
+		msgType, data, err := conn.ReadMessage()
+		if err != nil {
+			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
+				log.Printf("[relay] read error from %s: %v", client.ID, err)
+			}
+			break
+		}
+
+		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
+
+		switch msgType {
+		case websocket.TextMessage:
+			s.hub.HandleTextMessage(client, data)
+		case websocket.BinaryMessage:
+			s.hub.HandleBinaryMessage(client, data)
+		}
+	}
+}