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 }