helix-proxy/internal/auth/jwt_test.go

package auth

import (
	"context"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
	"time"
)

func TestJWTManager_GenerateValidate_Roundtrip(t *testing.T) {
	m := NewJWTManager("test-secret-123")
	token, err := m.GenerateToken(42, "user@example.com", "User", []string{"admin", "user"})
	if err != nil {
		t.Fatal(err)
	}
	if token == "" {
		t.Error("empty token")
	}

	claims, err := m.ValidateToken(token)
	if err != nil {
		t.Fatalf("validate: %v", err)
	}
	if claims.UserID != 42 || claims.Email != "user@example.com" || len(claims.Roles) != 2 {
		t.Errorf("claims: %+v", claims)
	}
	if claims.ExpiresAt == nil || time.Until(claims.ExpiresAt.Time) > time.Hour {
		t.Error("expiry not ~1h")
	}
}

func TestJWTManager_EmptySecretRandomRoundtrip(t *testing.T) {
	m := NewJWTManager("") // empty -> random in-memory secret
	token, _ := m.GenerateToken(1, "a@b", "A", nil)
	_, err := m.ValidateToken(token)
	if err != nil {
		t.Error("random secret should allow roundtrip")
	}
	m2 := NewJWTManager("")
	if _, err := m2.ValidateToken(token); err == nil {
		t.Error("separate empty managers must not share a well-known secret")
	}
}

func TestLoadOrCreateSecret_EnvWins(t *testing.T) {
	t.Setenv("JWT_SECRET", "from-env")
	got, err := LoadOrCreateSecret(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	if got != "from-env" {
		t.Fatalf("got %q", got)
	}
}

func TestLoadOrCreateSecret_Persists(t *testing.T) {
	t.Setenv("JWT_SECRET", "")
	dir := t.TempDir()
	a, err := LoadOrCreateSecret(dir)
	if err != nil {
		t.Fatal(err)
	}
	if a == "" {
		t.Fatal("empty secret")
	}
	b, err := LoadOrCreateSecret(dir)
	if err != nil {
		t.Fatal(err)
	}
	if a != b {
		t.Fatalf("secret not persisted: %q vs %q", a, b)
	}
	raw, err := os.ReadFile(filepath.Join(dir, ".jwt_secret"))
	if err != nil {
		t.Fatal(err)
	}
	if string(raw) == "" {
		t.Fatal("secret file empty")
	}
}

func TestJWTManager_BadToken(t *testing.T) {
	m := NewJWTManager("s")
	if _, err := m.ValidateToken("not.a.jwt"); err == nil {
		t.Error("bad token should err")
	}
	// wrong sig
	m2 := NewJWTManager("other")
	tok, _ := m2.GenerateToken(1, "e", "n", nil)
	if _, err := m.ValidateToken(tok); err == nil {
		t.Error("wrong sig should invalid")
	}
}

func TestJWTManager_Middleware(t *testing.T) {
	m := NewJWTManager("sec")
	token, _ := m.GenerateToken(7, "m@e", "M", []string{"admin"})

	// with valid bearer
	h := m.Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		c, ok := GetUserFromContext(r.Context())
		if !ok || c.UserID != 7 {
			t.Error("claims not in ctx")
		}
		w.WriteHeader(200)
	}))
	req := httptest.NewRequest("GET", "/", nil)
	req.Header.Set("Authorization", "Bearer "+token)
	rr := httptest.NewRecorder()
	h.ServeHTTP(rr, req)
	if rr.Code != 200 {
		t.Errorf("valid bearer: %d", rr.Code)
	}

	// no token: passes through (no 401)
	h2 := m.Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(204)
	}))
	req2 := httptest.NewRequest("GET", "/", nil)
	rr2 := httptest.NewRecorder()
	h2.ServeHTTP(rr2, req2)
	if rr2.Code != 204 {
		t.Error("no token should pass")
	}

	// bad token: 401
	req3 := httptest.NewRequest("GET", "/", nil)
	req3.Header.Set("Authorization", "Bearer bad")
	rr3 := httptest.NewRecorder()
	h.ServeHTTP(rr3, req3)
	if rr3.Code != 401 {
		t.Errorf("bad token 401: %d", rr3.Code)
	}

	// wrong scheme
	req4 := httptest.NewRequest("GET", "/", nil)
	req4.Header.Set("Authorization", "Basic foo")
	rr4 := httptest.NewRecorder()
	h.ServeHTTP(rr4, req4)
	if rr4.Code != 401 {
		t.Errorf("bad scheme: %d", rr4.Code)
	}
}

func TestRequireAuth(t *testing.T) {
	// direct, needs ctx with claims
	m := NewJWTManager("s")
	tok, _ := m.GenerateToken(99, "r@a", "R", nil)
	claims, _ := m.ValidateToken(tok)

	req := httptest.NewRequest("GET", "/", nil)
	// without ctx
	rr := httptest.NewRecorder()
	c, ok := RequireAuth(rr, req)
	if ok || c != nil || rr.Code != 401 {
		t.Error("require without claims should 401")
	}

	// with ctx set via context value
	req2 := httptest.NewRequest("GET", "/", nil)
	ctx := context.WithValue(req2.Context(), UserContextKey, claims)
	req2 = req2.WithContext(ctx)
	rr2 := httptest.NewRecorder()
	c2, ok2 := RequireAuth(rr2, req2)
	if !ok2 || c2.UserID != 99 {
		t.Error("require with claims failed")
	}
}