client2server/server/main.go

// client2server - Go server with WebSocket + Redpanda (Kafka-compatible)
// Copyright (c) 2026 Luis Rosales - MIT License
//
// Build: go build -o client2server-server .
// Run:   ./client2server-server
// Env:   REDPANDA_BROKERS=localhost:9092 TOKEN=*** PORT=3843
//
// WebSocket: ws://localhost:3843/ws
// HTTP API:  http://localhost:3843/api/{events,routers,command}, /health

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/coder/websocket"
	"github.com/google/uuid"
	"github.com/twmb/franz-go/pkg/kadm"
	"github.com/twmb/franz-go/pkg/kerr"
	"github.com/twmb/franz-go/pkg/kgo"
)

// ----------------------------------------------------------------------------
// Config
// ----------------------------------------------------------------------------

type Config struct {
	RedpandaBrokers []string
	Port            int
	Token           string
}

var cfg = Config{
	RedpandaBrokers: []string{"localhost:9092"},
	Port:            3843,
	Token:           "secret-token",
}

func loadConfigFromEnv() {
	cfg.RedpandaBrokers = strings.Split(getEnvStr("REDPANDA_BROKERS", "localhost:9092"), ",")
	cfg.Port = getEnvInt("PORT", cfg.Port)
	cfg.Token = getEnvStr("TOKEN", cfg.Token)
}

func getEnvStr(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}

func getEnvInt(key string, def int) int {
	if v := os.Getenv(key); v != "" {
		var n int
		if _, err := fmt.Sscanf(v, "%d", &n); err == nil {
			return n
		}
	}
	return def
}

// ----------------------------------------------------------------------------
// Domain types
// ----------------------------------------------------------------------------

// RouterEvent - generic event sent from a router to the server.
type RouterEvent struct {
	ID          string                 `json:"id"`
	RouterID    string                 `json:"router_id"`
	Hostname    string                 `json:"hostname,omitempty"`
	EventType   string                 `json:"event_type"`
	Timestamp   time.Time              `json:"timestamp"`
	Payload     map[string]interface{} `json:"payload"`
	ReceivedAt  time.Time              `json:"received_at"`
	Connection  string                 `json:"connection"` // "websocket" | "http"
}

// RouterCommand - command sent from server to a router.
type RouterCommand struct {
	ID       string            `json:"id"`
	RouterID string            `json:"router_id"`
	Command  string            `json:"command"`
	Args     map[string]string `json:"args,omitempty"`
	SentAt   time.Time         `json:"sent_at"`
}

// CommandResult - reply from a router after running a command.
type CommandResult struct {
	Success bool   `json:"success"`
	Output  string `json:"output"`
	Error   string `json:"error"`
}

// ----------------------------------------------------------------------------
// Router registry
// ----------------------------------------------------------------------------

type Router struct {
	ID        string
	LastSeen  time.Time
	Conn      *websocket.Conn
	Connected bool
	writeMu   sync.Mutex // serialise writes to the WS connection
}

var (
	routersMu      sync.RWMutex
	routers        = make(map[string]*Router)
	routerQueuesMu sync.Mutex
	routerQueues   = make(map[string][]RouterCommand) // queued while offline
	pendingMu      sync.Mutex
	pendingCmds    = make(map[string]chan CommandResult) // command_id -> result chan
	executedMu     sync.Mutex
	executedCmds   = make(map[string]time.Time) // idempotency: cmd_id -> last run
)

const (
	commandTimeout   = 30 * time.Second
	idempotencyTTL   = 5 * time.Minute
	wsWriteTimeout   = 10 * time.Second
	publishTimeout   = 3 * time.Second
	offlineThreshold = 60 * time.Second
)

// ----------------------------------------------------------------------------
// Redpanda (Kafka) client
// ----------------------------------------------------------------------------

var kcl *kgo.Client

func initRedpanda(ctx context.Context) error {
	cl, err := kgo.NewClient(
		kgo.SeedBrokers(cfg.RedpandaBrokers...),
		kgo.ClientID("client2server"),
		kgo.ProducerLinger(5*time.Millisecond),
		kgo.ProducerBatchCompression(kgo.SnappyCompression()),
	)
	if err != nil {
		return fmt.Errorf("kafka client: %w", err)
	}
	kcl = cl

	// Best-effort topic creation. Redpanda has auto-create enabled in dev, so
	// this is just to make sure they exist with sane defaults.
	adm := kadm.NewClient(cl)
	topics := []string{"router-events", "router-commands"}
	resp, err := adm.CreateTopics(ctx, 1, 1, nil, topics...)
	if err != nil {
		log.Printf("create topics admin call failed (non-fatal if auto-create is on): %v", err)
		return nil
	}
	for _, ct := range resp {
		if ct.Err != nil && !errors.Is(ct.Err, kerr.TopicAlreadyExists) {
			log.Printf("topic %s: %v", ct.Topic, ct.Err)
		}
	}
	return nil
}

func publish(ctx context.Context, topic string, key string, value any) error {
	data, err := json.Marshal(value)
	if err != nil {
		return fmt.Errorf("marshal: %w", err)
	}
	rec := &kgo.Record{Topic: topic, Key: []byte(key), Value: data}
	// Bound how long the HTTP handler can wait for the broker.
	pctx, cancel := context.WithTimeout(ctx, publishTimeout)
	defer cancel()
	res := kcl.ProduceSync(pctx, rec)
	if err := res.FirstErr(); err != nil {
		return err
	}
	return nil
}

// ----------------------------------------------------------------------------
// WebSocket handler
// ----------------------------------------------------------------------------

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	// Auth: token in query string (?token=...) or Authorization header
	token := r.URL.Query().Get("token")
	if token == "" {
		if h := r.Header.Get("Authorization"); strings.HasPrefix(h, "Bearer ") {
			token = strings.TrimPrefix(h, "Bearer ")
		}
	}
	if token != cfg.Token {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
		// Disable permessage-deflate to keep things simple on minimal routers
		CompressionMode: websocket.CompressionDisabled,
	})
	if err != nil {
		log.Printf("ws accept: %v", err)
		return
	}
	defer conn.Close(websocket.StatusNormalClosure, "bye")

	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()

	// First message must be a registration event
	firstMsg, err := readRouterMessage(ctx, conn)
	if err != nil {
		log.Printf("ws read reg: %v", err)
		return
	}
	var reg RouterEvent
	if err := json.Unmarshal(firstMsg, &reg); err != nil {
		log.Printf("ws reg parse: %v", err)
		return
	}

	routerID := reg.RouterID
	if routerID == "" {
		routerID = r.RemoteAddr
	}

	routersMu.Lock()
	routers[routerID] = &Router{
		ID:        routerID,
		LastSeen:  time.Now(),
		Conn:      conn,
		Connected: true,
	}
	routersMu.Unlock()

	log.Printf("router connected: %s (from %s)", routerID, r.RemoteAddr)
	flushQueuedCommands(ctx, routerID)

	// Message loop
	for {
		raw, err := readRouterMessage(ctx, conn)
		if err != nil {
			break
		}

		var ev RouterEvent
		if err := json.Unmarshal(raw, &ev); err != nil {
			log.Printf("[%s] bad event json: %v", routerID, err)
			continue
		}

		// Server-assigned fields
		ev.ID = uuid.New().String()
		ev.ReceivedAt = time.Now()
		ev.Connection = "websocket"
		if ev.Timestamp.IsZero() {
			ev.Timestamp = ev.ReceivedAt
		}

		// Command result handling
		if ev.EventType == "command_result" {
			if cid, _ := ev.Payload["command_id"].(string); cid != "" {
				deliverCommandResult(cid, ev.Payload)
				executedMu.Lock()
				executedCmds[cid] = time.Now()
				executedMu.Unlock()
			}
		}

		if err := publish(ctx, "router-events", routerID, ev); err != nil {
			log.Printf("publish event: %v", err)
		}
		log.Printf("[%s] %s", routerID, ev.EventType)

		routersMu.Lock()
		if r, ok := routers[routerID]; ok {
			r.LastSeen = time.Now()
		}
		routersMu.Unlock()
	}

	routersMu.Lock()
	if r, ok := routers[routerID]; ok {
		r.Connected = false
	}
	routersMu.Unlock()
	log.Printf("router disconnected: %s", routerID)
}

func readRouterMessage(ctx context.Context, conn *websocket.Conn) ([]byte, error) {
	// coder/websocket: Read returns a Message
	_, data, err := conn.Read(ctx)
	return data, err
}

// writeJSON serialises writes to the connection.
func (r *Router) writeJSON(ctx context.Context, v any) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	wctx, cancel := context.WithTimeout(ctx, wsWriteTimeout)
	defer cancel()
	r.writeMu.Lock()
	defer r.writeMu.Unlock()
	return r.Conn.Write(wctx, websocket.MessageText, data)
}

func deliverCommandResult(cmdID string, payload map[string]interface{}) {
	pendingMu.Lock()
	ch, ok := pendingCmds[cmdID]
	if ok {
		delete(pendingCmds, cmdID)
	}
	pendingMu.Unlock()
	if !ok {
		return
	}
	res := CommandResult{
		Success: payload["success"] == true,
	}
	if s, ok := payload["output"].(string); ok {
		res.Output = s
	}
	if s, ok := payload["error"].(string); ok {
		res.Error = s
	}
	select {
	case ch <- res:
	default:
	}
}

func flushQueuedCommands(ctx context.Context, routerID string) {
	routerQueuesMu.Lock()
	queue := routerQueues[routerID]
	delete(routerQueues, routerID)
	routerQueuesMu.Unlock()
	if len(queue) == 0 {
		return
	}
	log.Printf("flushing %d queued commands to %s", len(queue), routerID)
	for _, cmd := range queue {
		cmd.SentAt = time.Now()
		pendingMu.Lock()
		pendingCmds[cmd.ID] = make(chan CommandResult, 1)
		pendingMu.Unlock()
		if err := publish(ctx, "router-commands", routerID, cmd); err != nil {
			log.Printf("queue flush publish: %v", err)
		}
	}
}

// ----------------------------------------------------------------------------
// HTTP handlers
// ----------------------------------------------------------------------------

func handleHTTPEvent(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	if !authorised(r) {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	var ev RouterEvent
	if err := json.NewDecoder(r.Body).Decode(&ev); err != nil {
		http.Error(w, "invalid json", http.StatusBadRequest)
		return
	}
	ev.ID = uuid.New().String()
	ev.ReceivedAt = time.Now()
	ev.Connection = "http"
	if ev.Timestamp.IsZero() {
		ev.Timestamp = ev.ReceivedAt
	}
	if err := publish(r.Context(), "router-events", ev.RouterID, ev); err != nil {
		log.Printf("publish event: %v", err)
		http.Error(w, "publish failed", http.StatusBadGateway)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]string{
		"event_id":  ev.ID,
		"router_id": ev.RouterID,
		"status":    "accepted",
	})
}

func handleRouters(w http.ResponseWriter, r *http.Request) {
	type entry struct {
		ID       string    `json:"id"`
		LastSeen time.Time `json:"last_seen"`
		Online   bool      `json:"online"`
		Queued   int       `json:"queued_commands"`
	}
	routersMu.RLock()
	list := make([]entry, 0, len(routers))
	for id, rt := range routers {
		list = append(list, entry{
			ID:       id,
			LastSeen: rt.LastSeen,
			Online:   time.Since(rt.LastSeen) < offlineThreshold,
			Queued:   len(routerQueues[id]),
		})
	}
	routersMu.RUnlock()
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]interface{}{"routers": list})
}

func handleCommand(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	if !authorised(r) {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	var cmd RouterCommand
	if err := json.NewDecoder(r.Body).Decode(&cmd); err != nil {
		http.Error(w, "invalid json", http.StatusBadRequest)
		return
	}
	routerID := r.URL.Query().Get("router_id")
	if routerID == "" {
		routerID = cmd.RouterID
	}
	if routerID == "" {
		http.Error(w, "router_id required", http.StatusBadRequest)
		return
	}

	// Idempotency
	xecID := cmd.ID
	if xecID == "" {
		xecID = uuid.New().String()
	}
	executedMu.Lock()
	if last, exists := executedCmds[xecID]; exists && time.Since(last) < idempotencyTTL {
		executedMu.Unlock()
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]interface{}{
			"idempotent_reject":   true,
			"existing_command_id": xecID,
			"message":             "command already executed within TTL",
		})
		return
	}
	delete(executedCmds, xecID)
	executedMu.Unlock()

	cmd.ID = xecID
	cmd.RouterID = routerID
	cmd.SentAt = time.Now()

	routersMu.RLock()
	router, online := routers[routerID]
	routersMu.RUnlock()

	if !online || router == nil || !router.Connected {
		routerQueuesMu.Lock()
		routerQueues[routerID] = append(routerQueues[routerID], cmd)
		depth := len(routerQueues[routerID])
		routerQueuesMu.Unlock()
		log.Printf("queued cmd %s for %s (depth=%d)", cmd.Command, routerID, depth)
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]interface{}{
			"command_id":  cmd.ID,
			"status":      "queued",
			"queued_for":  routerID,
			"queue_depth": depth,
		})
		return
	}

	// Online: register waiter, publish, wait for result
	resultCh := make(chan CommandResult, 1)
	pendingMu.Lock()
	pendingCmds[cmd.ID] = resultCh
	pendingMu.Unlock()

	if err := publish(r.Context(), "router-commands", routerID, cmd); err != nil {
		pendingMu.Lock()
		delete(pendingCmds, cmd.ID)
		pendingMu.Unlock()
		log.Printf("publish command: %v", err)
		http.Error(w, "publish failed", http.StatusBadGateway)
		return
	}

	log.Printf("cmd %s -> %s (waiting)", cmd.Command, routerID)
	w.Header().Set("Content-Type", "application/json")

	select {
	case res := <-resultCh:
		_ = json.NewEncoder(w).Encode(map[string]interface{}{
			"command_id": cmd.ID,
			"status":     "completed",
			"success":    res.Success,
			"output":     res.Output,
			"error":      res.Error,
		})
	case <-time.After(commandTimeout):
		pendingMu.Lock()
		delete(pendingCmds, cmd.ID)
		pendingMu.Unlock()
		_ = json.NewEncoder(w).Encode(map[string]string{
			"command_id": cmd.ID,
			"status":     "timeout",
			"error":      "router did not respond",
		})
	}
}

func handleHealth(w http.ResponseWriter, r *http.Request) {
	routersMu.RLock()
	onlineCount := 0
	for _, rt := range routers {
		if time.Since(rt.LastSeen) < offlineThreshold {
			onlineCount++
		}
	}
	total := len(routers)
	routersMu.RUnlock()
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]interface{}{
		"status":         "ok",
		"routers":        total,
		"routers_online": onlineCount,
		"redpanda":       cfg.RedpandaBrokers,
	})
}

func authorised(r *http.Request) bool {
	h := r.Header.Get("Authorization")
	if h == cfg.Token { // legacy: raw token
		return true
	}
	if strings.HasPrefix(h, "Bearer ") && strings.TrimPrefix(h, "Bearer ") == cfg.Token {
		return true
	}
	return false
}

// ----------------------------------------------------------------------------
// Background cleanup
// ----------------------------------------------------------------------------

func startJanitor(ctx context.Context) {
	go func() {
		t := time.NewTicker(time.Minute)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				now := time.Now()
				executedMu.Lock()
				for id, ts := range executedCmds {
					if now.Sub(ts) > idempotencyTTL {
						delete(executedCmds, id)
					}
				}
				executedMu.Unlock()
			}
		}
	}()
}

// ----------------------------------------------------------------------------
// main
// ----------------------------------------------------------------------------

func main() {
	loadConfigFromEnv()
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.Printf("=== client2server Go Server ===")
	log.Printf("port=%d brokers=%v", cfg.Port, cfg.RedpandaBrokers)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	if err := initRedpanda(ctx); err != nil {
		log.Printf("redpanda init failed (continuing): %v", err)
	} else {
		defer kcl.Close()
	}
	startJanitor(ctx)

	mux := http.NewServeMux()
	mux.HandleFunc("/health", handleHealth)
	mux.HandleFunc("/api/events", handleHTTPEvent)
	mux.HandleFunc("/api/routers", handleRouters)
	mux.HandleFunc("/api/command", handleCommand)
	mux.HandleFunc("/ws", handleWebSocket)
	mux.HandleFunc("/", handleHealth)

	srv := &http.Server{
		Addr:         fmt.Sprintf(":%d", cfg.Port),
		Handler:      mux,
		ReadTimeout:  0, // WS connections are long-lived
		WriteTimeout: 0,
	}

	// Graceful shutdown
	go func() {
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
		<-sigCh
		log.Println("shutting down...")
		routersMu.RLock()
		for _, r := range routers {
			if r.Conn != nil {
				r.Conn.Close(websocket.StatusNormalClosure, "server shutdown")
			}
		}
		routersMu.RUnlock()
		shutdownCtx, c := context.WithTimeout(context.Background(), 5*time.Second)
		defer c()
		_ = srv.Shutdown(shutdownCtx)
		os.Exit(0)
	}()

	log.Printf("server ready on :%d", cfg.Port)
	if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
		log.Fatal(err)
	}
}