oto/services/core/internal/httpserver/server.go

package httpserver

import (
	"context"
	"encoding/json"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/toki/oto/services/core/internal/runnerregistry"
	otopb "github.com/toki/oto/services/core/oto"
)

// Server wraps the HTTP server for the OTO Core service.
type Server struct {
	httpServer *http.Server
}

// NewServer creates a new instance of Server.
func NewServer(addr string) *Server {
	return NewServerWithRegistry(addr, runnerregistry.New())
}

// NewServerWithRegistry creates a server using an injected runner registry.
func NewServerWithRegistry(addr string, registry *runnerregistry.Registry) *Server {
	mux := http.NewServeMux()

	// Register health and readiness endpoints
	mux.HandleFunc("/healthz", handleHealthz)
	mux.HandleFunc("/readyz", handleReadyz)
	mux.HandleFunc("/api/v1/runners/register", handleRunnerRegister(registry))
	mux.HandleFunc("/api/v1/runners/{id}/heartbeat", handleRunnerHeartbeat(registry))
	mux.HandleFunc("/api/v1/runners/{id}/disconnect", handleRunnerDisconnect(registry))
	mux.HandleFunc("/api/v1/runners/{id}", handleGetRunner(registry))

	return &Server{
		httpServer: &http.Server{
			Addr:    addr,
			Handler: mux,
		},
	}
}

// Start starts the HTTP server.
func (s *Server) Start() error {
	return s.httpServer.ListenAndServe()
}

// StartListener starts the HTTP server using a custom net.Listener.
// Useful for testing with dynamic ports.
func (s *Server) StartListener(ln net.Listener) error {
	return s.httpServer.Serve(ln)
}

// Shutdown gracefully shuts down the server.
func (s *Server) Shutdown(ctx context.Context) error {
	return s.httpServer.Shutdown(ctx)
}

func handleHealthz(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
		return
	}
	w.WriteHeader(http.StatusOK)
	w.Write([]byte("OK"))
}

func handleReadyz(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
		return
	}
	w.WriteHeader(http.StatusOK)
	w.Write([]byte("OK"))
}

func handleRunnerRegister(registry *runnerregistry.Registry) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
			return
		}

		var request otopb.RegisterRunnerRequest
		if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
			writeRunnerRegisterResponse(w, http.StatusBadRequest, &otopb.RegisterRunnerResponse{
				Accepted:     false,
				RejectReason: "invalid registration request",
			})
			return
		}

		response := registry.Register(&request)
		writeRunnerRegisterResponse(w, http.StatusOK, response)
	}
}

func handleRunnerHeartbeat(registry *runnerregistry.Registry) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
			return
		}

		runnerID := r.PathValue("id")
		if runnerID == "" {
			parts := strings.Split(r.URL.Path, "/")
			if len(parts) >= 5 {
				runnerID = parts[4]
			}
		}

		if runnerID == "" {
			writeHeartbeatResponse(w, http.StatusBadRequest, &otopb.HeartbeatResponse{
				Success:      false,
				ErrorMessage: "missing runner id",
			})
			return
		}

		var req otopb.HeartbeatRequest
		// Decode request if body is not empty
		if r.ContentLength > 0 || r.Body != nil {
			if err := json.NewDecoder(r.Body).Decode(&req); err != nil && err.Error() != "EOF" {
				writeHeartbeatResponse(w, http.StatusBadRequest, &otopb.HeartbeatResponse{
					Success:      false,
					ErrorMessage: "invalid heartbeat request",
				})
				return
			}
		}

		if req.RunnerId == "" {
			req.RunnerId = runnerID
		} else if req.RunnerId != runnerID {
			writeHeartbeatResponse(w, http.StatusBadRequest, &otopb.HeartbeatResponse{
				Success:      false,
				ErrorMessage: "runner id mismatch between path and body",
			})
			return
		}

		response := registry.Heartbeat(req.RunnerId, req.Status)
		if !response.Success {
			status := http.StatusNotFound
			if response.ErrorMessage != "unknown runner" {
				status = http.StatusBadRequest
			}
			writeHeartbeatResponse(w, status, response)
			return
		}

		writeHeartbeatResponse(w, http.StatusOK, response)
	}
}

func handleRunnerDisconnect(registry *runnerregistry.Registry) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
			return
		}

		runnerID := r.PathValue("id")
		if runnerID == "" {
			parts := strings.Split(r.URL.Path, "/")
			if len(parts) >= 5 {
				runnerID = parts[4]
			}
		}

		if runnerID == "" {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusBadRequest)
			_ = json.NewEncoder(w).Encode(map[string]interface{}{
				"success": false,
				"error":   "missing runner id",
			})
			return
		}

		success := registry.Disconnect(runnerID)
		w.Header().Set("Content-Type", "application/json")
		if !success {
			w.WriteHeader(http.StatusNotFound)
			_ = json.NewEncoder(w).Encode(map[string]interface{}{
				"success": false,
				"error":   "runner not found",
			})
			return
		}

		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(map[string]interface{}{
			"success": true,
		})
	}
}

func writeRunnerRegisterResponse(w http.ResponseWriter, status int, response *otopb.RegisterRunnerResponse) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(response)
}

func writeHeartbeatResponse(w http.ResponseWriter, status int, response *otopb.HeartbeatResponse) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(response)
}

func handleGetRunner(registry *runnerregistry.Registry) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
			return
		}

		runnerID := r.PathValue("id")
		if runnerID == "" {
			parts := strings.Split(r.URL.Path, "/")
			if len(parts) >= 5 {
				runnerID = parts[4]
			}
		}

		if runnerID == "" {
			http.Error(w, "missing runner id", http.StatusBadRequest)
			return
		}

		record, ok := registry.Snapshot(runnerID)
		if !ok {
			http.Error(w, "runner not found", http.StatusNotFound)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(map[string]interface{}{
			"runner_id":          record.RunnerID,
			"alias":              record.Alias,
			"protocol_version":   record.ProtocolVersion,
			"status":             record.Status,
			"accepted_at":        record.AcceptedAt.Format(time.RFC3339),
			"first_heartbeat_at": record.FirstHeartbeatAt.Format(time.RFC3339),
			"last_heartbeat_at":  record.LastHeartbeatAt.Format(time.RFC3339),
			"failure_reason":     record.FailureReason,
		})
	}
}