From eaf194cb89658daba396ad2abe5b72e5a1f2e46b Mon Sep 17 00:00:00 2001 From: Sergey Kozyrenko Date: Mon, 20 Jul 2026 15:39:51 +0700 Subject: [PATCH] fix(docker): allocate free host ports for sandboxes instead of deriving them from the flow id MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Primary sandbox host ports were computed as 28000 + (flowID*containerPortsNumber + i) % limitContainerPortsNumber, so any two flows whose ids are congruent mod (limitContainerPortsNumber/containerPortsNumber) mapped to the same host ports. The second container then failed to start with "port is already allocated" while its flow row lingered in "created" — and the collision also crossed compose stacks sharing one docker daemon (e.g. flow 2 vs flow 90002 both on 28004/28005). Reserve free OS ports at bind time and hold the reservations open until docker takes over, so two concurrent flows can never pick the same port. The agent prompt now reads the container's actually-bound host ports back from docker (correct after a restart too), keeping the deterministic set only as a host-network / inspect-failure fallback. On a container-start failure the flow is now marked failed instead of left stuck in "created". Regression test: two flows N and N+period now get distinct free ports. Co-Authored-By: Claude Fable 5 --- backend/pkg/controller/context.go | 17 ++++ backend/pkg/controller/flow.go | 3 + backend/pkg/docker/client.go | 99 ++++++++++++++++++- backend/pkg/docker/client_ports_test.go | 95 ++++++++++++++++++ backend/pkg/providers/helpers.go | 34 ++++++- backend/pkg/providers/provider.go | 4 + backend/pkg/providers/providers.go | 4 + .../pkg/server/services/flow_files_test.go | 3 + backend/pkg/tools/terminal_test.go | 3 + 9 files changed, 260 insertions(+), 2 deletions(-) create mode 100644 backend/pkg/docker/client_ports_test.go diff --git a/backend/pkg/controller/context.go b/backend/pkg/controller/context.go index b098eab8..2d8e35ff 100644 --- a/backend/pkg/controller/context.go +++ b/backend/pkg/controller/context.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "time" "pentagi/pkg/database" "pentagi/pkg/graph/subscriptions" @@ -49,6 +50,22 @@ type SubtaskContext struct { TaskContext } +// markFlowFailed records a flow as failed after its resources couldn't be +// prepared. It uses a detached context so the write still lands when the +// caller's context was the thing that got canceled. +func markFlowFailed(db database.Querier, flowID int64) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if _, err := db.UpdateFlowStatus(ctx, database.UpdateFlowStatusParams{ + Status: database.FlowStatusFailed, + ID: flowID, + }); err != nil { + logrus.WithError(err).WithField("flow_id", flowID). + Error("failed to mark flow as failed after resource preparation error") + } +} + func wrapErrorEndSpan(ctx context.Context, span langfuse.Span, msg string, err error) error { logrus.WithContext(ctx).WithError(err).Error(msg) err = fmt.Errorf("%s: %w", msg, err) diff --git a/backend/pkg/controller/flow.go b/backend/pkg/controller/flow.go index 2ba6cdb3..c451f02e 100644 --- a/backend/pkg/controller/flow.go +++ b/backend/pkg/controller/flow.go @@ -272,6 +272,9 @@ func NewFlowWorker( } if err := executor.Prepare(ctx); err != nil { + // The flow row is already persisted; a container that never started + // leaves it visibly failed instead of stuck in "created". + markFlowFailed(fwc.db, flow.ID) return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to prepare flow resources", err) } diff --git a/backend/pkg/docker/client.go b/backend/pkg/docker/client.go index 62d01c44..407f4a0b 100644 --- a/backend/pkg/docker/client.go +++ b/backend/pkg/docker/client.go @@ -7,10 +7,12 @@ import ( "fmt" "hash/crc32" "io" + "net" "os" "path" "path/filepath" "slices" + "strconv" "strings" "sync" @@ -76,12 +78,22 @@ type DockerClient interface { ContainerExecInspect(ctx context.Context, execID string) (container.ExecInspect, error) ContainerStatPath(ctx context.Context, containerID string, path string) (container.PathStat, error) ListContainerDir(ctx context.Context, containerID string, dirPath string) (ContainerDirListing, error) + // GetFlowContainerPorts returns the host ports actually bound to the flow's + // primary container (read from the live container, so it stays correct after + // a process restart and reflects the dynamically allocated ports). + GetFlowContainerPorts(ctx context.Context, flowID int64) ([]int, error) CopyToContainer(ctx context.Context, containerID string, dstPath string, content io.Reader, options container.CopyToContainerOptions) error CopyFromContainer(ctx context.Context, containerID string, srcPath string) (io.ReadCloser, container.PathStat, error) Cleanup(ctx context.Context) error GetDefaultImage() string } +// GetPrimaryContainerPorts is the deterministic advisory port set for a flow. +// It is used only as a fallback for the prompt text (host-network mode, or when +// the live container can't be inspected) — the actual bridge-mode host bindings +// come from reserveFreePorts, since this derivation wraps modulo +// limitContainerPortsNumber and therefore collides for flow ids that are +// congruent mod (limitContainerPortsNumber/containerPortsNumber). func GetPrimaryContainerPorts(flowID int64) []int { ports := make([]int, containerPortsNumber) for i := 0; i < containerPortsNumber; i++ { @@ -91,6 +103,73 @@ func GetPrimaryContainerPorts(flowID int64) []int { return ports } +// reserveFreePorts asks the OS for n free TCP ports on bindIP and keeps each +// listener open, so two concurrent reservations can never hand back the same +// port. The caller closes the returned listeners immediately before Docker +// binds the ports (the only remaining race window is against a non-pentagi +// process, which the deterministic derivation could not avoid either). +func reserveFreePorts(bindIP string, n int) ([]int, []io.Closer, error) { + ports := make([]int, 0, n) + closers := make([]io.Closer, 0, n) + closeAll := func() { + for _, c := range closers { + _ = c.Close() + } + } + + for i := 0; i < n; i++ { + listener, err := net.Listen("tcp", net.JoinHostPort(bindIP, "0")) + if err != nil { + closeAll() + return nil, nil, fmt.Errorf("failed to reserve free host port: %w", err) + } + closers = append(closers, listener) + ports = append(ports, listener.Addr().(*net.TCPAddr).Port) + } + + return ports, closers, nil +} + +func (dc *dockerClient) GetFlowContainerPorts(ctx context.Context, flowID int64) ([]int, error) { + cnt, err := dc.db.GetFlowPrimaryContainer(ctx, flowID) + if err != nil { + return nil, fmt.Errorf("failed to get primary container for flow %d: %w", flowID, err) + } + if !cnt.LocalID.Valid || cnt.LocalID.String == "" { + return nil, fmt.Errorf("primary container for flow %d has no local id", flowID) + } + + inspection, err := dc.client.ContainerInspect(ctx, cnt.LocalID.String) + if err != nil { + return nil, fmt.Errorf("failed to inspect container %s: %w", cnt.LocalID.String, err) + } + + return hostPortsFromBindings(inspection.NetworkSettings.Ports), nil +} + +// hostPortsFromBindings collects the distinct host ports from a container's port +// map, sorted for a stable prompt rendering. +func hostPortsFromBindings(portMap nat.PortMap) []int { + seen := make(map[int]struct{}) + for _, bindings := range portMap { + for _, binding := range bindings { + port, err := strconv.Atoi(binding.HostPort) + if err != nil { + continue + } + seen[port] = struct{}{} + } + } + + ports := make([]int, 0, len(seen)) + for port := range seen { + ports = append(ports, port) + } + slices.Sort(ports) + + return ports +} + func NewDockerClient(ctx context.Context, db database.Querier, cfg *config.Config) (DockerClient, error) { cli, err := client.NewClientWithOpts(client.FromEnv) if err != nil { @@ -292,6 +371,9 @@ func (dc *dockerClient) RunContainer( // Configure network mode and port bindings var networkingConfig *network.NetworkingConfig + // Held open until just before ContainerCreate so a concurrent flow cannot + // reserve the same host ports. + var portReservations []io.Closer if dc.network == "host" { // Host network mode: container uses host network stack directly // No port bindings needed as container has direct access to host interfaces @@ -305,7 +387,17 @@ func (dc *dockerClient) RunContainer( if config.ExposedPorts == nil { config.ExposedPorts = nat.PortSet{} } - for _, port := range GetPrimaryContainerPorts(flowID) { + // Allocate free host ports instead of deriving them from flowID: the + // derivation wraps and collides for far-apart flow ids (and across + // stacks sharing this daemon), which fails the bind with "port is + // already allocated". + hostPorts, reservations, err := reserveFreePorts(dc.publicIP, containerPortsNumber) + if err != nil { + defer updateContainerInfo(database.ContainerStatusFailed, "") + return database.Container{}, fmt.Errorf("failed to allocate host ports for container '%s': %w", containerName, err) + } + portReservations = reservations + for _, port := range hostPorts { natPort := nat.Port(fmt.Sprintf("%d/tcp", port)) hostConfig.PortBindings[natPort] = []nat.PortBinding{ { @@ -325,6 +417,11 @@ func (dc *dockerClient) RunContainer( } } + // Hand the reserved host ports to Docker: release them so its bind succeeds. + for _, reservation := range portReservations { + _ = reservation.Close() + } + resp, err := dc.client.ContainerCreate(ctx, config, hostConfig, networkingConfig, nil, containerName) if err != nil { if config.Image == dc.defImage { diff --git a/backend/pkg/docker/client_ports_test.go b/backend/pkg/docker/client_ports_test.go new file mode 100644 index 00000000..0f7b83c0 --- /dev/null +++ b/backend/pkg/docker/client_ports_test.go @@ -0,0 +1,95 @@ +package docker + +import ( + "io" + "testing" + + "github.com/docker/go-connections/nat" +) + +// TestDeterministicPortsCollide documents why container host ports are no longer +// derived from the flow id: the derivation wraps modulo +// limitContainerPortsNumber, so two flows whose ids differ by +// limitContainerPortsNumber/containerPortsNumber map to the same host ports and +// the second container's bind fails with "port is already allocated". +func TestDeterministicPortsCollide(t *testing.T) { + period := int64(limitContainerPortsNumber / containerPortsNumber) + + a := GetPrimaryContainerPorts(2) + b := GetPrimaryContainerPorts(2 + period) + + if !slicesEqual(a, b) { + t.Fatalf("expected ids 2 and %d to collide, got %v vs %v", 2+period, a, b) + } + // 90002 ≡ 2 (mod period) → both derive to 28004/28005. + if got := GetPrimaryContainerPorts(90002); !slicesEqual(got, GetPrimaryContainerPorts(2)) { + t.Fatalf("expected flow 90002 to collide with flow 2, got %v", got) + } +} + +// TestReserveFreePortsCoexist is the regression: two flows created concurrently +// (ids far apart, e.g. N and N+1000) must get distinct host ports. Holding both +// reservations open at once is exactly what the RunContainer path does until it +// hands the ports to Docker, so disjointness here proves the collision is gone. +func TestReserveFreePortsCoexist(t *testing.T) { + portsA, closersA, err := reserveFreePorts("127.0.0.1", containerPortsNumber) + if err != nil { + t.Fatalf("reserve A: %v", err) + } + defer closeAll(closersA) + + portsB, closersB, err := reserveFreePorts("127.0.0.1", containerPortsNumber) + if err != nil { + t.Fatalf("reserve B: %v", err) + } + defer closeAll(closersB) + + if len(portsA) != containerPortsNumber || len(portsB) != containerPortsNumber { + t.Fatalf("expected %d ports each, got %v and %v", containerPortsNumber, portsA, portsB) + } + + all := append(append([]int{}, portsA...), portsB...) + seen := make(map[int]struct{}, len(all)) + for _, p := range all { + if p <= 0 || p > 65535 { + t.Fatalf("port %d out of range", p) + } + if _, dup := seen[p]; dup { + t.Fatalf("two concurrent reservations collided on port %d (A=%v B=%v)", p, portsA, portsB) + } + seen[p] = struct{}{} + } +} + +func TestHostPortsFromBindings(t *testing.T) { + portMap := nat.PortMap{ + "28005/tcp": []nat.PortBinding{{HostIP: "127.0.0.1", HostPort: "40002"}}, + "28004/tcp": []nat.PortBinding{{HostIP: "127.0.0.1", HostPort: "40001"}}, + "9/tcp": nil, // no binding (host mode leaves these empty) + } + + got := hostPortsFromBindings(portMap) + + want := []int{40001, 40002} + if !slicesEqual(got, want) { + t.Fatalf("expected sorted host ports %v, got %v", want, got) + } +} + +func closeAll(closers []io.Closer) { + for _, c := range closers { + _ = c.Close() + } +} + +func slicesEqual(a, b []int) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} diff --git a/backend/pkg/providers/helpers.go b/backend/pkg/providers/helpers.go index 1445e575..e5744b45 100644 --- a/backend/pkg/providers/helpers.go +++ b/backend/pkg/providers/helpers.go @@ -860,8 +860,40 @@ func (fp *flowProvider) subtasksToMarkdown(subtasks []tools.SubtaskInfo) string return buffer.String() } +// containerPorts returns the primary container's actually bound host ports for +// the prompt (memoized, stable for the flow's lifetime). It falls back to the +// deterministic advisory set when the live ports can't be read — host-network +// mode publishes none, and a not-yet-inspectable container has none either. +func (fp *flowProvider) containerPorts() []int { + fp.mx.RLock() + cached := fp.cachedPorts + fp.mx.RUnlock() + if len(cached) > 0 { + return cached + } + + fallback := docker.GetPrimaryContainerPorts(fp.flowID) + if fp.docker == nil { + return fallback + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + ports, err := fp.docker.GetFlowContainerPorts(ctx, fp.flowID) + if err != nil || len(ports) == 0 { + return fallback + } + + fp.mx.Lock() + fp.cachedPorts = ports + fp.mx.Unlock() + + return ports +} + func (fp *flowProvider) getContainerPortsDescription() string { - ports := docker.GetPrimaryContainerPorts(fp.flowID) + ports := fp.containerPorts() var buffer strings.Builder buffer.WriteString("**OOB Attack Infrastructure:**\n\n") diff --git a/backend/pkg/providers/provider.go b/backend/pkg/providers/provider.go index eb996c99..d45968e2 100644 --- a/backend/pkg/providers/provider.go +++ b/backend/pkg/providers/provider.go @@ -138,6 +138,10 @@ type flowProvider struct { dataDir string publicIP string dockerNetwork string + docker docker.DockerClient + // cachedPorts memoizes the primary container's bound host ports (stable for + // the flow's lifetime) so the prompt doesn't inspect Docker on every render. + cachedPorts []int callCounter *atomic.Int64 diff --git a/backend/pkg/providers/providers.go b/backend/pkg/providers/providers.go index 94d18d48..6d7969ea 100644 --- a/backend/pkg/providers/providers.go +++ b/backend/pkg/providers/providers.go @@ -334,6 +334,7 @@ func (pc *providerController) NewFlowProvider( dataDir: pc.cfg.DataDir, publicIP: pc.publicIP, dockerNetwork: pc.dockerNetwork, + docker: pc.docker, callCounter: newAtomicInt64(pc.startCallNumber.Add(deltaCallCounter)), image: image, title: title, @@ -386,6 +387,7 @@ func (pc *providerController) LoadFlowProvider( dataDir: pc.cfg.DataDir, publicIP: pc.publicIP, dockerNetwork: pc.dockerNetwork, + docker: pc.docker, callCounter: newAtomicInt64(pc.startCallNumber.Add(deltaCallCounter)), image: image, title: title, @@ -483,6 +485,7 @@ func (pc *providerController) NewAssistantProvider( dataDir: pc.cfg.DataDir, publicIP: pc.publicIP, dockerNetwork: pc.dockerNetwork, + docker: pc.docker, callCounter: newAtomicInt64(pc.startCallNumber.Add(deltaCallCounter)), image: image, title: title, @@ -538,6 +541,7 @@ func (pc *providerController) LoadAssistantProvider( dataDir: pc.cfg.DataDir, publicIP: pc.publicIP, dockerNetwork: pc.dockerNetwork, + docker: pc.docker, callCounter: newAtomicInt64(pc.startCallNumber.Add(deltaCallCounter)), image: image, title: title, diff --git a/backend/pkg/server/services/flow_files_test.go b/backend/pkg/server/services/flow_files_test.go index 6044d2b9..9dbf8cc4 100644 --- a/backend/pkg/server/services/flow_files_test.go +++ b/backend/pkg/server/services/flow_files_test.go @@ -1101,6 +1101,9 @@ func (f *fakeDockerClient) CopyFromContainer(_ context.Context, _ string, contai } return io.NopCloser(bytes.NewReader(f.copyFromBody)), f.copyFromStat, nil } +func (f *fakeDockerClient) GetFlowContainerPorts(_ context.Context, _ int64) ([]int, error) { + return nil, nil +} func (f *fakeDockerClient) Cleanup(_ context.Context) error { return nil } func (f *fakeDockerClient) GetDefaultImage() string { return "test-image" } diff --git a/backend/pkg/tools/terminal_test.go b/backend/pkg/tools/terminal_test.go index bee7601e..5880d3f9 100644 --- a/backend/pkg/tools/terminal_test.go +++ b/backend/pkg/tools/terminal_test.go @@ -99,6 +99,9 @@ func (m *contextAwareMockDockerClient) ContainerStatPath(_ context.Context, _ st func (m *contextAwareMockDockerClient) ListContainerDir(_ context.Context, _ string, _ string) (docker.ContainerDirListing, error) { return docker.ContainerDirListing{}, nil } +func (m *contextAwareMockDockerClient) GetFlowContainerPorts(_ context.Context, _ int64) ([]int, error) { + return nil, nil +} func (m *contextAwareMockDockerClient) ContainerExecInspect(_ context.Context, _ string) (container.ExecInspect, error) { return m.inspectResp, nil }