pangolin/server/internal/auth/handler_test.go

package auth

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/wangjia/pangolin/server/internal/codes"
)

// newTestHandler builds a Handler + chi router for HTTP-level tests.
func newTestHandler(t *testing.T, cfg ServiceConfig) (*Service, http.Handler) {
	t.Helper()
	svc, _, _ := newService(t, cfg)
	h := NewHandler(svc)
	r := chi.NewRouter()
	h.RegisterRoutes(r)
	return svc, r
}

func doJSON(t *testing.T, h http.Handler, method, path string, body interface{}) *httptest.ResponseRecorder {
	t.Helper()
	var buf bytes.Buffer
	if body != nil {
		_ = json.NewEncoder(&buf).Encode(body)
	}
	req := httptest.NewRequest(method, path, &buf)
	req.RemoteAddr = "203.0.113.5:1234"
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return rec
}

func TestHandler_SendCode_204(t *testing.T) {
	_, h := newTestHandler(t, ServiceConfig{})
	rec := doJSON(t, h, http.MethodPost, "/auth/code", map[string]string{"email": "u@example.com"})
	if rec.Code != http.StatusNoContent {
		t.Fatalf("status = %d, want 204 (body %s)", rec.Code, rec.Body)
	}
}

func TestHandler_SendCode_RateLimited_RetryAfter(t *testing.T) {
	_, h := newTestHandler(t, ServiceConfig{EmailPerMinute: 1})
	_ = doJSON(t, h, http.MethodPost, "/auth/code", map[string]string{"email": "u@example.com"})
	rec := doJSON(t, h, http.MethodPost, "/auth/code", map[string]string{"email": "u@example.com"})
	if rec.Code != http.StatusTooManyRequests {
		t.Fatalf("status = %d, want 429", rec.Code)
	}
	ra := rec.Header().Get("Retry-After")
	if ra == "" {
		t.Fatal("missing Retry-After header")
	}
	if n, err := strconv.Atoi(ra); err != nil || n <= 0 {
		t.Fatalf("bad Retry-After %q", ra)
	}
}

func TestHandler_RegisterLoginRefresh_HTTP(t *testing.T) {
	svc, h := newTestHandler(t, ServiceConfig{})
	const email = "flow@example.com"

	// Send code, read it from Redis.
	_ = doJSON(t, h, http.MethodPost, "/auth/code", map[string]string{"email": email})
	code := codeInRedis(t, svc, email)

	// Register.
	rec := doJSON(t, h, http.MethodPost, "/auth/register", map[string]string{
		"email": email, "code": code, "password": "password123",
	})
	if rec.Code != http.StatusOK {
		t.Fatalf("register status = %d (body %s)", rec.Code, rec.Body)
	}
	var reg tokenPairResponse
	_ = json.Unmarshal(rec.Body.Bytes(), &reg)
	if reg.AccessToken == "" || reg.ExpiresIn != 900 {
		t.Fatalf("bad register response: %+v", reg)
	}

	// Login.
	rec = doJSON(t, h, http.MethodPost, "/auth/login", map[string]string{
		"email": email, "password": "password123",
	})
	if rec.Code != http.StatusOK {
		t.Fatalf("login status = %d (body %s)", rec.Code, rec.Body)
	}
	var login tokenPairResponse
	_ = json.Unmarshal(rec.Body.Bytes(), &login)

	// Refresh.
	rec = doJSON(t, h, http.MethodPost, "/auth/refresh", map[string]string{
		"refresh_token": login.RefreshToken,
	})
	if rec.Code != http.StatusOK {
		t.Fatalf("refresh status = %d (body %s)", rec.Code, rec.Body)
	}
}

func TestMiddleware_RequireAuth(t *testing.T) {
	rdb, _ := newMiniRedis(t)
	tm := newTokenManager(t, rdb, time.Now)

	// Protected handler that echoes the injected user id.
	var gotUID int64
	var gotOK bool
	protected := RequireAuth(tm)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotUID, gotOK = UserIDFromContext(r.Context())
		// Confirm the codes module reads the same value via its exported key.
		if v, ok := r.Context().Value(codes.CtxKeyUserID).(int64); !ok || v != gotUID {
			t.Errorf("codes.CtxKeyUserID mismatch: %v", v)
		}
		w.WriteHeader(http.StatusOK)
	}))

	// No token → 401.
	rec := httptest.NewRecorder()
	protected.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/me", nil))
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("no-token status = %d, want 401", rec.Code)
	}

	// Valid token → 200 with injected uid.
	pair, _ := tm.Issue(context.Background(), 77, "uuid-77")
	req := httptest.NewRequest(http.MethodGet, "/me", nil)
	req.Header.Set("Authorization", "Bearer "+pair.AccessToken)
	rec = httptest.NewRecorder()
	protected.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("valid-token status = %d, want 200", rec.Code)
	}
	if !gotOK || gotUID != 77 {
		t.Fatalf("injected uid = %d ok=%v, want 77", gotUID, gotOK)
	}

	// Refresh token must be rejected by access middleware.
	req = httptest.NewRequest(http.MethodGet, "/me", nil)
	req.Header.Set("Authorization", "Bearer "+pair.RefreshToken)
	rec = httptest.NewRecorder()
	protected.ServeHTTP(rec, req)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("refresh-as-access status = %d, want 401", rec.Code)
	}
}