mirror of
https://github.com/vxcontrol/pentagi.git
synced 2026-08-23 19:46:39 +00:00
fix(docker): allocate free host ports for sandboxes instead of deriving them from the flow id
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 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5
parent
2d1e677e75
commit
eaf194cb89
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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" }
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user