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