accounts/account/api/api_test.go

package api

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

	"github.com/gin-gonic/gin"
	"github.com/pquerna/otp"
	"github.com/pquerna/otp/totp"
)

type apiResponse struct {
	Message   string                 `json:"message"`
	Error     string                 `json:"error"`
	Token     string                 `json:"token"`
	MFAToken  string                 `json:"mfaToken"`
	User      map[string]interface{} `json:"user"`
	MFA       map[string]interface{} `json:"mfa"`
	Secret    string                 `json:"secret"`
	URI       string                 `json:"uri"`
	ExpiresAt string                 `json:"expiresAt"`
}

func decodeResponse(t *testing.T, rr *httptest.ResponseRecorder) apiResponse {
	t.Helper()
	var resp apiResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil {
		t.Fatalf("failed to decode response: %v", err)
	}
	return resp
}

func waitForStableTOTPWindow(t *testing.T) {
	t.Helper()
	const period int64 = 30
	remainder := time.Now().Unix() % period
	const buffer int64 = 10
	if remainder > period-buffer {
		sleep := (period - remainder) + 2
		if sleep > 0 {
			time.Sleep(time.Duration(sleep) * time.Second)
		}
	}
}

func TestRegisterEndpoint(t *testing.T) {
	gin.SetMode(gin.TestMode)

	router := gin.New()
	RegisterRoutes(router)

	payload := map[string]string{
		"name":     "Test User",
		"email":    "user@example.com",
		"password": "supersecure",
	}

	body, err := json.Marshal(payload)
	if err != nil {
		t.Fatalf("failed to marshal payload: %v", err)
	}

	req := httptest.NewRequest(http.MethodPost, "/api/auth/register", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")

	rr := httptest.NewRecorder()
	router.ServeHTTP(rr, req)

	if rr.Code != http.StatusCreated {
		t.Fatalf("expected status %d, got %d, body: %s", http.StatusCreated, rr.Code, rr.Body.String())
	}

	resp := decodeResponse(t, rr)
	if resp.User == nil {
		t.Fatalf("expected user object in response")
	}

	if email, ok := resp.User["email"].(string); !ok || email != payload["email"] {
		t.Fatalf("expected email %q, got %#v", payload["email"], resp.User["email"])
	}

	if id, ok := resp.User["id"].(string); !ok || id == "" {
		t.Fatalf("expected user id in response")
	} else if uuid, ok := resp.User["uuid"].(string); !ok || uuid != id {
		t.Fatalf("expected uuid to match id")
	}

	if mfaEnabled, ok := resp.User["mfaEnabled"].(bool); !ok || mfaEnabled {
		t.Fatalf("expected mfaEnabled to be false, got %#v", resp.User["mfaEnabled"])
	}

	mfaData, ok := resp.User["mfa"].(map[string]interface{})
	if !ok {
		t.Fatalf("expected mfa state in user payload")
	}
	if enabled, ok := mfaData["totpEnabled"].(bool); !ok || enabled {
		t.Fatalf("expected totpEnabled to be false, got %#v", mfaData["totpEnabled"])
	}
	if pending, ok := mfaData["totpPending"].(bool); !ok || pending {
		t.Fatalf("expected totpPending to be false, got %#v", mfaData["totpPending"])
	}
}

func TestMFATOTPFlow(t *testing.T) {
	gin.SetMode(gin.TestMode)

	router := gin.New()
	RegisterRoutes(router)

	registerPayload := map[string]string{
		"name":     "Login User",
		"email":    "login@example.com",
		"password": "supersecure",
	}
	registerBody, err := json.Marshal(registerPayload)
	if err != nil {
		t.Fatalf("failed to marshal registration payload: %v", err)
	}

	req := httptest.NewRequest(http.MethodPost, "/api/auth/register", bytes.NewReader(registerBody))
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()
	router.ServeHTTP(rr, req)
	if rr.Code != http.StatusCreated {
		t.Fatalf("expected registration to succeed, got %d", rr.Code)
	}

	loginPayload := map[string]string{
		"identifier": "Login User",
		"password":   registerPayload["password"],
	}
	loginBody, err := json.Marshal(loginPayload)
	if err != nil {
		t.Fatalf("failed to marshal login payload: %v", err)
	}

	req = httptest.NewRequest(http.MethodPost, "/api/auth/login", bytes.NewReader(loginBody))
	req.Header.Set("Content-Type", "application/json")
	rr = httptest.NewRecorder()
	router.ServeHTTP(rr, req)

	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("expected login to require mfa setup, got %d", rr.Code)
	}
	resp := decodeResponse(t, rr)
	if resp.Error != "mfa_setup_required" {
		t.Fatalf("expected mfa_setup_required error, got %q", resp.Error)
	}
	if resp.MFAToken == "" {
		t.Fatalf("expected mfa token in response")
	}

	provisionPayload := map[string]string{
		"token": resp.MFAToken,
	}
	provisionBody, err := json.Marshal(provisionPayload)
	if err != nil {
		t.Fatalf("failed to marshal provision payload: %v", err)
	}

	req = httptest.NewRequest(http.MethodPost, "/api/auth/mfa/totp/provision", bytes.NewReader(provisionBody))
	req.Header.Set("Content-Type", "application/json")
	rr = httptest.NewRecorder()
	router.ServeHTTP(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("expected provisioning success, got %d: %s", rr.Code, rr.Body.String())
	}
	resp = decodeResponse(t, rr)
	if resp.Secret == "" {
		t.Fatalf("expected totp secret in provisioning response")
	}
	if resp.URI == "" {
		t.Fatalf("expected otpauth uri in provisioning response")
	}
	secret := resp.Secret

	generateCode := func(offset time.Duration) string {
		code, err := totp.GenerateCodeCustom(secret, time.Now().UTC().Add(offset), totp.ValidateOpts{
			Period:    30,
			Skew:      1,
			Digits:    otp.DigitsSix,
			Algorithm: otp.AlgorithmSHA1,
		})
		if err != nil {
			t.Fatalf("failed to generate verification code: %v", err)
		}
		return code
	}

	waitForStableTOTPWindow(t)
	code := generateCode(-30 * time.Second)

	verifyPayload := map[string]string{
		"token": resp.MFAToken,
		"code":  code,
	}
	verifyBody, err := json.Marshal(verifyPayload)
	if err != nil {
		t.Fatalf("failed to marshal verify payload: %v", err)
	}

	req = httptest.NewRequest(http.MethodPost, "/api/auth/mfa/totp/verify", bytes.NewReader(verifyBody))
	req.Header.Set("Content-Type", "application/json")
	rr = httptest.NewRecorder()
	router.ServeHTTP(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("expected verification success, got %d: %s", rr.Code, rr.Body.String())
	}
	resp = decodeResponse(t, rr)
	if resp.Token == "" {
		t.Fatalf("expected session token after verification")
	}
	if resp.User == nil || resp.User["mfaEnabled"] != true {
		t.Fatalf("expected mfaEnabled true after verification")
	}

	sessionReq := httptest.NewRequest(http.MethodGet, "/api/auth/session", nil)
	sessionReq.Header.Set("Authorization", "Bearer "+resp.Token)
	sessionRec := httptest.NewRecorder()
	router.ServeHTTP(sessionRec, sessionReq)
	if sessionRec.Code != http.StatusOK {
		t.Fatalf("expected session lookup success, got %d", sessionRec.Code)
	}
	sessionResp := decodeResponse(t, sessionRec)
	if sessionResp.User == nil {
		t.Fatalf("expected user in session response")
	}
	if sessionResp.User["mfaEnabled"] != true {
		t.Fatalf("expected session user to have mfaEnabled true")
	}

	statusReq := httptest.NewRequest(http.MethodGet, "/api/auth/mfa/status", nil)
	statusReq.Header.Set("Authorization", "Bearer "+resp.Token)
	statusRec := httptest.NewRecorder()
	router.ServeHTTP(statusRec, statusReq)
	if statusRec.Code != http.StatusOK {
		t.Fatalf("expected status success, got %d", statusRec.Code)
	}

	loginWithTotp := func(body map[string]string) *httptest.ResponseRecorder {
		payload, err := json.Marshal(body)
		if err != nil {
			t.Fatalf("failed to marshal login payload: %v", err)
		}
		request := httptest.NewRequest(http.MethodPost, "/api/auth/login", bytes.NewReader(payload))
		request.Header.Set("Content-Type", "application/json")
		recorder := httptest.NewRecorder()
		router.ServeHTTP(recorder, request)
		return recorder
	}

	waitForStableTOTPWindow(t)
	totpCode := generateCode(-30 * time.Second)
	if ok, _ := totp.ValidateCustom(totpCode, secret, time.Now().UTC(), totp.ValidateOpts{
		Period:    30,
		Skew:      1,
		Digits:    otp.DigitsSix,
		Algorithm: otp.AlgorithmSHA1,
	}); !ok {
		t.Fatalf("locally generated totp code is invalid")
	}

	rr = loginWithTotp(map[string]string{
		"identifier": "Login User",
		"password":   registerPayload["password"],
		"totpCode":   totpCode,
	})
	if rr.Code != http.StatusOK {
		t.Fatalf("expected mfa login success, got %d: %s", rr.Code, rr.Body.String())
	}

	waitForStableTOTPWindow(t)
	totpCode = generateCode(0)
	if ok, _ := totp.ValidateCustom(totpCode, secret, time.Now().UTC(), totp.ValidateOpts{
		Period:    30,
		Skew:      1,
		Digits:    otp.DigitsSix,
		Algorithm: otp.AlgorithmSHA1,
	}); !ok {
		t.Fatalf("locally generated totp code is invalid (email login)")
	}

	rr = loginWithTotp(map[string]string{
		"identifier": registerPayload["email"],
		"totpCode":   totpCode,
	})
	if rr.Code != http.StatusOK {
		t.Fatalf("expected email+totp login success, got %d: %s", rr.Code, rr.Body.String())
	}
}