From 6002b16cd0e5ca9cb84f27bdb2c9191f2c95ca64 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:20:59 +0800 Subject: [PATCH] fix: make repeated RemDial calls in one process safe A second RemDial in the same process must never wedge the caller or corrupt the first channel: - the FFI RemDial and the REM proxy client waited on agent Init with no exit condition; a HandlerInit failure (e.g. the local endpoint is already bound by another agent in the process) closes the agent without ever setting Init, hanging the caller forever. Both now wait via agent.WaitInit, which also fails on agent close and on timeout. - the handler-exit path deleted the agent's registry entry unconditionally, which could evict a same-alias successor registered after Close ran; it now uses CompareAndDelete. Add runner-level regression tests for the second-dial contract: a dial whose alias is still registered is refused with a controlled error instead of panicking in newAgent, and a re-dial after the first agent closes re-establishes the channel cleanly. Add WaitInit unit tests. --- agent/agent.go | 17 ++++++ agent/waitinit_test.go | 68 +++++++++++++++++++++++ cmd/export/export.go | 22 ++++++-- runner/redial_test.go | 123 +++++++++++++++++++++++++++++++++++++++++ runner/runner.go | 5 ++ runner/runner_std.go | 11 ++-- 6 files changed, 235 insertions(+), 11 deletions(-) create mode 100644 agent/waitinit_test.go create mode 100644 runner/redial_test.go 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",