silo-server/internal/api/handlers/diagnostics_chunked_test.go

package handlers

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"

	"github.com/go-chi/chi/v5"

	apimw "github.com/Silo-Server/silo-server/internal/api/middleware"
	"github.com/Silo-Server/silo-server/internal/auth"
	"github.com/Silo-Server/silo-server/internal/diagnostics"
)

var errFakeStatusUnavailable = errors.New("status store temporarily unavailable")

// newChunkedTestHandler returns a handler whose chunk spool lives in the
// test's temp dir so parallel test runs never share state.
func newChunkedTestHandler(t *testing.T, service DiagnosticsService) *DiagnosticsHandler {
	t.Helper()
	handler := NewDiagnosticsHandler(service)
	handler.chunkSessions = newDiagnosticsChunkSessions(t.TempDir())
	return handler
}

func chunkedRouter(handler *DiagnosticsHandler) chi.Router {
	r := chi.NewRouter()
	r.Post("/uploads", handler.HandleChunkedUploadInit)
	r.Put("/uploads/{upload_id}/chunks/{chunk_index}", handler.HandleChunkedUploadChunk)
	r.Post("/uploads/{upload_id}/complete", handler.HandleChunkedUploadComplete)
	r.Delete("/uploads/{upload_id}", handler.HandleChunkedUploadAbort)
	return r
}

func doChunked(t *testing.T, router chi.Router, method, path string, body []byte, claims *auth.Claims) *httptest.ResponseRecorder {
	t.Helper()
	req := httptest.NewRequest(method, path, bytes.NewReader(body))
	req.Header.Set("Authorization", "Bearer token")
	if method == http.MethodPost && strings.HasSuffix(path, "/uploads") {
		req.Header.Set("Content-Type", "application/json")
	}
	if claims != nil {
		req = req.WithContext(apimw.SetClaims(req.Context(), claims))
	}
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)
	return rec
}

func initChunkedUpload(t *testing.T, router chi.Router, bundleBytes int, claims *auth.Claims) diagnosticsChunkInitResponse {
	t.Helper()
	body := []byte(fmt.Sprintf(`{"manifest":{"ok":true},"bundle_bytes":%d}`, bundleBytes))
	rec := doChunked(t, router, http.MethodPost, "/uploads", body, claims)
	if rec.Code != http.StatusCreated {
		t.Fatalf("init status = %d, want %d; body=%s", rec.Code, http.StatusCreated, rec.Body.String())
	}
	var resp diagnosticsChunkInitResponse
	if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
		t.Fatalf("decode init response: %v", err)
	}
	return resp
}

func TestDiagnosticsChunkedUploadHappyPath(t *testing.T) {
	service := newFakeDiagnosticsService()
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	bundle := bytes.Repeat([]byte("b"), int(diagnostics.UploadChunkBytes)+512)
	session := initChunkedUpload(t, router, len(bundle), accessClaims())
	if session.ChunkBytes != diagnostics.UploadChunkBytes {
		t.Fatalf("chunk_bytes = %d, want %d", session.ChunkBytes, diagnostics.UploadChunkBytes)
	}
	if session.TotalChunks != 2 {
		t.Fatalf("total_chunks = %d, want 2", session.TotalChunks)
	}

	for index := 0; index < session.TotalChunks; index++ {
		start := int64(index) * session.ChunkBytes
		end := start + session.ChunkBytes
		if end > int64(len(bundle)) {
			end = int64(len(bundle))
		}
		rec := doChunked(t, router, http.MethodPut,
			fmt.Sprintf("/uploads/%s/chunks/%d", session.UploadID, index),
			bundle[start:end], accessClaims())
		if rec.Code != http.StatusOK {
			t.Fatalf("chunk %d status = %d; body=%s", index, rec.Code, rec.Body.String())
		}
	}

	rec := doChunked(t, router, http.MethodPost, "/uploads/"+session.UploadID+"/complete", nil, accessClaims())
	if rec.Code != http.StatusCreated {
		t.Fatalf("complete status = %d; body=%s", rec.Code, rec.Body.String())
	}
	var resp diagnostics.IngestResult
	if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
		t.Fatalf("decode complete response: %v", err)
	}
	if resp.ShortID != "SILO-ABCDEF123456" {
		t.Fatalf("short_id = %q", resp.ShortID)
	}
	if service.ingestCalls != 1 {
		t.Fatalf("ingest calls = %d, want 1", service.ingestCalls)
	}
	if got := string(service.lastManifest); got != `{"ok":true}` {
		t.Fatalf("manifest passed to ingest = %q", got)
	}
	if int(service.lastBundleBytes) != len(bundle) {
		t.Fatalf("bundle bytes passed to ingest = %d, want %d", service.lastBundleBytes, len(bundle))
	}

	// The ingest slot must be free again after completion.
	release, acquired := handler.inflight.acquire(42)
	if !acquired {
		t.Fatal("in-flight slot still held after complete")
	}
	release()
}

func TestDiagnosticsChunkedUploadInitOverBundleLimit(t *testing.T) {
	service := newFakeDiagnosticsService()
	service.status.MaxBundleBytes = 1024
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	rec := doChunked(t, router, http.MethodPost, "/uploads",
		[]byte(`{"manifest":{"ok":true},"bundle_bytes":2048}`), accessClaims())
	assertDiagnosticsError(t, rec, http.StatusRequestEntityTooLarge, "too_large")

	// Rejected init must not leave the limiter slot held.
	release, acquired := handler.inflight.acquire(42)
	if !acquired {
		t.Fatal("in-flight slot held after rejected init")
	}
	release()
}

func TestDiagnosticsChunkedUploadInitDisabled(t *testing.T) {
	service := newFakeDiagnosticsService()
	service.status.Status = diagnostics.StatusDisabled
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	rec := doChunked(t, router, http.MethodPost, "/uploads",
		[]byte(`{"manifest":{"ok":true},"bundle_bytes":10}`), accessClaims())
	assertDiagnosticsError(t, rec, http.StatusForbidden, "disabled")
}

func TestDiagnosticsChunkedUploadForeignSessionRejected(t *testing.T) {
	service := newFakeDiagnosticsService()
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	session := initChunkedUpload(t, router, 10, accessClaims())

	other := &auth.Claims{UserID: 7, TokenType: auth.TokenTypeAccess}
	rec := doChunked(t, router, http.MethodPut, "/uploads/"+session.UploadID+"/chunks/0",
		[]byte("0123456789"), other)
	if rec.Code != http.StatusNotFound {
		t.Fatalf("foreign chunk status = %d, want 404; body=%s", rec.Code, rec.Body.String())
	}
	rec = doChunked(t, router, http.MethodPost, "/uploads/"+session.UploadID+"/complete", nil, other)
	if rec.Code != http.StatusNotFound {
		t.Fatalf("foreign complete status = %d, want 404", rec.Code)
	}
	// A foreign abort reports success but must not destroy the session.
	rec = doChunked(t, router, http.MethodDelete, "/uploads/"+session.UploadID, nil, other)
	if rec.Code != http.StatusNoContent {
		t.Fatalf("foreign abort status = %d, want 204", rec.Code)
	}
	rec = doChunked(t, router, http.MethodPut, "/uploads/"+session.UploadID+"/chunks/0",
		[]byte("0123456789"), accessClaims())
	if rec.Code != http.StatusOK {
		t.Fatalf("owner chunk after foreign abort = %d, want 200; body=%s", rec.Code, rec.Body.String())
	}
}

func TestDiagnosticsChunkedUploadIncompleteComplete(t *testing.T) {
	service := newFakeDiagnosticsService()
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	bundle := bytes.Repeat([]byte("b"), int(diagnostics.UploadChunkBytes)+512)
	session := initChunkedUpload(t, router, len(bundle), accessClaims())

	rec := doChunked(t, router, http.MethodPut, "/uploads/"+session.UploadID+"/chunks/0",
		bundle[:diagnostics.UploadChunkBytes], accessClaims())
	if rec.Code != http.StatusOK {
		t.Fatalf("chunk status = %d", rec.Code)
	}

	rec = doChunked(t, router, http.MethodPost, "/uploads/"+session.UploadID+"/complete", nil, accessClaims())
	if rec.Code != http.StatusConflict {
		t.Fatalf("incomplete complete status = %d, want 409; body=%s", rec.Code, rec.Body.String())
	}
	if service.ingestCalls != 0 {
		t.Fatalf("ingest calls = %d, want 0", service.ingestCalls)
	}
}

func TestDiagnosticsChunkedUploadReinitReplacesPreviousSession(t *testing.T) {
	service := newFakeDiagnosticsService()
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	first := initChunkedUpload(t, router, 10, accessClaims())
	second := initChunkedUpload(t, router, 10, accessClaims())
	if first.UploadID == second.UploadID {
		t.Fatal("re-init returned the same session id")
	}

	// The replaced session is gone; the new one accepts chunks.
	rec := doChunked(t, router, http.MethodPut, "/uploads/"+first.UploadID+"/chunks/0",
		[]byte("0123456789"), accessClaims())
	if rec.Code != http.StatusNotFound {
		t.Fatalf("replaced-session chunk status = %d, want 404", rec.Code)
	}
	rec = doChunked(t, router, http.MethodPut, "/uploads/"+second.UploadID+"/chunks/0",
		[]byte("0123456789"), accessClaims())
	if rec.Code != http.StatusOK {
		t.Fatalf("new-session chunk status = %d, want 200; body=%s", rec.Code, rec.Body.String())
	}

	rec = doChunked(t, router, http.MethodDelete, "/uploads/"+second.UploadID, nil, accessClaims())
	if rec.Code != http.StatusNoContent {
		t.Fatalf("abort status = %d, want 204", rec.Code)
	}
}

func TestDiagnosticsChunkedCompleteBusyKeepsSessionRetryable(t *testing.T) {
	service := newFakeDiagnosticsService()
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	session := initChunkedUpload(t, router, 10, accessClaims())
	rec := doChunked(t, router, http.MethodPut, "/uploads/"+session.UploadID+"/chunks/0",
		[]byte("0123456789"), accessClaims())
	if rec.Code != http.StatusOK {
		t.Fatalf("chunk status = %d", rec.Code)
	}

	// Occupy the user's ingest slot: complete must answer busy without
	// consuming the spooled session.
	release, acquired := handler.inflight.acquire(42)
	if !acquired {
		t.Fatal("could not occupy in-flight slot")
	}
	rec = doChunked(t, router, http.MethodPost, "/uploads/"+session.UploadID+"/complete", nil, accessClaims())
	assertDiagnosticsError(t, rec, http.StatusServiceUnavailable, "busy")
	release()

	rec = doChunked(t, router, http.MethodPost, "/uploads/"+session.UploadID+"/complete", nil, accessClaims())
	if rec.Code != http.StatusCreated {
		t.Fatalf("retried complete status = %d; body=%s", rec.Code, rec.Body.String())
	}
	if service.ingestCalls != 1 {
		t.Fatalf("ingest calls = %d, want 1", service.ingestCalls)
	}
}

func TestDiagnosticsChunkedUploadOversizedChunkRejected(t *testing.T) {
	service := newFakeDiagnosticsService()
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	bundle := bytes.Repeat([]byte("b"), int(diagnostics.UploadChunkBytes)*2)
	session := initChunkedUpload(t, router, len(bundle), accessClaims())

	oversized := bytes.Repeat([]byte("x"), int(diagnostics.UploadChunkBytes)+1)
	rec := doChunked(t, router, http.MethodPut, "/uploads/"+session.UploadID+"/chunks/0",
		oversized, accessClaims())
	if rec.Code != http.StatusBadRequest && rec.Code != http.StatusRequestEntityTooLarge {
		t.Fatalf("oversized chunk status = %d, want 400 or 413; body=%s", rec.Code, rec.Body.String())
	}
	if service.ingestCalls != 0 {
		t.Fatalf("ingest calls = %d, want 0", service.ingestCalls)
	}
}

func TestDiagnosticsChunkedUploadManifestTooLarge(t *testing.T) {
	service := newFakeDiagnosticsService()
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	manifest := `{"pad":"` + strings.Repeat("x", int(diagnostics.MaxManifestBytes)) + `"}`
	rec := doChunked(t, router, http.MethodPost, "/uploads",
		[]byte(`{"manifest":`+manifest+`,"bundle_bytes":10}`), accessClaims())
	assertDiagnosticsError(t, rec, http.StatusRequestEntityTooLarge, "too_large")
}

func TestDiagnosticsChunkedCompleteIngestErrorMapsLikeSingleShot(t *testing.T) {
	service := newFakeDiagnosticsService()
	service.ingestErr = diagnostics.ErrStaleConsent
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	session := initChunkedUpload(t, router, 10, accessClaims())
	rec := doChunked(t, router, http.MethodPut, "/uploads/"+session.UploadID+"/chunks/0",
		[]byte("0123456789"), accessClaims())
	if rec.Code != http.StatusOK {
		t.Fatalf("chunk status = %d", rec.Code)
	}

	rec = doChunked(t, router, http.MethodPost, "/uploads/"+session.UploadID+"/complete", nil, accessClaims())
	assertDiagnosticsError(t, rec, http.StatusBadRequest, "stale_consent")

	// The failed session is consumed; a retry needs a fresh init and the slot
	// must be free for it.
	initChunkedUpload(t, router, 10, accessClaims())
}

func TestDiagnosticsChunkedUploadConcurrentInitsRespectSlotAtomically(t *testing.T) {
	service := newFakeDiagnosticsService()
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	// Fire concurrent inits for the same user. The reservation must admit at
	// most one at a time — the losers get busy — so a single account can
	// never fan out to multiple sessions and eat the global cap.
	const attempts = 8
	var wg sync.WaitGroup
	codes := make([]int, attempts)
	for i := 0; i < attempts; i++ {
		wg.Add(1)
		go func(slot int) {
			defer wg.Done()
			rec := doChunked(t, router, http.MethodPost, "/uploads",
				[]byte(`{"manifest":{"ok":true},"bundle_bytes":10}`), accessClaims())
			codes[slot] = rec.Code
		}(i)
	}
	wg.Wait()

	created := 0
	for _, code := range codes {
		switch code {
		case http.StatusCreated:
			created++
		case http.StatusServiceUnavailable:
		default:
			t.Fatalf("unexpected init status %d", code)
		}
	}
	if created < 1 {
		t.Fatal("no init succeeded")
	}
	// However the race resolved, the user must end up with at most one live
	// session and no leaked cap slots: after aborting it, a fresh init and
	// full upload must succeed.
	handler.chunkSessions.mu.Lock()
	owned := len(handler.chunkSessions.owners)
	reserved := len(handler.chunkSessions.reserved)
	handler.chunkSessions.mu.Unlock()
	if owned > 1 || reserved != 0 {
		t.Fatalf("owners = %d (want ≤1), reserved = %d (want 0)", owned, reserved)
	}

	session := initChunkedUpload(t, router, 10, accessClaims())
	rec := doChunked(t, router, http.MethodPut, "/uploads/"+session.UploadID+"/chunks/0",
		[]byte("0123456789"), accessClaims())
	if rec.Code != http.StatusOK {
		t.Fatalf("post-race chunk status = %d; body=%s", rec.Code, rec.Body.String())
	}
	rec = doChunked(t, router, http.MethodPost, "/uploads/"+session.UploadID+"/complete", nil, accessClaims())
	if rec.Code != http.StatusCreated {
		t.Fatalf("post-race complete status = %d; body=%s", rec.Code, rec.Body.String())
	}
}

func TestDiagnosticsChunkedCompleteTransientStatusErrorKeepsSession(t *testing.T) {
	service := newFakeDiagnosticsService()
	handler := newChunkedTestHandler(t, service)
	router := chunkedRouter(handler)

	session := initChunkedUpload(t, router, 10, accessClaims())
	rec := doChunked(t, router, http.MethodPut, "/uploads/"+session.UploadID+"/chunks/0",
		[]byte("0123456789"), accessClaims())
	if rec.Code != http.StatusOK {
		t.Fatalf("chunk status = %d", rec.Code)
	}

	// A transient status failure must answer 500 WITHOUT destroying the fully
	// uploaded spool; the retried complete succeeds with no re-upload.
	service.setStatusErr(errFakeStatusUnavailable)
	rec = doChunked(t, router, http.MethodPost, "/uploads/"+session.UploadID+"/complete", nil, accessClaims())
	assertDiagnosticsError(t, rec, http.StatusInternalServerError, "internal_error")

	service.setStatusErr(nil)
	rec = doChunked(t, router, http.MethodPost, "/uploads/"+session.UploadID+"/complete", nil, accessClaims())
	if rec.Code != http.StatusCreated {
		t.Fatalf("retried complete status = %d; body=%s", rec.Code, rec.Body.String())
	}
}

func TestDiagnosticsStatusAdvertisesUploadChunkBytes(t *testing.T) {
	service := newFakeDiagnosticsService()
	service.status.UploadChunkBytes = diagnostics.UploadChunkBytes
	handler := newChunkedTestHandler(t, service)

	req := httptest.NewRequest(http.MethodGet, "/api/v1/diagnostics/status", nil)
	req.Header.Set("Authorization", "Bearer token")
	req = req.WithContext(apimw.SetClaims(req.Context(), accessClaims()))
	rec := httptest.NewRecorder()
	handler.HandleStatus(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d", rec.Code)
	}
	var payload map[string]any
	if err := json.NewDecoder(rec.Body).Decode(&payload); err != nil {
		t.Fatalf("decode status: %v", err)
	}
	got, ok := payload["upload_chunk_bytes"].(float64)
	if !ok || int64(got) != diagnostics.UploadChunkBytes {
		t.Fatalf("upload_chunk_bytes = %v, want %d", payload["upload_chunk_bytes"], diagnostics.UploadChunkBytes)
	}
}