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

package httpserver

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"

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

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",
				Error: &otopb.ProtocolError{
					Code:    "invalid_registration_request",
					Message: "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
		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 == "" {
			writeResponse(w, http.StatusBadRequest, runnerDisconnectToJSON(false, "missing runner id"))
			return
		}

		success := registry.Disconnect(runnerID)
		if !success {
			writeResponse(w, http.StatusNotFound, runnerDisconnectToJSON(false, "runner not found"))
			return
		}

		writeResponse(w, http.StatusOK, runnerDisconnectToJSON(true, ""))
	}
}

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
		}

		writeResponse(w, http.StatusOK, runnerRecordToJSON(record))
	}
}

func shellEscape(s string) string {
	return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"
}

func handleRunnerBootstrapCommand(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.BootstrapCommandRequest
		if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
			writeResponse(w, http.StatusBadRequest, errorToJSON("invalid bootstrap command request"))
			return
		}

		runnerID := strings.TrimSpace(request.GetRunnerId())
		token := strings.TrimSpace(request.GetEnrollmentToken())
		if runnerID == "" || token == "" {
			writeResponse(w, http.StatusBadRequest, errorToJSON("missing runner id or enrollment token"))
			return
		}

		scheme := "http"
		if r.TLS != nil {
			scheme = "https"
		}
		serverURL := scheme + "://" + r.Host

		u, err := url.Parse(serverURL)
		if err != nil || u.Host == "" || u.Scheme == "" {
			writeResponse(w, http.StatusBadRequest, errorToJSON("invalid server URL"))
			return
		}

		// Host validation to prevent command injection and host header spoofing
		for _, char := range u.Host {
			if !((char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') || (char >= '0' && char <= '9') ||
				char == '.' || char == '-' || char == ':' || char == '[' || char == ']') {
				writeResponse(w, http.StatusBadRequest, errorToJSON("invalid characters in server Host"))
				return
			}
		}

		releaseBaseURL := os.Getenv("OTO_RUNNER_RELEASE_BASE_URL")
		if releaseBaseURL == "" {
			if scheme == "https" {
				releaseBaseURL = serverURL + "/releases"
			} else {
				writeResponse(w, http.StatusBadRequest, errorToJSON("release base URL must use HTTPS (set OTO_RUNNER_RELEASE_BASE_URL environment variable)"))
				return
			}
		} else {
			if !strings.HasPrefix(releaseBaseURL, "https://") {
				writeResponse(w, http.StatusBadRequest, errorToJSON("OTO_RUNNER_RELEASE_BASE_URL must use https:// scheme"))
				return
			}
		}

		escapedScriptURL := shellEscape(serverURL + "/bootstrap/oto-agent.sh")
		escapedServerURL := shellEscape(serverURL)
		escapedRunnerID := shellEscape(runnerID)
		escapedToken := shellEscape(token)
		escapedReleaseURL := shellEscape(releaseBaseURL)

		bootstrapCmd := fmt.Sprintf(
			"curl -fsSL %s | bash -s -- --server-url %s --agent-id %s --enrollment-token %s --release-base-url %s",
			escapedScriptURL, escapedServerURL, escapedRunnerID, escapedToken, escapedReleaseURL,
		)

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(&otopb.BootstrapCommandResponse{
			BootstrapCommand: bootstrapCmd,
		})
	}
}

func handleServeBootstrapScript() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
			return
		}
		paths := []string{
			"../../apps/runner/assets/script/shell/oto_agent_bootstrap.sh",
			"../../../../apps/runner/assets/script/shell/oto_agent_bootstrap.sh",
			"apps/runner/assets/script/shell/oto_agent_bootstrap.sh",
		}
		var content []byte
		var err error
		for _, p := range paths {
			content, err = os.ReadFile(p)
			if err == nil {
				break
			}
		}
		if err != nil {
			http.Error(w, "Bootstrap script not found: "+err.Error(), http.StatusNotFound)
			return
		}
		w.Header().Set("Content-Type", "application/x-sh")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(content)
	}
}

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)
}