fix: make repeated RemDial calls in one process safe

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() {

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)
+	}
+}

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

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")
+	}
+}

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)

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",