OpenFlare/internal/apps/openflare/apiutil/middleware_test.go

// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0

package apiutil

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

	"github.com/Rain-kl/Wavelet/internal/apps/admin"
	"github.com/Rain-kl/Wavelet/internal/apps/oauth"
	"github.com/Rain-kl/Wavelet/internal/common/response"
	"github.com/Rain-kl/Wavelet/internal/config"
	"github.com/Rain-kl/Wavelet/internal/db"
	"github.com/Rain-kl/Wavelet/internal/db/idgen"
	"github.com/Rain-kl/Wavelet/internal/model"
	"github.com/Rain-kl/Wavelet/internal/testhelper"
	"github.com/gin-contrib/sessions"
	"github.com/gin-contrib/sessions/cookie"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func setupAdminMiddlewareTest(t *testing.T) (*gin.Engine, *gorm.DB, func()) {
	t.Helper()

	dbConn, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
		DisableForeignKeyConstraintWhenMigrating: true,
	})
	require.NoError(t, err)
	require.NoError(t, dbConn.AutoMigrate(&model.User{}, &model.AccessToken{}))
	db.SetDB(dbConn)

	sessionCookieName := "test_admin_middleware_session"
	if config.Config.App.SessionCookieName != "" {
		sessionCookieName = config.Config.App.SessionCookieName
	}
	store := cookie.NewStore([]byte("test_admin_middleware_session_secret"))
	store.Options(oauth.GetSessionOptions(3600))
	engine := testhelper.NewTestGinEngine(sessions.Sessions(sessionCookieName, store))
	protected := engine.Group("/protected", AdminMiddlewares()...)
	protected.GET("", func(c *gin.Context) {
		c.JSON(http.StatusOK, response.OK(gin.H{"ok": true}))
	})

	cleanup := func() {
		db.SetDB(nil)
	}

	return engine, dbConn, cleanup
}

func seedUser(t *testing.T, dbConn *gorm.DB, username string, isAdmin bool) *model.User {
	t.Helper()

	user := &model.User{
		ID:       idgen.NextUint64ID(),
		Username: username,
		Nickname: username,
		Email:    username + "@openflare.test",
		IsActive: true,
		IsAdmin:  isAdmin,
	}
	require.NoError(t, dbConn.Create(user).Error)
	return user
}

func seedAccessToken(t *testing.T, dbConn *gorm.DB, user *model.User, isAdmin bool) string {
	t.Helper()

	token, err := model.GenerateTokenString()
	require.NoError(t, err)
	require.NoError(t, dbConn.Create(&model.AccessToken{
		UserID:      user.ID,
		Name:        user.Username + "-token",
		TokenHash:   model.HashToken(token),
		MaskedToken: model.MaskTokenString(token),
		IsAdmin:     isAdmin,
	}).Error)
	return token
}

func decodeResponse(t *testing.T, rec *httptest.ResponseRecorder) response.Any {
	t.Helper()

	var resp response.Any
	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
	return resp
}

func TestAdminRequiredUnauthenticated(t *testing.T) {
	engine, _, cleanup := setupAdminMiddlewareTest(t)
	defer cleanup()

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/protected", nil)
	engine.ServeHTTP(rec, req)

	assert.Equal(t, http.StatusUnauthorized, rec.Code)
	resp := decodeResponse(t, rec)
	assert.NotEmpty(t, resp.ErrorMsg)
}

func TestAdminRequiredNonAdminToken(t *testing.T) {
	engine, dbConn, cleanup := setupAdminMiddlewareTest(t)
	defer cleanup()

	user := seedUser(t, dbConn, "regular", false)
	token := seedAccessToken(t, dbConn, user, false)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/protected", nil)
	req.Header.Set("X-Access-Token", token)
	engine.ServeHTTP(rec, req)

	assert.Equal(t, http.StatusNotFound, rec.Code)
	resp := decodeResponse(t, rec)
	assert.Equal(t, admin.TokenAdminRequired, resp.ErrorMsg)
}

func TestAdminRequiredAdminWithoutTokenAdmin(t *testing.T) {
	engine, dbConn, cleanup := setupAdminMiddlewareTest(t)
	defer cleanup()

	user := seedUser(t, dbConn, "admin-no-token-admin", true)
	token := seedAccessToken(t, dbConn, user, false)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/protected", nil)
	req.Header.Set("X-Access-Token", token)
	engine.ServeHTTP(rec, req)

	assert.Equal(t, http.StatusNotFound, rec.Code)
	resp := decodeResponse(t, rec)
	assert.Equal(t, admin.TokenAdminRequired, resp.ErrorMsg)
}

func TestAdminRequiredAdminWithTokenAdmin(t *testing.T) {
	engine, dbConn, cleanup := setupAdminMiddlewareTest(t)
	defer cleanup()

	user := seedUser(t, dbConn, "admin", true)
	token := seedAccessToken(t, dbConn, user, true)

	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/protected", nil)
	req.Header.Set("X-Access-Token", token)
	engine.ServeHTTP(rec, req)

	assert.Equal(t, http.StatusOK, rec.Code)
	resp := decodeResponse(t, rec)
	assert.Empty(t, resp.ErrorMsg)
}