usb-server/internal/relay/diag_test.go

package relay

import (
	"bytes"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func newDiagServer() *Server {
	return &Server{hub: NewHub(), diag: newDiagStore()}
}

func TestDiagStoreAndFetch(t *testing.T) {
	s := newDiagServer()
	body := []byte(`{"tool":{"os":"windows"}}`)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPut, "/diag/report-1", bytes.NewReader(body))
	s.handleDiag(rec, req)

	if rec.Code != http.StatusCreated {
		t.Fatalf("store returned %d, want %d", rec.Code, http.StatusCreated)
	}

	rec = httptest.NewRecorder()
	req = httptest.NewRequest(http.MethodGet, "/diag/report-1", nil)
	s.handleDiag(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("fetch returned %d, want 200", rec.Code)
	}
	if !bytes.Equal(rec.Body.Bytes(), body) {
		t.Errorf("fetched %q, want %q", rec.Body.String(), body)
	}
	if ct := rec.Header().Get("Content-Type"); ct != "application/json" {
		t.Errorf("content type %q, want application/json", ct)
	}
}

func TestDiagMissingReportIs404(t *testing.T) {
	s := newDiagServer()

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/diag/nope", nil)
	s.handleDiag(rec, req)

	if rec.Code != http.StatusNotFound {
		t.Errorf("got %d, want 404", rec.Code)
	}
}

func TestDiagRejectsBadPaths(t *testing.T) {
	s := newDiagServer()

	for _, path := range []string{"/diag/", "/diag/a/b"} {
		rec := httptest.NewRecorder()
		req := httptest.NewRequest(http.MethodGet, path, nil)
		s.handleDiag(rec, req)

		if rec.Code != http.StatusBadRequest {
			t.Errorf("%s returned %d, want 400", path, rec.Code)
		}
	}
}

// The endpoint takes uploads from anyone who can reach the relay, so it must
// bound what one caller can make it hold.
func TestDiagRejectsOversizedReport(t *testing.T) {
	s := newDiagServer()

	huge := bytes.Repeat([]byte("x"), maxDiagSize+100)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPut, "/diag/big", bytes.NewReader(huge))
	s.handleDiag(rec, req)

	if rec.Code != http.StatusRequestEntityTooLarge {
		t.Errorf("got %d, want 413", rec.Code)
	}

	// And it must not have been stored anyway.
	if _, ok := s.diag.get("big"); ok {
		t.Error("an oversized report was stored")
	}
}

// A lying Content-Length must not get past the limit either.
func TestDiagLimitIgnoresContentLength(t *testing.T) {
	s := newDiagServer()

	huge := bytes.Repeat([]byte("x"), maxDiagSize+100)
	req := httptest.NewRequest(http.MethodPut, "/diag/liar", bytes.NewReader(huge))
	req.ContentLength = 10 // claims to be small

	rec := httptest.NewRecorder()
	s.handleDiag(rec, req)

	if rec.Code == http.StatusCreated {
		t.Error("an oversized body was accepted because it claimed to be small")
	}
}

func TestDiagRejectsEmptyReport(t *testing.T) {
	s := newDiagServer()

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPut, "/diag/empty", bytes.NewReader(nil))
	s.handleDiag(rec, req)

	if rec.Code != http.StatusBadRequest {
		t.Errorf("got %d, want 400", rec.Code)
	}
}

func TestDiagRejectsOtherMethods(t *testing.T) {
	s := newDiagServer()

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodDelete, "/diag/x", nil)
	s.handleDiag(rec, req)

	if rec.Code != http.StatusMethodNotAllowed {
		t.Errorf("got %d, want 405", rec.Code)
	}
}

// The store is bounded, so a stream of uploads cannot grow it without limit.
func TestDiagEvictsOldestWhenFull(t *testing.T) {
	store := newDiagStore()

	for i := 0; i < maxDiagReports+5; i++ {
		store.put(fmt.Sprintf("report-%d", i), []byte("{}"), "127.0.0.1")
		// Ordering by timestamp needs the timestamps to differ.
		time.Sleep(time.Millisecond)
	}

	store.mu.Lock()
	count := len(store.reports)
	store.mu.Unlock()

	if count > maxDiagReports {
		t.Errorf("store holds %d reports, limit is %d", count, maxDiagReports)
	}

	// The newest must have survived; the very first must not have.
	if _, ok := store.get(fmt.Sprintf("report-%d", maxDiagReports+4)); !ok {
		t.Error("the most recent report was evicted")
	}
	if _, ok := store.get("report-0"); ok {
		t.Error("the oldest report survived eviction")
	}
}

func TestDiagExpiresOldReports(t *testing.T) {
	store := newDiagStore()

	store.put("old", []byte("{}"), "127.0.0.1")

	// Backdate it past the TTL.
	store.mu.Lock()
	store.reports["old"].stored = time.Now().Add(-diagTTL - time.Minute)
	store.mu.Unlock()

	if _, ok := store.get("old"); ok {
		t.Error("a report past its TTL was still served")
	}
}

func TestClientIPPrefersForwardedHeader(t *testing.T) {
	tests := []struct {
		name  string
		setup func(*http.Request)
		want  string
	}{
		{
			name:  "remote address",
			setup: func(r *http.Request) { r.RemoteAddr = "203.0.113.7:12345" },
			want:  "203.0.113.7",
		},
		{
			name: "forwarded header wins",
			setup: func(r *http.Request) {
				r.RemoteAddr = "10.0.0.1:12345"
				r.Header.Set("X-Forwarded-For", "203.0.113.7, 10.0.0.2")
			},
			want: "203.0.113.7",
		},
		{
			name: "real ip header",
			setup: func(r *http.Request) {
				r.RemoteAddr = "10.0.0.1:12345"
				r.Header.Set("X-Real-IP", "203.0.113.9")
			},
			want: "203.0.113.9",
		},
		{
			name: "garbage header falls back",
			setup: func(r *http.Request) {
				r.RemoteAddr = "203.0.113.7:12345"
				r.Header.Set("X-Forwarded-For", "not-an-ip")
			},
			want: "203.0.113.7",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, "/", nil)
			req.Header = http.Header{}
			tt.setup(req)

			if got := clientIP(req); got != tt.want {
				t.Errorf("clientIP() = %q, want %q", got, tt.want)
			}
		})
	}
}

func TestRetentionNoteIsMentionedOnStore(t *testing.T) {
	s := newDiagServer()

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPut, "/diag/x", strings.NewReader("{}"))
	s.handleDiag(rec, req)

	if !strings.Contains(rec.Body.String(), RetentionNote) {
		t.Errorf("the response does not say how long the report is kept: %q", rec.Body.String())
	}
}