Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cmd/maxx/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -463,7 +463,7 @@ func main() {
requestTracker := core.NewRequestTracker()

// Create handlers
proxyHandler := handler.NewProxyHandler(clientAdapter, requestExecutor, cachedSessionRepo, tokenAuthMiddleware)
proxyHandler := handler.NewProxyHandler(clientAdapter, requestExecutor, cachedSessionRepo, settingRepo, tokenAuthMiddleware)
proxyHandler.SetRequestTracker(requestTracker)
adminHandler := handler.NewAdminHandler(adminService, backupService, logPath)
selfServiceHandler := handler.NewSelfServiceHandler(adminService)
Expand Down
2 changes: 1 addition & 1 deletion internal/core/database.go
Original file line number Diff line number Diff line change
Expand Up @@ -452,7 +452,7 @@ func InitializeServerComponents(

log.Printf("[Core] Creating handlers")
tokenAuthMiddleware := handler.NewTokenAuthMiddleware(repos.CachedAPITokenRepo, repos.SettingRepo)
proxyHandler := handler.NewProxyHandler(clientAdapter, exec, repos.CachedSessionRepo, tokenAuthMiddleware)
proxyHandler := handler.NewProxyHandler(clientAdapter, exec, repos.CachedSessionRepo, repos.SettingRepo, tokenAuthMiddleware)
modelsHandler := handler.NewModelsHandler(
repos.ResponseModelRepo,
repos.CachedProviderRepo,
Expand Down
1 change: 1 addition & 0 deletions internal/domain/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -688,6 +688,7 @@ const (
SettingKeyAutoSortCodex = "auto_sort_codex" // 是否自动排序 Codex 路由,"true" 或 "false"
SettingKeyCodexInstructionsEnabled = "codex_instructions_enabled" // 是否启用 Codex 官方 instructions,"true" 或 "false"
SettingKeyPayloadOverrideRules = "payload_override_rules" // 请求 payload 覆盖规则(JSON 数组)
SettingKeyProxyRequestsDisabled = "proxy_requests_disabled" // 是否全局禁用代理请求,"true" 或 "false",默认 "false"
SettingKeyEnablePprof = "enable_pprof" // 是否启用 pprof 性能分析,"true" 或 "false",默认 "false"
SettingKeyPprofPort = "pprof_port" // pprof 服务端口,默认 6060
SettingKeyPprofPassword = "pprof_password" // pprof 访问密码,为空表示不需要密码
Expand Down
22 changes: 22 additions & 0 deletions internal/handler/proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,13 @@ import (
"github.com/awsl-project/maxx/internal/domain"
"github.com/awsl-project/maxx/internal/executor"
"github.com/awsl-project/maxx/internal/flow"
"github.com/awsl-project/maxx/internal/repository"
"github.com/awsl-project/maxx/internal/repository/cached"
"github.com/awsl-project/maxx/internal/systemsettingcache"
)

const proxyRequestsDisabledMessage = "proxy requests are temporarily disabled by admin"

// RequestTracker interface for tracking active requests
type RequestTracker interface {
Add() bool
Expand All @@ -32,6 +36,7 @@ type ProxyHandler struct {
clientAdapter *client.Adapter
executor *executor.Executor
sessionRepo *cached.SessionRepository
settingRepo repository.SystemSettingRepository
tokenAuth *TokenAuthMiddleware
tracker RequestTracker
trackerMu sync.RWMutex
Expand All @@ -44,12 +49,14 @@ func NewProxyHandler(
clientAdapter *client.Adapter,
exec *executor.Executor,
sessionRepo *cached.SessionRepository,
settingRepo repository.SystemSettingRepository,
tokenAuth *TokenAuthMiddleware,
) *ProxyHandler {
h := &ProxyHandler{
clientAdapter: clientAdapter,
executor: exec,
sessionRepo: sessionRepo,
settingRepo: settingRepo,
tokenAuth: tokenAuth,
engine: flow.NewEngine(),
}
Expand Down Expand Up @@ -103,6 +110,13 @@ func (h *ProxyHandler) ingress(c *flow.Ctx) {
return
}

if h.isProxyRequestsDisabled() {
log.Printf("[Proxy] Rejecting request because proxy kill switch is enabled: %s %s", r.Method, r.URL.Path)
writeError(w, http.StatusServiceUnavailable, proxyRequestsDisabledMessage)
c.Abort()
return
}

if strings.HasPrefix(r.URL.Path, "/v1/responses") {
r.URL.Path = strings.TrimPrefix(r.URL.Path, "/v1")
}
Expand Down Expand Up @@ -303,6 +317,14 @@ func normalizeOpenAIChatCompletionsPayload(body []byte) ([]byte, bool) {
return converted, true
}

func (h *ProxyHandler) isProxyRequestsDisabled() bool {
return isBooleanSystemSettingEnabled(h.settingRepo, domain.SettingKeyProxyRequestsDisabled)
}

func isBooleanSystemSettingEnabled(repo repository.SystemSettingRepository, key string) bool {
return systemsettingcache.GetBoolean(repo, key)
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

// Helper functions

func writeError(w http.ResponseWriter, status int, message string) {
Expand Down
96 changes: 96 additions & 0 deletions internal/handler/proxy_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package handler

import (
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"strconv"
Expand All @@ -10,8 +11,34 @@ import (
"time"

"github.com/awsl-project/maxx/internal/domain"
"github.com/awsl-project/maxx/internal/systemsettingcache"
)

type proxyBooleanSettingRepo struct {
values []string
errs []error
reads int
}

func (r *proxyBooleanSettingRepo) Get(key string) (string, error) {
r.reads++
idx := r.reads - 1
if idx < len(r.errs) && r.errs[idx] != nil {
return "", r.errs[idx]
}
if idx < len(r.values) {
return r.values[idx], nil
}
if len(r.values) > 0 {
return r.values[len(r.values)-1], nil
}
return "", nil
}

func (r *proxyBooleanSettingRepo) Set(key, value string) error { return nil }
func (r *proxyBooleanSettingRepo) GetAll() ([]*domain.SystemSetting, error) { return nil, nil }
func (r *proxyBooleanSettingRepo) Delete(key string) error { return nil }

func TestWriteError(t *testing.T) {
rec := httptest.NewRecorder()
writeError(rec, http.StatusBadRequest, "bad request")
Expand Down Expand Up @@ -123,3 +150,72 @@ func TestWriteStreamErrorPreservesStatusAndRetryAfter(t *testing.T) {
t.Fatalf("stream body = %q, want error event", rec.Body.String())
}
}

func TestProxyHandlerRejectsRequestsWhenKillSwitchEnabled(t *testing.T) {
repo := &settingsTestRepo{values: map[string]string{domain.SettingKeyProxyRequestsDisabled: "true"}}
h := NewProxyHandler(nil, nil, nil, repo, nil)

req := httptest.NewRequest(http.MethodPost, "http://example.test/v1/chat/completions", strings.NewReader(`{"model":"gpt-5.4","messages":[]}`))
rec := httptest.NewRecorder()

h.ServeHTTP(rec, req)

if rec.Code != http.StatusServiceUnavailable {
t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusServiceUnavailable, rec.Body.String())
}
if !strings.Contains(rec.Body.String(), proxyRequestsDisabledMessage) {
t.Fatalf("body = %q, want disabled message", rec.Body.String())
}
}

func TestIsBooleanSystemSettingEnabledDefaultsFalse(t *testing.T) {
systemsettingcache.Invalidate(domain.SettingKeyProxyRequestsDisabled)
repo := &settingsTestRepo{}
if isBooleanSystemSettingEnabled(repo, domain.SettingKeyProxyRequestsDisabled) {
t.Fatal("expected missing setting to default to false")
}
}

func TestIsBooleanSystemSettingEnabledCachesFreshValue(t *testing.T) {
oldTTL := systemsettingcache.BooleanTTL
systemsettingcache.BooleanTTL = time.Hour
defer func() { systemsettingcache.BooleanTTL = oldTTL }()

key := domain.SettingKeyProxyRequestsDisabled
systemsettingcache.Invalidate(key)
repo := &proxyBooleanSettingRepo{values: []string{"true"}}

if !isBooleanSystemSettingEnabled(repo, key) {
t.Fatal("expected first read to return true")
}
if !isBooleanSystemSettingEnabled(repo, key) {
t.Fatal("expected cached read to return true")
}
if repo.reads != 1 {
t.Fatalf("reads = %d, want 1", repo.reads)
}
}

func TestIsBooleanSystemSettingEnabledFallsBackToLastKnownValueOnReadError(t *testing.T) {
oldTTL := systemsettingcache.BooleanTTL
systemsettingcache.BooleanTTL = time.Nanosecond
defer func() { systemsettingcache.BooleanTTL = oldTTL }()

key := domain.SettingKeyProxyRequestsDisabled
systemsettingcache.Invalidate(key)
repo := &proxyBooleanSettingRepo{
values: []string{"true"},
errs: []error{nil, errors.New("db temporarily unavailable")},
}

if !isBooleanSystemSettingEnabled(repo, key) {
t.Fatal("expected first read to return true")
}
time.Sleep(time.Millisecond)
if !isBooleanSystemSettingEnabled(repo, key) {
t.Fatal("expected cached true value on refresh error")
}
if repo.reads != 2 {
t.Fatalf("reads = %d, want 2", repo.reads)
}
}
3 changes: 3 additions & 0 deletions internal/service/admin.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import (
"github.com/awsl-project/maxx/internal/payloadoverride"
"github.com/awsl-project/maxx/internal/pricing"
"github.com/awsl-project/maxx/internal/repository"
"github.com/awsl-project/maxx/internal/systemsettingcache"
"github.com/awsl-project/maxx/internal/version"
)

Expand Down Expand Up @@ -486,6 +487,7 @@ func (s *AdminService) UpdateSetting(key, value string) error {
if err := s.settingRepo.Set(key, value); err != nil {
return err
}
systemsettingcache.Invalidate(key)
if key == domain.SettingKeyPayloadOverrideRules {
payloadoverride.InvalidateGlobalSettingsCache()
}
Expand All @@ -507,6 +509,7 @@ func (s *AdminService) DeleteSetting(key string) error {
if err := s.settingRepo.Delete(key); err != nil {
return err
}
systemsettingcache.Invalidate(key)
if key == domain.SettingKeyPayloadOverrideRules {
payloadoverride.InvalidateGlobalSettingsCache()
}
Expand Down
18 changes: 18 additions & 0 deletions internal/service/system_setting_validation.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
package service

import (
"fmt"
"strings"

"github.com/awsl-project/maxx/internal/domain"
"github.com/awsl-project/maxx/internal/payloadoverride"
)
Expand All @@ -9,7 +12,22 @@ func validateSystemSettingValue(key, value string) error {
switch key {
case domain.SettingKeyPayloadOverrideRules:
return payloadoverride.ValidateRulesJSON(value)
case domain.SettingKeyProxyRequestsDisabled:
return validateBooleanSystemSetting(value)
default:
return nil
}
}

func validateBooleanSystemSetting(value string) error {
if strings.TrimSpace(value) != value {
return fmt.Errorf("%w: boolean setting must not contain surrounding whitespace", domain.ErrInvalidInput)
}

switch value {
case "true", "false":
return nil
default:
return fmt.Errorf("%w: boolean setting must be \"true\" or \"false\"", domain.ErrInvalidInput)
}
}
83 changes: 83 additions & 0 deletions internal/service/system_setting_validation_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,10 @@ package service
import (
"errors"
"testing"
"time"

"github.com/awsl-project/maxx/internal/domain"
"github.com/awsl-project/maxx/internal/systemsettingcache"
)

type stubSystemSettingRepo struct {
Expand Down Expand Up @@ -48,6 +50,87 @@ func TestAdminServiceUpdateSettingRejectsInvalidPayloadOverrideRules(t *testing.
}
}

func TestAdminServiceUpdateSettingRejectsInvalidProxyRequestsDisabledValue(t *testing.T) {
repo := &stubSystemSettingRepo{}
svc := &AdminService{settingRepo: repo}

err := svc.UpdateSetting(domain.SettingKeyProxyRequestsDisabled, "maybe")
if !errors.Is(err, domain.ErrInvalidInput) {
t.Fatalf("expected invalid input error, got %v", err)
}
if len(repo.values) != 0 {
t.Fatalf("expected invalid setting not to be persisted")
}
}

func TestAdminServiceUpdateSettingAcceptsValidProxyRequestsDisabledValue(t *testing.T) {
repo := &stubSystemSettingRepo{}
svc := &AdminService{settingRepo: repo}

err := svc.UpdateSetting(domain.SettingKeyProxyRequestsDisabled, "true")
if err != nil {
t.Fatalf("expected no error, got %v", err)
}
if got := repo.values[domain.SettingKeyProxyRequestsDisabled]; got != "true" {
t.Fatalf("stored value = %q, want true", got)
}
}

func TestAdminServiceUpdateSettingRejectsProxyRequestsDisabledValueWithWhitespace(t *testing.T) {
repo := &stubSystemSettingRepo{}
svc := &AdminService{settingRepo: repo}

err := svc.UpdateSetting(domain.SettingKeyProxyRequestsDisabled, " true ")
if !errors.Is(err, domain.ErrInvalidInput) {
t.Fatalf("expected invalid input error, got %v", err)
}
if len(repo.values) != 0 {
t.Fatalf("expected invalid setting not to be persisted")
}
}

func TestAdminServiceUpdateSettingInvalidatesProxyBooleanCache(t *testing.T) {
oldTTL := systemsettingcache.BooleanTTL
systemsettingcache.BooleanTTL = time.Hour
defer func() { systemsettingcache.BooleanTTL = oldTTL }()
Comment on lines +92 to +95

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

补上全局缓存清理,避免测试之间相互污染。

这两条用例会修改 systemsettingcache.BooleanTTL 并填充全局缓存,但现在只恢复了 TTL,没有在结束时清掉 proxy_requests_disabled 的缓存值。后续同包测试如果也读取这个 key,结果会依赖执行顺序。

🧪 建议修复
 func TestAdminServiceUpdateSettingInvalidatesProxyBooleanCache(t *testing.T) {
 	oldTTL := systemsettingcache.BooleanTTL
 	systemsettingcache.BooleanTTL = time.Hour
-	defer func() { systemsettingcache.BooleanTTL = oldTTL }()
+	t.Cleanup(func() {
+		systemsettingcache.BooleanTTL = oldTTL
+		systemsettingcache.Invalidate(domain.SettingKeyProxyRequestsDisabled)
+	})
@@
 func TestAdminServiceDeleteSettingInvalidatesProxyBooleanCache(t *testing.T) {
 	oldTTL := systemsettingcache.BooleanTTL
 	systemsettingcache.BooleanTTL = time.Hour
-	defer func() { systemsettingcache.BooleanTTL = oldTTL }()
+	t.Cleanup(func() {
+		systemsettingcache.BooleanTTL = oldTTL
+		systemsettingcache.Invalidate(domain.SettingKeyProxyRequestsDisabled)
+	})

Also applies to: 100-110, 113-116, 121-131

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@internal/service/system_setting_validation_test.go` around lines 92 - 95,
这些测试在修改 systemsettingcache.BooleanTTL 并写入全局缓存(键 "proxy_requests_disabled"),但只恢复了
TTL,没有清理缓存,导致跨测试污染。为每个相关测试(如
TestAdminServiceUpdateSettingInvalidatesProxyBooleanCache 及其他列出的用例)在
setup/teardown 或在 defer 中恢复 BooleanTTL 后也显式删除或清空该缓存键(例如调用相应的缓存清理方法,像
systemsettingcache.Delete("proxy_requests_disabled") 或
systemsettingcache.ClearKey("proxy_requests_disabled"),或使用 package 提供的全局 Clear
方法),确保每个测试结束时全局缓存回到干净状态以避免顺序依赖。


repo := &stubSystemSettingRepo{values: map[string]string{domain.SettingKeyProxyRequestsDisabled: "false"}}
svc := &AdminService{settingRepo: repo}

systemsettingcache.Invalidate(domain.SettingKeyProxyRequestsDisabled)
if systemsettingcache.GetBoolean(repo, domain.SettingKeyProxyRequestsDisabled) {
t.Fatal("expected initial cached value to be false")
}

if err := svc.UpdateSetting(domain.SettingKeyProxyRequestsDisabled, "true"); err != nil {
t.Fatalf("expected update to succeed, got %v", err)
}
if !systemsettingcache.GetBoolean(repo, domain.SettingKeyProxyRequestsDisabled) {
t.Fatal("expected cache invalidation to expose updated true value")
}
}

func TestAdminServiceDeleteSettingInvalidatesProxyBooleanCache(t *testing.T) {
oldTTL := systemsettingcache.BooleanTTL
systemsettingcache.BooleanTTL = time.Hour
defer func() { systemsettingcache.BooleanTTL = oldTTL }()

repo := &stubSystemSettingRepo{values: map[string]string{domain.SettingKeyProxyRequestsDisabled: "true"}}
svc := &AdminService{settingRepo: repo}

systemsettingcache.Invalidate(domain.SettingKeyProxyRequestsDisabled)
if !systemsettingcache.GetBoolean(repo, domain.SettingKeyProxyRequestsDisabled) {
t.Fatal("expected initial cached value to be true")
}

if err := svc.DeleteSetting(domain.SettingKeyProxyRequestsDisabled); err != nil {
t.Fatalf("expected delete to succeed, got %v", err)
}
if systemsettingcache.GetBoolean(repo, domain.SettingKeyProxyRequestsDisabled) {
t.Fatal("expected cache invalidation to expose deleted false value")
}
}

func TestBackupServiceImportSystemSettingsSkipsPayloadOverrideRules(t *testing.T) {
repo := &stubSystemSettingRepo{
values: map[string]string{},
Expand Down
Loading
Loading