the-other-dude/winbox-worker/cmd/worker/main.go

package main

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/the-other-dude/winbox-worker/internal/session"
)

func envInt(key string, def int) int {
	if v := os.Getenv(key); v != "" {
		if n, err := strconv.Atoi(v); err == nil {
			return n
		}
	}
	return def
}

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

func main() {
	slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
		Level: slog.LevelInfo,
	})))

	cfg := session.Config{
		MaxSessions: envInt("MAX_CONCURRENT_SESSIONS", 10),
		DisplayMin:  100,
		DisplayMax:  119,
		WSPortMin:   10100,
		WSPortMax:   10119,
		IdleTimeout: envInt("IDLE_TIMEOUT_SECONDS", 600),
		MaxLifetime: envInt("MAX_LIFETIME_SECONDS", 7200),
		WinBoxPath:  envStr("WINBOX_PATH", "/opt/winbox/WinBox"),
		BindAddr:    envStr("BIND_ADDR", "0.0.0.0"),
	}

	mgr := session.NewManager(cfg)
	mgr.CleanupOrphans()

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

	mux := http.NewServeMux()

	mux.HandleFunc("POST /sessions", func(w http.ResponseWriter, r *http.Request) {
		var req session.CreateRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeJSON(w, http.StatusBadRequest, session.ErrorResponse{Error: "invalid request body"})
			return
		}

		if !mgr.HasCapacity() {
			writeJSON(w, http.StatusServiceUnavailable, session.ErrorResponse{
				Error:       "capacity",
				MaxSessions: cfg.MaxSessions,
			})
			return
		}

		resp, err := mgr.CreateSession(req)
		req.Username = ""
		req.Password = ""

		if err != nil {
			slog.Error("create session failed", "err", err)
			if strings.Contains(err.Error(), "capacity") {
				writeJSON(w, http.StatusServiceUnavailable, session.ErrorResponse{
					Error:       "capacity",
					MaxSessions: cfg.MaxSessions,
				})
				return
			}
			writeJSON(w, http.StatusInternalServerError, session.ErrorResponse{Error: "launch failed"})
			return
		}

		writeJSON(w, http.StatusCreated, resp)
	})

	mux.HandleFunc("DELETE /sessions/{id}", func(w http.ResponseWriter, r *http.Request) {
		id := r.PathValue("id")
		if err := mgr.TerminateSession(id); err != nil {
			writeJSON(w, http.StatusInternalServerError, session.ErrorResponse{Error: err.Error()})
			return
		}
		writeJSON(w, http.StatusOK, map[string]string{"status": "terminated"})
	})

	mux.HandleFunc("GET /sessions/{id}", func(w http.ResponseWriter, r *http.Request) {
		id := r.PathValue("id")
		resp, err := mgr.GetSession(id)
		if err != nil {
			writeJSON(w, http.StatusNotFound, session.ErrorResponse{Error: "not found"})
			return
		}
		writeJSON(w, http.StatusOK, resp)
	})

	mux.HandleFunc("GET /sessions", func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, http.StatusOK, mgr.ListSessions())
	})

	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, http.StatusOK, map[string]any{
			"status":    "ok",
			"sessions":  mgr.SessionCount(),
			"capacity":  cfg.MaxSessions,
			"available": cfg.MaxSessions - mgr.SessionCount(),
		})
	})

	handler := provenanceMiddleware(mux)

	listenAddr := envStr("LISTEN_ADDR", ":9090")
	srv := &http.Server{
		Addr:         listenAddr,
		Handler:      handler,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 30 * time.Second,
	}

	go func() {
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
		<-sigCh
		slog.Info("shutting down worker")
		cancel()

		for _, s := range mgr.ListSessions() {
			mgr.TerminateSession(s.WorkerSessionID)
		}

		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer shutdownCancel()
		srv.Shutdown(shutdownCtx)
	}()

	slog.Info("winbox-worker starting", "addr", listenAddr, "max_sessions", cfg.MaxSessions)
	if err := srv.ListenAndServe(); err != http.ErrServerClosed {
		slog.Error("server error", "err", err)
		os.Exit(1)
	}
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(v)
}

func provenanceMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		svc := r.Header.Get("X-Internal-Service")
		if svc == "" && !strings.HasPrefix(r.URL.Path, "/healthz") {
			slog.Warn("request missing X-Internal-Service header", "path", r.URL.Path, "remote", r.RemoteAddr)
		}
		next.ServeHTTP(w, r)
	})
}