feat: publish sanitized custom.97 source

middleware/header_nav_test.go

@@ -4,12 +4,14 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/QuantumNous/new-api/common"
 	"github.com/QuantumNous/new-api/model"
 	"github.com/QuantumNous/new-api/service"
 	"github.com/gin-gonic/gin"
 	"github.com/glebarez/sqlite"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"gorm.io/gorm"
 )
@@ -50,9 +52,10 @@ func performHeaderNavRequest(t *testing.T, handler gin.HandlerFunc, authenticate
 		previousDB, previousRedis := model.DB, common.RedisEnabled
 		db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
 		require.NoError(t, err)
-		require.NoError(t, db.AutoMigrate(&model.User{}))
+		require.NoError(t, db.AutoMigrate(&model.User{}, &model.Option{}))
 		model.DB = db
 		common.RedisEnabled = false
+		require.NoError(t, model.EnsureLegacyAccessTokenRetireAt(time.Now().Unix()))
 		t.Cleanup(func() {
 			model.DB = previousDB
 			common.RedisEnabled = previousRedis
@@ -96,6 +99,79 @@ func TestHeaderNavModuleAuthRejectsDisabledPricing(t *testing.T) {
 	require.Equal(t, http.StatusForbidden, recorder.Code)
 }
 
+func TestDisabledPricingRequiresRootAuthentication(t *testing.T) {
+	for _, tc := range []struct {
+		name       string
+		module     string
+		options    string
+		role       int
+		credential string
+		wantStatus int
+	}{
+		{name: "root", role: common.RoleRootUser, wantStatus: http.StatusOK},
+		{name: "root with login required", options: `{"pricing":{"enabled":false,"requireAuth":true}}`, role: common.RoleRootUser, wantStatus: http.StatusOK},
+		{name: "root with legacy disabled setting", options: `{"pricing":false}`, role: common.RoleRootUser, wantStatus: http.StatusOK},
+		{name: "ordinary user", role: common.RoleCommonUser, wantStatus: http.StatusForbidden},
+		{name: "non-root administrator", role: common.RoleAdminUser, wantStatus: http.StatusForbidden},
+		{name: "expired root token", role: common.RoleRootUser, credential: "expired", wantStatus: http.StatusUnauthorized},
+		{name: "tampered root token", role: common.RoleRootUser, credential: "tampered", wantStatus: http.StatusUnauthorized},
+		{name: "revoked root session", role: common.RoleRootUser, credential: "revoked", wantStatus: http.StatusUnauthorized},
+		{name: "disabled root account", role: common.RoleRootUser, credential: "banned", wantStatus: http.StatusUnauthorized},
+		{name: "root cannot bypass disabled rankings", module: "rankings", options: `{"rankings":false}`, role: common.RoleRootUser, wantStatus: http.StatusForbidden},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			setupDashboardAuthMiddlewareTest(t)
+			gin.SetMode(gin.TestMode)
+			options := tc.options
+			if options == "" {
+				options = `{"pricing":{"enabled":false,"requireAuth":false}}`
+			}
+			withHeaderNavModules(t, options)
+			user := model.User{
+				Username: "pricing-reader", Password: "unused-password-hash", Role: tc.role,
+				Status: common.UserStatusEnabled, Group: "default", AuthVersion: 1,
+			}
+			require.NoError(t, model.DB.Create(&user).Error)
+			bundle, err := service.CreateLoginSession(user.Id, "password", "127.0.0.1", "header-nav-test")
+			require.NoError(t, err)
+			token := bundle.AccessToken
+			switch tc.credential {
+			case "expired":
+				identity, internal, err := service.ParseDashboardAccessToken(token)
+				require.NoError(t, err)
+				require.True(t, internal)
+				token = issueExpiredDashboardAccessToken(t, identity)
+			case "tampered":
+				token = tamperDashboardToken(token)
+			case "revoked":
+				revoked, err := model.RevokeUserSession(user.Id, bundle.Session.SID, "test")
+				require.NoError(t, err)
+				require.True(t, revoked)
+			case "banned":
+				require.NoError(t, model.DB.Model(&user).Update("status", common.UserStatusDisabled).Error)
+			}
+
+			module := tc.module
+			if module == "" {
+				module = "pricing"
+			}
+			handled := false
+			router := gin.New()
+			router.GET("/api/"+module, HeaderNavModuleAuth(module), func(c *gin.Context) {
+				handled = true
+				assert.Equal(t, user.Id, c.GetInt("id"))
+				assert.Equal(t, common.RoleRootUser, c.GetInt("role"))
+				c.JSON(http.StatusOK, gin.H{"success": true})
+			})
+
+			response := middlewareBearerRequest(router, "/api/"+module, token)
+
+			assert.Equal(t, tc.wantStatus, response.Code)
+			assert.Equal(t, tc.wantStatus == http.StatusOK, handled)
+		})
+	}
+}
+
 func TestHeaderNavModuleAuthRequiresLoginForPricing(t *testing.T) {
 	raw := `{"pricing":{"enabled":true,"requireAuth":true}}`
 	withHeaderNavModules(t, raw)