Files
gochat/internal/handler/api/v1/mfa_handler_test.go
T
2026-06-04 15:44:48 +08:00

664 lines
19 KiB
Go

package v1
import (
"bytes"
"crypto/hmac"
"crypto/sha1"
"encoding/base32"
"encoding/binary"
"encoding/json"
"fmt"
"math"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/suite"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"gorm.io/gorm/logger"
"github.com/gochat/gochat/internal/auth"
"github.com/gochat/gochat/internal/model"
"github.com/gochat/gochat/pkg/response"
)
// --- MFA Handler Test Suite ---
// Uses real SQLite DB + real MFAService + httptest.
type MFAHandlerTestSuite struct {
suite.Suite
db *gorm.DB
router *gin.Engine
handler *MFAHandler
mfaService *auth.MFAService
user *model.User
account *model.Account
userID uint
accountID uint
}
func TestMFAHandlerSuite(t *testing.T) {
suite.Run(t, new(MFAHandlerTestSuite))
}
func (s *MFAHandlerTestSuite) SetupSuite() {
gin.SetMode(gin.TestMode)
// Create in-memory SQLite DB
db, err := gorm.Open(sqlite.Open("file::memory:"), &gorm.Config{
Logger: logger.Default.LogMode(logger.Silent),
})
s.Require().NoError(err, "failed to open SQLite test database")
// Migrate models needed for MFA operations
s.Require().NoError(db.AutoMigrate(
&model.Account{},
&model.User{},
))
s.db = db
// Create real MFA service backed by the test DB
s.mfaService = auth.NewMFAService(db)
s.handler = NewMFAHandler(s.mfaService)
// Create test account
account := &model.Account{Name: "MFATestAccount"}
s.Require().NoError(db.Create(account).Error)
s.account = account
s.accountID = account.ID
// Create test user belonging to the account
user := &model.User{
AccountID: account.ID,
Name: "MFA Test User",
Email: "mfatest@example.com",
Password: "hashedpassword123",
Provider: "email",
Role: "agent",
Active: true,
}
s.Require().NoError(db.Create(user).Error)
s.user = user
s.userID = user.ID
// Setup router with middleware that injects user_id into context
s.setupRouter(s.userID)
}
func (s *MFAHandlerTestSuite) setupRouter(userID uint) {
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user_id", userID)
c.Next()
})
mfaGroup := r.Group("/api/v1/auth/mfa")
{
mfaGroup.POST("/enable", s.handler.EnableMFA)
mfaGroup.POST("/verify", s.handler.VerifyMFA)
mfaGroup.POST("/disable", s.handler.DisableMFA)
}
s.router = r
}
func (s *MFAHandlerTestSuite) SetupTest() {
// Hard cleanup: delete all users and accounts, then recreate
s.db.Exec("DELETE FROM users")
s.db.Exec("DELETE FROM accounts")
// Recreate test data
account := &model.Account{Name: "MFATestAccount"}
s.Require().NoError(s.db.Create(account).Error)
s.account = account
s.accountID = account.ID
user := &model.User{
AccountID: account.ID,
Name: "MFA Test User",
Email: "mfatest@example.com",
Password: "hashedpassword123",
Provider: "email",
Role: "agent",
Active: true,
}
s.Require().NoError(s.db.Create(user).Error)
s.user = user
s.userID = user.ID
// Re-setup router with the new user ID
s.setupRouter(s.userID)
}
func (s *MFAHandlerTestSuite) TearDownSuite() {
if s.db != nil {
sqlDB, err := s.db.DB()
if err == nil {
sqlDB.Close()
}
}
}
// --- Helper to make requests and parse responses ---
func (s *MFAHandlerTestSuite) doRequest(method, path, body string) *httptest.ResponseRecorder {
var reqBody *bytes.Buffer
if body != "" {
reqBody = bytes.NewBufferString(body)
} else {
reqBody = bytes.NewBufferString("")
}
w := httptest.NewRecorder()
req := httptest.NewRequest(method, path, reqBody)
if body != "" {
req.Header.Set("Content-Type", "application/json")
}
s.router.ServeHTTP(w, req)
return w
}
func (s *MFAHandlerTestSuite) parseResponse(w *httptest.ResponseRecorder) response.APIResponse {
var resp response.APIResponse
s.NoError(json.Unmarshal(w.Body.Bytes(), &resp))
return resp
}
// --- Helper to generate a valid TOTP code for a secret ---
// Uses the same algorithm as auth.validateTOTP/generateTOTP to compute a valid code.
func (s *MFAHandlerTestSuite) generateValidTOTPCode(secret string) string {
cfg := auth.DefaultTOTPConfig()
return computeTOTPCode(secret, cfg)
}
func computeTOTPCode(secret string, cfg auth.TOTPConfig) string {
key, err := decodeBase32NoPad(secret)
if err != nil {
return ""
}
now := time.Now().Unix()
period := int64(cfg.Period)
timeCounter := now / period
return generateTOTPFromKey(key, timeCounter, cfg)
}
func decodeBase32NoPad(secret string) ([]byte, error) {
return base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(strings.ToUpper(secret))
}
func generateTOTPFromKey(key []byte, timeCounter int64, cfg auth.TOTPConfig) string {
// Encode time counter as 8-byte big-endian
buf := make([]byte, 8)
binary.BigEndian.PutUint64(buf, uint64(timeCounter))
// HMAC-SHA1
h := hmac.New(sha1.New, key)
h.Write(buf)
hash := h.Sum(nil)
// Dynamic truncation per RFC 4226
offset := hash[len(hash)-1] & 0x0f
truncated := (int32(hash[offset]&0x7f) << 24) |
(int32(hash[offset+1]&0xff) << 16) |
(int32(hash[offset+2]&0xff) << 8) |
(int32(hash[offset+3]&0xff))
// Modulo 10^digits
mod := int32(math.Pow10(cfg.Digits))
code := truncated % mod
// Format with leading zeros
return fmt.Sprintf("%0*d", cfg.Digits, code)
}
func jsonBody(data map[string]interface{}) string {
b, err := json.Marshal(data)
if err != nil {
return ""
}
return string(b)
}
// ============================================================
// EnableMFA tests
// ============================================================
func (s *MFAHandlerTestSuite) TestEnable_Success() {
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/enable", "{}")
s.Equal(http.StatusOK, w.Code)
resp := s.parseResponse(w)
s.True(resp.Success)
dataMap, ok := resp.Data.(map[string]interface{})
s.True(ok)
// Response should contain totp_secret and qr_uri
s.NotEmpty(dataMap["totp_secret"])
s.NotEmpty(dataMap["qr_uri"])
s.Contains(dataMap["qr_uri"], "otpauth://totp/")
s.Contains(dataMap["qr_uri"], dataMap["totp_secret"])
}
func (s *MFAHandlerTestSuite) TestEnable_NoBody() {
// EnableMFARequest has no required fields, empty body should still work
// since the handler doesn't even call ShouldBindJSON
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/enable", "")
// With empty body and no Content-Type, the handler doesn't bind JSON,
// so it just uses user_id from context → should succeed
s.Equal(http.StatusOK, w.Code)
resp := s.parseResponse(w)
s.True(resp.Success)
}
func (s *MFAHandlerTestSuite) TestEnable_InvalidJSON() {
// Enable handler does NOT call ShouldBindJSON at all — it only uses
// c.GetUint("user_id") and service calls. So invalid JSON in the body
// won't cause a binding error. With a valid user_id in context,
// this should succeed regardless of body content.
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/enable", "{invalid}")
// The handler ignores the body entirely, so with valid user_id it succeeds
s.Equal(http.StatusOK, w.Code)
}
func (s *MFAHandlerTestSuite) TestEnable_Unauthorized_NoUserID() {
// Create router without user_id middleware — user_id will be 0
r := gin.New()
r.POST("/api/v1/auth/mfa/enable", s.handler.EnableMFA)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/mfa/enable", bytes.NewBufferString("{}"))
req.Header.Set("Content-Type", "application/json")
r.ServeHTTP(w, req)
s.Equal(http.StatusUnauthorized, w.Code)
var respStruct response.APIResponse
s.NoError(json.Unmarshal(w.Body.Bytes(), &respStruct))
s.False(respStruct.Success)
s.NotNil(respStruct.Error)
s.Equal(response.ErrUnauthorized, respStruct.Error.Code)
}
func (s *MFAHandlerTestSuite) TestEnable_UserNotFound() {
// Router that sets a non-existent user ID
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user_id", uint(9999)) // non-existent user
c.Next()
})
r.POST("/api/v1/auth/mfa/enable", s.handler.EnableMFA)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/mfa/enable", bytes.NewBufferString("{}"))
req.Header.Set("Content-Type", "application/json")
r.ServeHTTP(w, req)
// IsMFAEnabled will fail to find the user → 500 Internal Server Error
s.Equal(http.StatusUnprocessableEntity, w.Code)
var respStruct response.APIResponse
s.NoError(json.Unmarshal(w.Body.Bytes(), &respStruct))
s.False(respStruct.Success)
s.NotNil(respStruct.Error)
s.Equal(response.ErrInternal, respStruct.Error.Code)
}
func (s *MFAHandlerTestSuite) TestEnable_AlreadyEnabled() {
// First enable MFA for the user
user := s.user
user.TOTPSecret = "JBSWY3DPEHPK3PXP"
user.TOTPEnabled = true
s.Require().NoError(s.db.Save(user).Error)
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/enable", "{}")
s.Equal(http.StatusConflict, w.Code)
resp := s.parseResponse(w)
s.False(resp.Success)
s.NotNil(resp.Error)
s.Equal(response.ErrConflict, resp.Error.Code)
}
// ============================================================
// VerifyMFA tests
// ============================================================
func (s *MFAHandlerTestSuite) TestVerify_Success() {
// Step 1: Generate secret
secret, _, err := s.mfaService.GenerateTOTPSecret(s.userID)
s.Require().NoError(err)
// Step 2: Compute a valid TOTP code for the secret
code := s.generateValidTOTPCode(secret)
// Step 3: Verify with secret + code
body := jsonBody(map[string]interface{}{
"totp_secret": secret,
"totp_code": code,
})
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/verify", body)
s.Equal(http.StatusOK, w.Code)
resp := s.parseResponse(w)
s.True(resp.Success)
dataMap, ok := resp.Data.(map[string]interface{})
s.True(ok)
s.Equal("MFA enabled successfully", dataMap["message"])
s.Equal(true, dataMap["mfa_enabled"])
// Verify that TOTPEnabled is now true in DB
var updatedUser model.User
s.Require().NoError(s.db.First(&updatedUser, s.userID).Error)
s.True(updatedUser.TOTPEnabled)
s.Equal(secret, updatedUser.TOTPSecret)
}
func (s *MFAHandlerTestSuite) TestVerify_InvalidJSON() {
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/verify", "{invalid}")
s.Equal(http.StatusBadRequest, w.Code)
resp := s.parseResponse(w)
s.False(resp.Success)
s.NotNil(resp.Error)
s.Equal(response.ErrValidation, resp.Error.Code)
}
func (s *MFAHandlerTestSuite) TestVerify_MissingTOTPSecret() {
// When totp_secret is missing from JSON, ShouldBindJSON fails with
// binding:"required" validation error → handler returns ErrValidation
body := jsonBody(map[string]interface{}{
"totp_code": "123456",
})
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/verify", body)
// Missing required binding field → ShouldBindJSON error → 400 VALIDATION_ERROR
s.Equal(http.StatusBadRequest, w.Code)
resp := s.parseResponse(w)
s.False(resp.Success)
s.NotNil(resp.Error)
s.Equal(response.ErrValidation, resp.Error.Code)
}
func (s *MFAHandlerTestSuite) TestVerify_MissingTOTPCode() {
secret, _, err := s.mfaService.GenerateTOTPSecret(s.userID)
s.Require().NoError(err)
// When totp_code is missing from JSON, ShouldBindJSON fails with
// binding:"required" validation error → handler returns ErrValidation
body := jsonBody(map[string]interface{}{
"totp_secret": secret,
})
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/verify", body)
// Missing required binding field → ShouldBindJSON error → 400 VALIDATION_ERROR
s.Equal(http.StatusBadRequest, w.Code)
resp := s.parseResponse(w)
s.False(resp.Success)
s.NotNil(resp.Error)
s.Equal(response.ErrValidation, resp.Error.Code)
}
func (s *MFAHandlerTestSuite) TestVerify_InvalidTOTPCode() {
secret, _, err := s.mfaService.GenerateTOTPSecret(s.userID)
s.Require().NoError(err)
body := jsonBody(map[string]interface{}{
"totp_secret": secret,
"totp_code": "000000", // definitely wrong code
})
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/verify", body)
// ValidateTOTPCode returns false → 400 Bad Request
s.Equal(http.StatusBadRequest, w.Code)
resp := s.parseResponse(w)
s.False(resp.Success)
s.NotNil(resp.Error)
s.Equal(response.ErrBadRequest, resp.Error.Code)
}
func (s *MFAHandlerTestSuite) TestVerify_Unauthorized_NoUserID() {
r := gin.New()
r.POST("/api/v1/auth/mfa/verify", s.handler.VerifyMFA)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/mfa/verify", bytes.NewBufferString(`{"totp_secret":"abc","totp_code":"123456"}`))
req.Header.Set("Content-Type", "application/json")
r.ServeHTTP(w, req)
s.Equal(http.StatusUnauthorized, w.Code)
var respStruct response.APIResponse
s.NoError(json.Unmarshal(w.Body.Bytes(), &respStruct))
s.False(respStruct.Success)
s.Equal(response.ErrUnauthorized, respStruct.Error.Code)
}
// ============================================================
// DisableMFA tests
// ============================================================
func (s *MFAHandlerTestSuite) TestDisable_Success() {
// First, enable MFA for the user so we can disable it
secret, _, err := s.mfaService.GenerateTOTPSecret(s.userID)
s.Require().NoError(err)
// Enable TOTP via service directly
s.Require().NoError(s.mfaService.EnableTOTP(s.userID, secret))
// Now the user has TOTP enabled. Generate a current valid code for disable.
disableCode := s.generateValidTOTPCode(secret)
body := jsonBody(map[string]interface{}{
"totp_code": disableCode,
})
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/disable", body)
s.Equal(http.StatusOK, w.Code)
resp := s.parseResponse(w)
s.True(resp.Success)
dataMap, ok := resp.Data.(map[string]interface{})
s.True(ok)
s.Equal("MFA disabled successfully", dataMap["message"])
s.Equal(false, dataMap["mfa_enabled"])
// Verify that TOTPEnabled is now false in DB
var updatedUser model.User
s.Require().NoError(s.db.First(&updatedUser, s.userID).Error)
s.False(updatedUser.TOTPEnabled)
s.Empty(updatedUser.TOTPSecret)
}
func (s *MFAHandlerTestSuite) TestDisable_InvalidJSON() {
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/disable", "{invalid}")
s.Equal(http.StatusBadRequest, w.Code)
resp := s.parseResponse(w)
s.False(resp.Success)
s.NotNil(resp.Error)
s.Equal(response.ErrValidation, resp.Error.Code)
}
func (s *MFAHandlerTestSuite) TestDisable_MissingTOTPCode() {
// First enable MFA
secret, _, err := s.mfaService.GenerateTOTPSecret(s.userID)
s.Require().NoError(err)
s.Require().NoError(s.mfaService.EnableTOTP(s.userID, secret))
// When totp_code is missing from JSON, ShouldBindJSON fails with
// binding:"required" validation error → handler returns ErrValidation
body := jsonBody(map[string]interface{}{})
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/disable", body)
// Missing required binding field → ShouldBindJSON error → 400 VALIDATION_ERROR
s.Equal(http.StatusBadRequest, w.Code)
resp := s.parseResponse(w)
s.False(resp.Success)
s.NotNil(resp.Error)
s.Equal(response.ErrValidation, resp.Error.Code)
}
func (s *MFAHandlerTestSuite) TestDisable_InvalidTOTPCode() {
// First enable MFA
secret, _, err := s.mfaService.GenerateTOTPSecret(s.userID)
s.Require().NoError(err)
s.Require().NoError(s.mfaService.EnableTOTP(s.userID, secret))
body := jsonBody(map[string]interface{}{
"totp_code": "000000", // definitely wrong
})
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/disable", body)
// DisableTOTP → VerifyTOTPCode → invalid code → error → 400 Bad Request
s.Equal(http.StatusBadRequest, w.Code)
resp := s.parseResponse(w)
s.False(resp.Success)
s.NotNil(resp.Error)
s.Equal(response.ErrBadRequest, resp.Error.Code)
}
func (s *MFAHandlerTestSuite) TestDisable_Unauthorized_NoUserID() {
r := gin.New()
r.POST("/api/v1/auth/mfa/disable", s.handler.DisableMFA)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/mfa/disable", bytes.NewBufferString(`{"totp_code":"123456"}`))
req.Header.Set("Content-Type", "application/json")
r.ServeHTTP(w, req)
s.Equal(http.StatusUnauthorized, w.Code)
var respStruct response.APIResponse
s.NoError(json.Unmarshal(w.Body.Bytes(), &respStruct))
s.False(respStruct.Success)
s.Equal(response.ErrUnauthorized, respStruct.Error.Code)
}
func (s *MFAHandlerTestSuite) TestDisable_MFANotEnabled() {
// User does not have MFA enabled — DisableTOTP calls VerifyTOTPCode
// which checks user.TOTPEnabled == false → error
body := jsonBody(map[string]interface{}{
"totp_code": "123456",
})
w := s.doRequest(http.MethodPost, "/api/v1/auth/mfa/disable", body)
// VerifyTOTPCode will return error "mfa not enabled for user" → 400 Bad Request
s.Equal(http.StatusBadRequest, w.Code)
resp := s.parseResponse(w)
s.False(resp.Success)
s.NotNil(resp.Error)
s.Equal(response.ErrBadRequest, resp.Error.Code)
}
// ============================================================
// MFAStatus tests (bonus coverage for the status endpoint)
// ============================================================
func (s *MFAHandlerTestSuite) TestStatus_MFADisabled() {
// Setup router with status route
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user_id", s.userID)
c.Next()
})
r.GET("/api/v1/auth/mfa/status", s.handler.MFAStatus)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/v1/auth/mfa/status", nil)
r.ServeHTTP(w, req)
s.Equal(http.StatusOK, w.Code)
var respStruct response.APIResponse
s.NoError(json.Unmarshal(w.Body.Bytes(), &respStruct))
s.True(respStruct.Success)
dataMap, ok := respStruct.Data.(map[string]interface{})
s.True(ok)
s.Equal(false, dataMap["mfa_enabled"])
}
func (s *MFAHandlerTestSuite) TestStatus_MFAEnabled() {
// Enable MFA first
secret, _, err := s.mfaService.GenerateTOTPSecret(s.userID)
s.Require().NoError(err)
s.Require().NoError(s.mfaService.EnableTOTP(s.userID, secret))
// Setup router with status route
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set("user_id", s.userID)
c.Next()
})
r.GET("/api/v1/auth/mfa/status", s.handler.MFAStatus)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/v1/auth/mfa/status", nil)
r.ServeHTTP(w, req)
s.Equal(http.StatusOK, w.Code)
var respStruct response.APIResponse
s.NoError(json.Unmarshal(w.Body.Bytes(), &respStruct))
s.True(respStruct.Success)
dataMap, ok := respStruct.Data.(map[string]interface{})
s.True(ok)
s.Equal(true, dataMap["mfa_enabled"])
}
func (s *MFAHandlerTestSuite) TestStatus_Unauthorized_NoUserID() {
r := gin.New()
r.GET("/api/v1/auth/mfa/status", s.handler.MFAStatus)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/v1/auth/mfa/status", nil)
r.ServeHTTP(w, req)
s.Equal(http.StatusUnauthorized, w.Code)
var respStruct response.APIResponse
s.NoError(json.Unmarshal(w.Body.Bytes(), &respStruct))
s.False(respStruct.Success)
s.Equal(response.ErrUnauthorized, respStruct.Error.Code)
}
// Ensure unused import warning doesn't cause issues
var _ = assert.Equal