diff --git a/server/listen/http_server/setup_security.go b/server/listen/http_server/setup_security.go index 81269f4..272b188 100644 --- a/server/listen/http_server/setup_security.go +++ b/server/listen/http_server/setup_security.go @@ -33,6 +33,7 @@ type setupAttempt struct { type setupSecurity struct { mu sync.Mutex + requestMu sync.Mutex token string tokenHash [32]byte consumed bool @@ -143,6 +144,13 @@ func (s *setupSecurity) authorize(next http.HandlerFunc) http.HandlerFunc { http.Error(w, "forbidden", http.StatusForbidden) return } + // 一次性令牌必须串行消费请求,避免两个并发请求在状态切换前同时进入初始化处理器。 + s.requestMu.Lock() + defer s.requestMu.Unlock() + if !s.validToken(token) { + http.Error(w, "forbidden", http.StatusForbidden) + return + } s.clearFailures(ip) next.ServeHTTP(w, r) if config.IsSetupReady() { diff --git a/server/listen/http_server/setup_security_test.go b/server/listen/http_server/setup_security_test.go index 0eed4af..deb8a12 100644 --- a/server/listen/http_server/setup_security_test.go +++ b/server/listen/http_server/setup_security_test.go @@ -4,6 +4,7 @@ import ( "net/http" "net/http/httptest" "os" + "sync" "testing" "git.pcux.cn/xuans/Smail/config" @@ -86,6 +87,46 @@ func TestSetupSecurityAcceptsTokenOnce(t *testing.T) { } } +func TestSetupSecuritySerializesConcurrentTokenRequests(t *testing.T) { + token := "0123456789abcdefghijklmnop" + security := withSetupSecurity(t, token) + started := make(chan struct{}) + continueRequest := make(chan struct{}) + var calls int + var callsMu sync.Mutex + handler := security.authorize(func(w http.ResponseWriter, r *http.Request) { + callsMu.Lock() + calls++ + callsMu.Unlock() + close(started) + <-continueRequest + config.IsInit = true + w.WriteHeader(http.StatusNoContent) + }) + + results := make(chan int, 2) + for i := 0; i < 2; i++ { + go func() { + request := setupRequest(http.MethodPost) + request.Header.Set("X-Setup-Token", token) + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, request) + results <- recorder.Code + }() + } + + <-started + close(continueRequest) + first, second := <-results, <-results + if (first != http.StatusNoContent && second != http.StatusNoContent) || (first == http.StatusNoContent && second == http.StatusNoContent) { + t.Fatalf("concurrent setup results = %d, %d; want exactly one success", first, second) + } + callsMu.Lock() + defer callsMu.Unlock() + if calls != 1 { + t.Fatalf("setup handler calls = %d, want 1", calls) + } +} func TestSetupSecurityBlocksUnauthorizedSource(t *testing.T) { token := "0123456789abcdefghijklmnop" security := withSetupSecurity(t, token) diff --git a/server/services/setup/finish.go b/server/services/setup/finish.go index 319adb6..924e454 100644 --- a/server/services/setup/finish.go +++ b/server/services/setup/finish.go @@ -4,14 +4,23 @@ import ( "git.pcux.cn/xuans/Smail/config" "git.pcux.cn/xuans/Smail/signal" "git.pcux.cn/xuans/Smail/utils/errors" + "sync" ) +var finishMu sync.Mutex + // Finish 标记初始化完成 func Finish() error { + finishMu.Lock() + defer finishMu.Unlock() cfg, err := config.ReadConfig() if err != nil { return errors.Wrap(err) } + if cfg.IsSetupReady() { + config.IsInit = true + return nil + } cfg.SetupState = config.SetupStateReady cfg.IsInit = true @@ -21,6 +30,10 @@ func Finish() error { } config.IsInit = true // 初始化完成 - signal.InitChan <- true + select { + case signal.InitChan <- true: + default: + // 初始化通知只需唤醒一次启动流程,避免重复完成请求阻塞响应。 + } return nil }