diff --git a/internal/agent/agent.go b/internal/agent/agent.go index f8333e3..e296620 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -8,6 +8,7 @@ import ( "bytes" "crypto/tls" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -80,6 +81,18 @@ func New(cfg *Config) *Agent { } } +// statusError is a non-200 response from the master. It keeps the status +// code addressable so the registration retry loop can tell a wrong-path +// 404 (master_url missing the panel base path) apart from other failures. +type statusError struct { + code int + body string +} + +func (e *statusError) Error() string { + return fmt.Sprintf("register failed with status %d: %s", e.code, e.body) +} + // Register registers the agent with the master func (a *Agent) Register() error { url := fmt.Sprintf("%s/panel/api/agent/register", a.config.MasterURL) @@ -103,21 +116,83 @@ func (a *Agent) Register() error { if resp.StatusCode != http.StatusOK { data, _ := io.ReadAll(resp.Body) - return fmt.Errorf("register failed with status %d: %s", resp.StatusCode, string(data)) + return &statusError{code: resp.StatusCode, body: string(data)} } logger.Info("Agent registered successfully with master") return nil } +// registrationBackoff returns the delay before the next registration +// attempt: 5s, 15s, 30s, then capped at 60s. +func registrationBackoff(failures int) time.Duration { + switch { + case failures <= 1: + return 5 * time.Second + case failures == 2: + return 15 * time.Second + case failures == 3: + return 30 * time.Second + default: + return 60 * time.Second + } +} + +// registerWait sleeps between registration attempts, returning false when +// the agent is stopping. A package var so tests can drive the retry loop +// without real delays. +var registerWait = func(stopCh chan struct{}, d time.Duration) bool { + select { + case <-stopCh: + return false + case <-time.After(d): + return true + } +} + +// errRegistrationAborted marks a registration loop cut short by Stop. +var errRegistrationAborted = errors.New("registration aborted: agent stopping") + +// registerWithRetry keeps trying to register until the master accepts or +// the agent is stopped. A one-shot failure used to kill the agent +// outright — a typo'd master_url (classically: the panel base path +// missing, so the master answers 404) left the node dead and "unknown" +// on the panel with nothing retrying. Registration is idempotent on the +// master, so retrying with backoff is safe. +func (a *Agent) registerWithRetry() error { + failures := 0 + for { + err := a.Register() + if err == nil { + return nil + } + failures++ + var st *statusError + if errors.As(err, &st) && st.code == http.StatusNotFound { + logger.Warningf("主控返回 404:检查 master_url 是否包含面板基础路径(如 https://host/rogue),当前 master_url=%s", a.config.MasterURL) + } + delay := registrationBackoff(failures) + logger.Warningf("Registration attempt %d failed: %v; retrying in %s", failures, err, delay) + if !registerWait(a.stopCh, delay) { + return errRegistrationAborted + } + } +} + // Run starts the agent main loop func (a *Agent) Run() error { logger.Info("Starting 3x-ui agent...") logger.Infof("Master: %s, Node: %s", a.config.MasterURL, a.config.NodeName) - // Initial registration - if err := a.Register(); err != nil { - return fmt.Errorf("initial registration failed: %w", err) + // Initial registration, retried with backoff: a node that cannot + // register yet (master down, URL just fixed, network still coming up) + // must not die on the first failure. + if err := a.registerWithRetry(); err != nil { + if errors.Is(err, errRegistrationAborted) { + logger.Info("Agent stopping...") + return nil + } + return err } // Start heartbeat ticker diff --git a/internal/agent/agent_register_test.go b/internal/agent/agent_register_test.go new file mode 100644 index 0000000..84a6e8d --- /dev/null +++ b/internal/agent/agent_register_test.go @@ -0,0 +1,135 @@ +package agent + +import ( + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/mhsanaei/3x-ui/v3/internal/logger" +) + +// stubRegisterWait replaces the retry sleep with a recorder that never +// actually sleeps. allowWaits is how many waits return true (keep +// retrying) before returning false (agent stopping). +func stubRegisterWait(t *testing.T, allowWaits int) *[]time.Duration { + t.Helper() + delays := &[]time.Duration{} + orig := registerWait + registerWait = func(_ chan struct{}, d time.Duration) bool { + *delays = append(*delays, d) + return len(*delays) <= allowWaits + } + t.Cleanup(func() { registerWait = orig }) + return delays +} + +func captureLogs(t *testing.T) *[]string { + t.Helper() + var mu sync.Mutex + lines := &[]string{} + prev := logger.SetLogHook(func(_, msg string) { + mu.Lock() + *lines = append(*lines, msg) + mu.Unlock() + }) + t.Cleanup(func() { logger.SetLogHook(prev) }) + return lines +} + +// scriptedMaster answers register requests from a status-code script; +// once the script is exhausted it answers 200. +func scriptedMaster(t *testing.T, script []int, hits *int) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/panel/api/agent/register") { + w.WriteHeader(http.StatusNotFound) + return + } + *hits++ + code := http.StatusOK + if *hits <= len(script) { + code = script[*hits-1] + } + w.WriteHeader(code) + fmt.Fprint(w, `{}`) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestRegistrationBackoffSequence(t *testing.T) { + want := []time.Duration{5, 15, 30, 60, 60, 60} + for i, w := range want { + if got := registrationBackoff(i + 1); got != w*time.Second { + t.Errorf("registrationBackoff(%d) = %s, want %ss", i+1, got, w) + } + } +} + +func TestRegisterWithRetrySucceedsAfter404s(t *testing.T) { + hits := 0 + srv := scriptedMaster(t, []int{404, 404}, &hits) + delays := stubRegisterWait(t, 10) + logs := captureLogs(t) + + ag := New(&Config{MasterURL: srv.URL, Token: "tok", NodeName: "n1"}) + if err := ag.registerWithRetry(); err != nil { + t.Fatalf("registerWithRetry: %v", err) + } + if hits != 3 { + t.Fatalf("register hits = %d, want 3 (404, 404, 200)", hits) + } + if len(*delays) != 2 || (*delays)[0] != 5*time.Second || (*delays)[1] != 15*time.Second { + t.Fatalf("retry delays = %v, want [5s 15s]", *delays) + } + joined := strings.Join(*logs, "\n") + if !strings.Contains(joined, "检查 master_url 是否包含面板基础路径") { + t.Errorf("404 hint missing from logs:\n%s", joined) + } + if !strings.Contains(joined, "retrying in 5s") { + t.Errorf("retry log line missing:\n%s", joined) + } +} + +func TestRegisterWithRetryAbortsOnStop(t *testing.T) { + hits := 0 + srv := scriptedMaster(t, []int{500, 500, 500, 500, 500, 500, 500, 500}, &hits) + // First wait says "stopping": the loop must give up immediately + // instead of burning through the script. + delays := stubRegisterWait(t, 0) + + ag := New(&Config{MasterURL: srv.URL, Token: "tok", NodeName: "n1"}) + err := ag.registerWithRetry() + if !errors.Is(err, errRegistrationAborted) { + t.Fatalf("registerWithRetry err = %v, want errRegistrationAborted", err) + } + if hits != 1 { + t.Fatalf("register hits = %d, want 1 (no retry after stop)", hits) + } + if len(*delays) != 1 { + t.Fatalf("waits = %d, want 1", len(*delays)) + } +} + +func TestRegisterStatusErrorCarriesCode(t *testing.T) { + hits := 0 + srv := scriptedMaster(t, []int{404}, &hits) + ag := New(&Config{MasterURL: srv.URL, Token: "tok", NodeName: "n1"}) + err := ag.Register() + var st *statusError + if !errors.As(err, &st) { + t.Fatalf("Register err = %v (%T), want *statusError", err, err) + } + if st.code != http.StatusNotFound { + t.Fatalf("statusError.code = %d, want 404", st.code) + } + // Error text must stay the historical one — scripts may grep for it. + if !strings.Contains(err.Error(), "register failed with status 404") { + t.Fatalf("error text changed: %q", err.Error()) + } +}