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

package httpserver

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

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

func TestHandleHealthz(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
	rr := httptest.NewRecorder()

	handleHealthz(rr, req)

	if status := rr.Code; status != http.StatusOK {
		t.Errorf("handleHealthz returned wrong status code: got %v want %v", status, http.StatusOK)
	}

	expected := "OK"
	if body := rr.Body.String(); body != expected {
		t.Errorf("handleHealthz returned unexpected body: got %v want %v", body, expected)
	}
}

func TestHandleHealthz_MethodNotAllowed(t *testing.T) {
	req := httptest.NewRequest(http.MethodPost, "/healthz", nil)
	rr := httptest.NewRecorder()

	handleHealthz(rr, req)

	if status := rr.Code; status != http.StatusMethodNotAllowed {
		t.Errorf("handleHealthz for POST returned wrong status code: got %v want %v", status, http.StatusMethodNotAllowed)
	}
}

func TestHandleReadyz(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/readyz", nil)
	rr := httptest.NewRecorder()

	handleReadyz(rr, req)

	if status := rr.Code; status != http.StatusOK {
		t.Errorf("handleReadyz returned wrong status code: got %v want %v", status, http.StatusOK)
	}

	expected := "OK"
	if body := rr.Body.String(); body != expected {
		t.Errorf("handleReadyz returned unexpected body: got %v want %v", body, expected)
	}
}

func TestHandleReadyz_MethodNotAllowed(t *testing.T) {
	req := httptest.NewRequest(http.MethodPost, "/readyz", nil)
	rr := httptest.NewRecorder()

	handleReadyz(rr, req)

	if status := rr.Code; status != http.StatusMethodNotAllowed {
		t.Errorf("handleReadyz for POST returned wrong status code: got %v want %v", status, http.StatusMethodNotAllowed)
	}
}

func TestServerMux(t *testing.T) {
	registry := runnerregistry.New()
	server := NewServerWithRegistry(":0", registry)
	mux := server.httpServer.Handler.(*http.ServeMux)

	// Test healthz routing
	reqHealth := httptest.NewRequest(http.MethodGet, "/healthz", nil)
	rrHealth := httptest.NewRecorder()
	mux.ServeHTTP(rrHealth, reqHealth)
	if rrHealth.Code != http.StatusOK {
		t.Errorf("mux /healthz failed: got status %v", rrHealth.Code)
	}

	// Test readyz routing
	reqReady := httptest.NewRequest(http.MethodGet, "/readyz", nil)
	rrReady := httptest.NewRecorder()
	mux.ServeHTTP(rrReady, reqReady)
	if rrReady.Code != http.StatusOK {
		t.Errorf("mux /readyz failed: got status %v", rrReady.Code)
	}

	// Test fallback/not found routing
	reqNotFound := httptest.NewRequest(http.MethodGet, "/unknown", nil)
	rrNotFound := httptest.NewRecorder()
	mux.ServeHTTP(rrNotFound, reqNotFound)
	if rrNotFound.Code != http.StatusNotFound {
		t.Errorf("mux /unknown routing status code: got %v want %v", rrNotFound.Code, http.StatusNotFound)
	}
}

func TestHandleRunnerRegisterAcceptsAndStoresRunner(t *testing.T) {
	registry := runnerregistry.New()
	body := bytes.NewBufferString(`{
		"enrollment_token":"token-123",
		"runner_id":"runner-123",
		"alias":"linux-build",
		"protocol_version":"oto.runner.v1",
		"capability":{"name":"oto-runner","version":"1.0.0"},
		"command_catalog":{"command_types":["Shell","Git"]}
	}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/register", body)
	rr := httptest.NewRecorder()

	handleRunnerRegister(registry)(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v; body=%s", rr.Code, http.StatusOK, rr.Body.String())
	}
	var response otopb.RegisterRunnerResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if !response.GetAccepted() {
		t.Fatalf("accepted=false, reason=%q", response.GetRejectReason())
	}
	record, ok := registry.Snapshot("runner-123")
	if !ok {
		t.Fatal("registered runner was not stored")
	}
	if record.Alias != "linux-build" {
		t.Fatalf("alias = %q, want linux-build", record.Alias)
	}
	if got := record.CommandTypes; len(got) != 2 || got[0] != "Shell" || got[1] != "Git" {
		t.Fatalf("command types = %#v, want [Shell Git]", got)
	}
}

func TestHandleRunnerRegisterRejectsMissingToken(t *testing.T) {
	registry := runnerregistry.New()
	body := bytes.NewBufferString(`{
		"runner_id":"runner-123",
		"protocol_version":"oto.runner.v1"
	}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/register", body)
	rr := httptest.NewRecorder()

	handleRunnerRegister(registry)(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}
	var response otopb.RegisterRunnerResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if response.GetAccepted() {
		t.Fatal("missing token registration was accepted")
	}
	if response.GetRejectReason() != "missing enrollment token" {
		t.Fatalf("reject reason = %q, want missing enrollment token", response.GetRejectReason())
	}
}

func TestHandleRunnerRegisterMethodNotAllowed(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/api/v1/runners/register", nil)
	rr := httptest.NewRecorder()

	handleRunnerRegister(runnerregistry.New())(rr, req)

	if rr.Code != http.StatusMethodNotAllowed {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusMethodNotAllowed)
	}
}

func TestHandleRunnerHeartbeat(t *testing.T) {
	registry := runnerregistry.New()

	// 1. Heartbeat on unregistered runner should fail (404)
	body := bytes.NewBufferString(`{"runner_id":"runner-123", "status":1}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/heartbeat", body)
	rr := httptest.NewRecorder()

	handleRunnerHeartbeat(registry)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("status = %v, want %v for unregistered runner", rr.Code, http.StatusNotFound)
	}

	// 2. Register runner first
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
	})

	// 3. Valid Heartbeat on registered runner should succeed
	body = bytes.NewBufferString(`{"runner_id":"runner-123", "status":1}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/heartbeat", body)
	rr = httptest.NewRecorder()

	handleRunnerHeartbeat(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}

	var response otopb.HeartbeatResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode heartbeat response: %v", err)
	}
	if !response.GetSuccess() {
		t.Fatal("heartbeat failed")
	}

	// Verify status is online
	record, _ := registry.Snapshot("runner-123")
	if record.Status != runnerregistry.StatusOnline {
		t.Fatalf("status = %q, want %q", record.Status, runnerregistry.StatusOnline)
	}
}

func TestHandleRunnerDisconnect(t *testing.T) {
	registry := runnerregistry.New()

	// 1. Disconnect unregistered runner should fail (404)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/disconnect", nil)
	rr := httptest.NewRecorder()

	handleRunnerDisconnect(registry)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusNotFound)
	}

	// 2. Register runner
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
	})

	// 3. Disconnect registered runner should succeed
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/disconnect", nil)
	rr = httptest.NewRecorder()

	handleRunnerDisconnect(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}

	var res struct {
		Success bool `json:"success"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &res); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if !res.Success {
		t.Fatal("expected success=true")
	}

	record, _ := registry.Snapshot("runner-123")
	if record.Status != runnerregistry.StatusDisconnected {
		t.Fatalf("status = %q, want %q", record.Status, runnerregistry.StatusDisconnected)
	}
}

func TestHandleGetRunner(t *testing.T) {
	registry := runnerregistry.New()

	// 1. Get unregistered runner should return 404
	req := httptest.NewRequest(http.MethodGet, "/api/v1/runners/runner-123", nil)
	rr := httptest.NewRecorder()

	handleGetRunner(registry)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusNotFound)
	}

	// 2. Register runner
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		Alias:           "build-linux",
		ProtocolVersion: "oto.runner.v1",
	})

	// 3. Get registered runner should succeed (200) and return JSON record status
	req = httptest.NewRequest(http.MethodGet, "/api/v1/runners/runner-123", nil)
	rr = httptest.NewRecorder()

	handleGetRunner(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}

	var res map[string]interface{}
	if err := json.Unmarshal(rr.Body.Bytes(), &res); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if res["runner_id"] != "runner-123" {
		t.Fatalf("runner_id = %v", res["runner_id"])
	}
	if res["status"] != runnerregistry.StatusAccepted {
		t.Fatalf("status = %v", res["status"])
	}
}

func TestHandleRunnerHeartbeatMismatch(t *testing.T) {
	registry := runnerregistry.New()

	// Register runner-a
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-a",
		ProtocolVersion: "oto.runner.v1",
	})
	// Register runner-b
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-b",
		ProtocolVersion: "oto.runner.v1",
	})

	// Heartbeat with mismatched runner ID in body vs path
	body := bytes.NewBufferString(`{"runner_id":"runner-b", "status":1}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-a/heartbeat", body)
	rr := httptest.NewRecorder()

	handleRunnerHeartbeat(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusBadRequest)
	}

	var response otopb.HeartbeatResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if response.GetSuccess() {
		t.Fatal("expected heartbeat to fail due to ID mismatch")
	}
	if response.GetErrorMessage() != "runner id mismatch between path and body" {
		t.Fatalf("error message = %q", response.GetErrorMessage())
	}

	// Verify runner-b remains accepted (not revived/online)
	rec, _ := registry.Snapshot("runner-b")
	if rec.Status != runnerregistry.StatusAccepted {
		t.Fatalf("runner-b status = %q, want %q", rec.Status, runnerregistry.StatusAccepted)
	}
}

func TestHandleRunnerBootstrapCommand(t *testing.T) {
	registry := runnerregistry.New()

	// 1. Missing runner_id or enrollment_token should be rejected
	bodyMissing := bytes.NewBufferString(`{"runner_id":"","enrollment_token":""}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyMissing)
	rr := httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("status = %v, want %v for missing runner_id/token", rr.Code, http.StatusBadRequest)
	}

	// 2. Missing token only should be rejected
	bodyNoToken := bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":""}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyNoToken)
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("status = %v, want %v for missing token", rr.Code, http.StatusBadRequest)
	}

	// 3. Request should be rejected if server is HTTP and OTO_RUNNER_RELEASE_BASE_URL is not set
	t.Setenv("OTO_RUNNER_RELEASE_BASE_URL", "")
	bodyValid := bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyValid)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("expected 400 Bad Request when HTTP and no release override, got %v", rr.Code)
	}

	// 4. Request should be rejected if OTO_RUNNER_RELEASE_BASE_URL is not HTTPS
	t.Setenv("OTO_RUNNER_RELEASE_BASE_URL", "http://example.com/releases")
	bodyValid = bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyValid)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("expected 400 Bad Request when override is not HTTPS, got %v", rr.Code)
	}

	// 5. Valid request should return correct escaped command when OTO_RUNNER_RELEASE_BASE_URL is set to HTTPS
	t.Setenv("OTO_RUNNER_RELEASE_BASE_URL", "https://example.com/releases")
	bodyValid = bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyValid)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}

	var response otopb.BootstrapCommandResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode bootstrap command response: %v", err)
	}

	expectedCmd := "curl -fsSL 'http://localhost:8080/bootstrap/oto-agent.sh' | bash -s -- --server-url 'http://localhost:8080' --agent-id 'runner-123' --enrollment-token 'token-123' --release-base-url 'https://example.com/releases'"
	if response.GetBootstrapCommand() != expectedCmd {
		t.Fatalf("bootstrap_command = %q, want %q", response.GetBootstrapCommand(), expectedCmd)
	}

	// 6. Request with shell metacharacters in runner_id/token should be safely escaped
	bodyMalicious := bytes.NewBufferString(`{"runner_id":"runner; rm -rf /","enrollment_token":"token'$(say hello)'"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyMalicious)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}

	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode bootstrap command response: %v", err)
	}

	expectedEscapedCmd := "curl -fsSL 'http://localhost:8080/bootstrap/oto-agent.sh' | bash -s -- --server-url 'http://localhost:8080' --agent-id 'runner; rm -rf /' --enrollment-token 'token'\\''$(say hello)'\\''' --release-base-url 'https://example.com/releases'"
	if response.GetBootstrapCommand() != expectedEscapedCmd {
		t.Fatalf("bootstrap_command = %q, want %q", response.GetBootstrapCommand(), expectedEscapedCmd)
	}

	// 7. Malicious Host header with invalid characters should be rejected
	bodyValid = bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyValid)
	req.Host = "localhost; rm -rf /"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("expected 400 Bad Request for malicious Host header, got %v", rr.Code)
	}
}

func TestHandleServeBootstrapScript(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/bootstrap/oto-agent.sh", nil)
	rr := httptest.NewRecorder()

	handleServeBootstrapScript()(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v; body = %s", rr.Code, http.StatusOK, rr.Body.String())
	}

	if contentType := rr.Header().Get("Content-Type"); contentType != "application/x-sh" {
		t.Fatalf("Content-Type = %q, want application/x-sh", contentType)
	}

	if len(rr.Body.Bytes()) == 0 {
		t.Fatal("empty bootstrap script returned")
	}
}