diff --git a/agent/agent.go b/agent/agent.go index 4a4eb4a..3427281 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -288,6 +288,23 @@ func (agent *Agent) Handler() error { return agent.handleMessage() } +// WaitInit blocks until the agent finishes synchronous handler init, is +// closed, or the timeout elapses, and reports whether Init was reached. +// A HandlerInit failure closes the agent without ever setting Init, so +// waiting on Init alone would block the caller forever. +func (agent *Agent) WaitInit(timeout time.Duration) bool { + deadline := time.Now().Add(timeout) + for { + if agent.Init { + return true + } + if agent.IsClosed() || time.Now().After(deadline) { + return false + } + time.Sleep(50 * time.Millisecond) + } +} + // StartBackgroundLoops starts the monitor and stream acceptor goroutines. // Safe to call multiple times (uses sync.Once internally). func (agent *Agent) StartBackgroundLoops() { diff --git a/agent/waitinit_test.go b/agent/waitinit_test.go new file mode 100644 index 0000000..b5148c3 --- /dev/null +++ b/agent/waitinit_test.go @@ -0,0 +1,68 @@ +//go:build !tinygo + +package agent + +import ( + "testing" + "time" + + "github.com/chainreactors/rem/protocol/core" +) + +// TestWaitInitReturnsTrueAfterInit covers the success path: the handler +// finishes synchronous init and waiting callers unblock. +func TestWaitInitReturnsTrueAfterInit(t *testing.T) { + a, err := NewAgent(&Config{Alias: "waitinit-success", Type: core.CLIENT}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { a.Close(nil) }) + + go func() { + time.Sleep(100 * time.Millisecond) + a.Init = true + }() + if !a.WaitInit(5 * time.Second) { + t.Fatal("WaitInit should report true once Init is set") + } +} + +// TestWaitInitReturnsFalseWhenAgentClosesBeforeInit is the regression for +// the wedged RemDial wait: a HandlerInit failure closes the agent without +// ever setting Init, and callers waiting on Init alone would block forever. +func TestWaitInitReturnsFalseWhenAgentClosesBeforeInit(t *testing.T) { + a, err := NewAgent(&Config{Alias: "waitinit-closed", Type: core.CLIENT}) + if err != nil { + t.Fatal(err) + } + + go func() { + time.Sleep(100 * time.Millisecond) + a.Close(nil) + }() + start := time.Now() + if a.WaitInit(5 * time.Second) { + t.Fatal("WaitInit should report false for an agent closed before Init") + } + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("WaitInit blocked %v after the agent closed", elapsed) + } +} + +// TestWaitInitTimeout is the backstop: an agent that neither initializes +// nor closes fails the wait after the timeout instead of blocking forever. +func TestWaitInitTimeout(t *testing.T) { + a, err := NewAgent(&Config{Alias: "waitinit-timeout", Type: core.CLIENT}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { a.Close(nil) }) + + start := time.Now() + if a.WaitInit(300 * time.Millisecond) { + t.Fatal("WaitInit should report false on timeout") + } + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("WaitInit ignored the timeout, blocked %v", elapsed) + } +} diff --git a/cmd/export/export.go b/cmd/export/export.go index 8e914f6..6873b8a 100644 --- a/cmd/export/export.go +++ b/cmd/export/export.go @@ -7,6 +7,7 @@ package main import "C" import ( "context" + "fmt" "io" "math/rand" "net" @@ -23,6 +24,10 @@ import ( "github.com/kballard/go-shellquote" ) +// initWaitTimeout bounds how long RemDial waits for the dialed agent to +// finish synchronous handler init before failing the call. +const initWaitTimeout = 30 * time.Second + func init() { utils.Log = logs.NewLogger(100) } @@ -84,19 +89,24 @@ func RemDial(cmdline *C.char) (*C.char, C.int) { utils.Log.Errorf("[rem] handler exited: %v", err) } a.Close(err) - agent.Agents.Map.Delete(a.ID) + // CompareAndDelete (not an unconditional Delete): a same-alias agent + // from a later RemDial may have registered after Close ran, and its + // registry entry must not be evicted. + agent.Agents.CompareAndDelete(a.ID, a) // Do NOT reconnect here — the caller (malefic Rust) manages // reconnection by re-calling RemDial. Having two reconnection // loops (Go + Rust) causes competing agents that kill each // other's memory listeners via CreateChannel race. }() - for { - if a.Init { - break - } else { - time.Sleep(100 * time.Millisecond) + // Wait for synchronous init, but never wait forever: a HandlerInit + // failure (e.g. the local endpoint is already bound by another agent + // in this process) closes the agent without setting Init. + if !a.WaitInit(initWaitTimeout) { + if !a.IsClosed() { + a.Close(fmt.Errorf("agent init timeout")) } + return nil, ErrDialFailed } return C.CString(a.ID), 0 diff --git a/runner/redial_test.go b/runner/redial_test.go new file mode 100644 index 0000000..6045bad --- /dev/null +++ b/runner/redial_test.go @@ -0,0 +1,123 @@ +//go:build !tinygo + +package runner + +import ( + "fmt" + "net" + "testing" + "time" + + "github.com/chainreactors/rem/agent" + "github.com/chainreactors/rem/protocol/core" + _ "github.com/chainreactors/rem/protocol/serve/socks" + _ "github.com/chainreactors/rem/protocol/tunnel/tcp" +) + +// dialWithAlias runs the client side of the FFI RemDial flow: build a fresh +// console from the command line and dial the server console with a fixed +// agent alias. +func dialWithAlias(t *testing.T, link, alias string) (*Console, *agent.Agent, error) { + t.Helper() + client, err := NewConsoleWithCMD(fmt.Sprintf("-c %s -l socks5://127.0.0.1:0 -a %s --no-sub", link, alias)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = client.Close() }) + a, err := client.Dial(client.ConsoleURL) + return client, a, err +} + +func runHandler(t *testing.T, a *agent.Agent) { + t.Helper() + go func() { _ = a.Handler() }() + if !a.WaitInit(10 * time.Second) { + t.Fatalf("agent %s failed to initialize", a.ID) + } + t.Cleanup(func() { a.Close(nil) }) +} + +// TestSecondDialWithLiveAliasIsRefusedWithoutPanic reproduces the production +// second-rem_dial scenario: a second dial in the same process while the +// first agent is still registered under the same alias. Before the console +// lifecycle fixes, the ignored agent.NewAgent error panicked with a nil +// pointer dereference in newAgent and killed the whole process. The second +// dial must instead fail with a controlled error and leave the first +// channel untouched. +func TestSecondDialWithLiveAliasIsRefusedWithoutPanic(t *testing.T) { + // The first rem_dial's channel: an incumbent agent registered under the + // alias and still alive. The refusal happens in newAgent before any + // network I/O, so no server is needed. + incumbent, err := agent.NewAgent(&agent.Config{Alias: "redial-live-alias", Type: core.CLIENT}) + if err != nil { + t.Fatal(err) + } + if err := agent.Agents.Add(incumbent); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { incumbent.Close(nil); agent.Agents.CompareAndDelete(incumbent.ID, incumbent) }) + + client, err := NewConsoleWithCMD("-c tcp://nonenonenonenone@127.0.0.1:1/?wrapper=raw -l socks5://127.0.0.1:0 -a redial-live-alias --no-sub") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = client.Close() }) + + a, err := client.Dial(client.ConsoleURL) + if err == nil { + if a != nil { + a.Close(nil) + } + t.Fatal("second dial with a live alias should have been refused") + } + if a != nil { + t.Fatal("refused dial returned a non-nil agent") + } + if incumbent.IsClosed() { + t.Fatal("refused second dial killed the incumbent agent") + } + if current, ok := agent.Agents.Get(incumbent.ID); !ok || current != incumbent { + t.Fatal("incumbent agent lost its registry entry after refused dial") + } +} + +// TestDialWithAliasReestablishesAfterClose covers the other half of the +// second-dial contract: once the first agent is closed (handler exit or +// explicit cleanup), a re-dial with the same alias must succeed and fully +// re-initialize. +// +// The server console runs in a subprocess because the agent registry is +// process-global: an in-process server would register its server-side mirror +// agent under the same alias as the client agent and the two would collide. +func TestDialWithAliasReestablishesAfterClose(t *testing.T) { + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + port := l.Addr().(*net.TCPAddr).Port + if err := l.Close(); err != nil { + t.Fatal(err) + } + startHelper(t, fmt.Sprintf("-s tcp://redialkey@127.0.0.1:%d/?wrapper=raw --no-sub", port)) + waitForTCP(t, fmt.Sprintf("127.0.0.1:%d", port), 10*time.Second) + link := fmt.Sprintf("tcp://redialkey@127.0.0.1:%d/?wrapper=raw", port) + + _, first, err := dialWithAlias(t, link, "redial-reestablish-alias") + if err != nil { + t.Fatalf("first dial: %v", err) + } + runHandler(t, first) + first.Close(nil) + + _, second, err := dialWithAlias(t, link, "redial-reestablish-alias") + if err != nil { + t.Fatalf("re-dial after close: %v", err) + } + runHandler(t, second) + if second.ID != first.ID { + t.Fatalf("re-dialed agent ID = %q, want the alias %q", second.ID, first.ID) + } + if current, ok := agent.Agents.Get(second.ID); !ok || current != second { + t.Fatal("re-dialed agent is not the registered one") + } +} diff --git a/runner/runner.go b/runner/runner.go index f21eaf1..4d08f0d 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -6,6 +6,7 @@ import ( "os" "runtime/debug" "sync" + "time" "github.com/chainreactors/rem/agent" "github.com/chainreactors/rem/protocol/core" @@ -14,6 +15,10 @@ import ( "github.com/kballard/go-shellquote" ) +// initWaitTimeout bounds how long client-side callers wait for a dialed +// agent to finish synchronous handler init before failing. +const initWaitTimeout = 30 * time.Second + func NewConsoleWithCMD(cmdline string) (*Console, error) { var option Options args, err := shellquote.Split(cmdline) diff --git a/runner/runner_std.go b/runner/runner_std.go index 5682921..7f88a55 100644 --- a/runner/runner_std.go +++ b/runner/runner_std.go @@ -12,7 +12,6 @@ import ( "net/url" "os/user" "strings" - "time" "github.com/chainreactors/logs" "github.com/chainreactors/proxyclient" @@ -46,11 +45,13 @@ func newRemProxyClient(proxyURL *url.URL, _ proxyclient.Dial) (proxyclient.Dial, } }() - for { - if a.Init { - break + // Never wait forever: a HandlerInit failure closes the agent + // without setting Init. + if !a.WaitInit(initWaitTimeout) { + if !a.IsClosed() { + a.Close(fmt.Errorf("agent init timeout")) } - time.Sleep(100 * time.Millisecond) + return nil, fmt.Errorf("rem agent failed to initialize") } memURL := &url.URL{ Scheme: "memory",