Wait for proxy readiness before starting updater

internal/infra/proxy.go

@@ -3,12 +3,14 @@ package infra
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log"
 	"os"
 	"path"
 	"path/filepath"
+	"time"
 
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/mount"
@@ -20,6 +22,9 @@ import (
 )
 
 const proxyCertPath = "/usr/local/share/ca-certificates/custom-ca-cert.crt"
+const proxyReadyPort = 1080
+const proxyReadyTimeout = 60 * time.Second
+const proxyReadyPollInterval = 100 * time.Millisecond
 
 // ProxyImageName is the default Docker image used by the proxy
 const ProxyImageName = "ghcr.io/dependabot/proxy:latest"
@@ -132,6 +137,26 @@ func NewProxy(ctx context.Context, cli *client.Client, params *RunParams, nets *
 	return proxy, nil
 }
 
+func (p *Proxy) WaitUntilReady(ctx context.Context) error {
+	readyCtx, cancel := context.WithTimeout(ctx, proxyReadyTimeout)
+	defer cancel()
+
+	if err := waitForPortUntil(readyCtx, proxyReadyPollInterval, func(ctx context.Context) (bool, error) {
+		return isPortListening(ctx, p.cli, p.containerID, proxyReadyPort)
+	}); err != nil {
+		return proxyReadinessError(ctx, err)
+	}
+	return nil
+}
+
+func proxyReadinessError(ctx context.Context, err error) error {
+	format := "proxy did not start listening on port %d within %s: "
+	if errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil {
+		return fmt.Errorf(format+"%v", proxyReadyPort, proxyReadyTimeout, err)
+	}
+	return fmt.Errorf(format+"%w", proxyReadyPort, proxyReadyTimeout, err)
+}
+
 // proxyEnv builds the environment variables passed to the proxy container.
 func proxyEnv(apiURL string) []string {
 	env := []string{

internal/infra/proxy_test.go

@@ -1,9 +1,12 @@
 package infra
 
 import (
+	"context"
+	"errors"
 	"os"
 	"strings"
 	"testing"
+	"time"
 )
 
 func envValue(env []string, key string) (string, bool) {
@@ -161,3 +164,83 @@ func Test_proxyEnv_DependabotAPIURL(t *testing.T) {
 		}
 	})
 }
+
+func Test_waitForPortUntil(t *testing.T) {
+	t.Run("returns when a connection succeeds", func(t *testing.T) {
+		attempts := 0
+		err := waitForPortUntil(t.Context(), time.Millisecond, func(context.Context) (bool, error) {
+			attempts++
+			return attempts == 3, nil
+		})
+
+		if err != nil {
+			t.Fatalf("waitForPortUntil returned unexpected error: %v", err)
+		}
+		if attempts != 3 {
+			t.Fatalf("expected 3 connection attempts, got %d", attempts)
+		}
+	})
+
+	t.Run("returns probe errors", func(t *testing.T) {
+		probeErr := errors.New("docker exec failed")
+		err := waitForPortUntil(t.Context(), time.Millisecond, func(context.Context) (bool, error) {
+			return false, probeErr
+		})
+
+		if !errors.Is(err, probeErr) {
+			t.Fatalf("expected probe error, got %v", err)
+		}
+	})
+
+	t.Run("returns when the context is cancelled", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(t.Context())
+		attempts := 0
+		err := waitForPortUntil(ctx, time.Hour, func(context.Context) (bool, error) {
+			attempts++
+			cancel()
+			return false, nil
+		})
+
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("expected context cancellation, got %v", err)
+		}
+		if attempts != 1 {
+			t.Fatalf("expected 1 connection attempt, got %d", attempts)
+		}
+	})
+}
+
+func Test_proxyReadinessError(t *testing.T) {
+	t.Run("does not expose the internal readiness deadline", func(t *testing.T) {
+		err := proxyReadinessError(t.Context(), context.DeadlineExceeded)
+
+		if errors.Is(err, context.DeadlineExceeded) {
+			t.Fatalf("expected internal readiness deadline to be hidden, got %v", err)
+		}
+		if !strings.Contains(err.Error(), context.DeadlineExceeded.Error()) {
+			t.Fatalf("expected readiness error to retain deadline details, got %v", err)
+		}
+	})
+
+	t.Run("preserves a parent deadline", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(t.Context(), 0)
+		defer cancel()
+		<-ctx.Done()
+
+		err := proxyReadinessError(ctx, ctx.Err())
+
+		if !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatalf("expected parent deadline to remain identifiable, got %v", err)
+		}
+	})
+
+	t.Run("preserves probe errors", func(t *testing.T) {
+		probeErr := errors.New("docker exec failed")
+
+		err := proxyReadinessError(t.Context(), probeErr)
+
+		if !errors.Is(err, probeErr) {
+			t.Fatalf("expected probe error to remain identifiable, got %v", err)
+		}
+	})
+}

internal/infra/run.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -430,16 +431,18 @@ func runContainers(ctx context.Context, params RunParams) (err error) {
 		return err
 	}
 	defer func() {
-		if proxyErr := prox.Close(); proxyErr != nil {
-			err = proxyErr
-		}
+		err = errors.Join(err, prox.Close())
 	}()
 
 	// proxy logs interfere with debugging output
 	if !params.Debug {
 		go prox.TailLogs(ctx, cli)
 	}
 
+	if err = prox.WaitUntilReady(ctx); err != nil {
+		return err
+	}
+
 	var collector *Collector
 	if params.CollectorConfigPath != "" {
 		collector, err = NewCollector(ctx, cli, networks, &params, prox)

internal/infra/updater.go

@@ -539,33 +539,13 @@ func waitForPort(ctx context.Context, cli *client.Client, containerID string, po
 	const maxAttempts = 5
 	const sleepDuration = time.Second
 
-	// check /proc/net/tcp for the requested port; n.b., it is hex encoded and 4 characters wide
-	testCmd := fmt.Sprintf("test -f /proc/net/tcp && grep ' *\\d+: [A-F0-9]{8}:%04X ' /proc/net/tcp >/dev/null 2>&1", port)
-
 	for i := range maxAttempts {
-		execCreate, err := cli.ContainerExecCreate(ctx, containerID, container.ExecOptions{
-			AttachStdout: false,
-			AttachStderr: false,
-			User:         root,
-			Cmd:          []string{"/bin/sh", "-c", testCmd},
-		})
-		if err != nil {
-			return fmt.Errorf("failed to create exec for port check: %w", err)
-		}
-
-		execResp, err := cli.ContainerExecAttach(ctx, execCreate.ID, container.ExecAttachOptions{})
+		listening, err := isPortListening(ctx, cli, containerID, port)
 		if err != nil {
-			return fmt.Errorf("failed to attach to exec for port check: %w", err)
+			return err
 		}
 
-		// wait for completion and check the exit code
-		execResp.Close()
-		execInspect, err := cli.ContainerExecInspect(ctx, execCreate.ID)
-		if err != nil {
-			return fmt.Errorf("failed to inspect exec: %w", err)
-		}
-
-		if execInspect.ExitCode == 0 {
+		if listening {
 			// port is listening
 			log.Printf("  port %d is listening after %d attempts", port, i+1)
 
@@ -582,3 +562,76 @@ func waitForPort(ctx context.Context, cli *client.Client, containerID string, po
 
 	return fmt.Errorf("port %d is not listening after %d attempts", port, maxAttempts)
 }
+
+func waitForPortUntil(
+	ctx context.Context,
+	pollInterval time.Duration,
+	probe func(context.Context) (bool, error),
+) error {
+	for attempt := 1; ; attempt++ {
+		listening, err := probe(ctx)
+		if err != nil {
+			return err
+		}
+		if listening {
+			log.Printf("  proxy is listening after %d attempts", attempt)
+			return nil
+		}
+
+		timer := time.NewTimer(pollInterval)
+		select {
+		case <-ctx.Done():
+			timer.Stop()
+			return ctx.Err()
+		case <-timer.C:
+		}
+	}
+}
+
+func isPortListening(ctx context.Context, cli *client.Client, containerID string, port int) (bool, error) {
+	// The CLI already starts proxy images through sh. Use only shell built-ins
+	// here so readiness does not require additional executables in custom images.
+	testCmd := fmt.Sprintf(`
+for file in /proc/net/tcp /proc/net/tcp6; do
+  [ -r "$file" ] || continue
+  while read -r _ local_address _ state _; do
+    case "$local_address:$state" in
+      *:%04X:0A) exit 0 ;;
+    esac
+  done < "$file"
+done
+exit 1`, port)
+	return containerCommandSucceeded(ctx, cli, containerID, testCmd)
+}
+
+func containerCommandSucceeded(ctx context.Context, cli *client.Client, containerID, command string) (bool, error) {
+	return containerExecSucceeded(ctx, cli, containerID, []string{"sh", "-c", command})
+}
+
+func containerExecSucceeded(ctx context.Context, cli *client.Client, containerID string, command []string) (bool, error) {
+	execCreate, err := cli.ContainerExecCreate(ctx, containerID, container.ExecOptions{
+		AttachStdout: true,
+		AttachStderr: true,
+		User:         root,
+		Cmd:          command,
+	})
+	if err != nil {
+		return false, fmt.Errorf("failed to create container check: %w", err)
+	}
+
+	execResp, err := cli.ContainerExecAttach(ctx, execCreate.ID, container.ExecAttachOptions{})
+	if err != nil {
+		return false, fmt.Errorf("failed to attach to container check: %w", err)
+	}
+	if _, err = io.Copy(io.Discard, execResp.Reader); err != nil {
+		execResp.Close()
+		return false, fmt.Errorf("failed to wait for container check: %w", err)
+	}
+	execResp.Close()
+
+	execInspect, err := cli.ContainerExecInspect(ctx, execCreate.ID)
+	if err != nil {
+		return false, fmt.Errorf("failed to inspect container check: %w", err)
+	}
+	return execInspect.ExitCode == 0, nil
+}

testdata/scripts/proxy.txt

@@ -13,7 +13,9 @@ stderr 'proxy \| custom-ca-cert\.crt'
 stderr 'proxy \| I am a certificate'
 
 # Test that the CLI exits with non-zero if the proxy does too.
-! dependabot update go_modules dependabot/cli --proxy-cert crash --updater-image proxy-updater --proxy-image dummy-proxy --proxy-username user --proxy-password pass
+! dependabot update go_modules dependabot/cli --proxy-cert crash --updater-image proxy-updater --proxy-image dummy-proxy
+stderr 'proxy did not start listening on port 1080'
+stderr 'proxy container exited with non-zero exit code'
 
 exec docker rmi -f proxy-updater dummy-proxy
 
@@ -24,21 +26,63 @@ crash
 I am a certificate
 
 -- Dockerfile.proxy --
+FROM golang:1.26 AS builder
+
+WORKDIR /src
+COPY proxy.go .
+RUN CGO_ENABLED=0 go build -o /dependabot-proxy proxy.go
+
 FROM ubuntu:22.04
 
-COPY --chmod=755 update-ca-certificates /usr/bin/update-ca-certificates
-COPY --chmod=755 dependabot-proxy /dependabot-proxy
+RUN mkdir -p /custom/bin \
+    && cp /bin/sh /custom/bin/sh \
+    && ! command -v nc
 
--- dependabot-proxy --
-#!/usr/bin/env bash
+ENV PATH="/custom/bin:${PATH}"
 
-echo "Proxy is running"
-echo "$(</config.json)"
+COPY --chmod=755 update-ca-certificates /usr/bin/update-ca-certificates
+COPY --from=builder /dependabot-proxy /dependabot-proxy
+RUN rm /bin/sh \
+    && [ ! -e /bin/sh ]
+
+-- proxy.go --
+package main
+
+import (
+	"fmt"
+	"net"
+	"os"
+	"time"
+)
+
+func main() {
+	fmt.Println("Proxy is running")
+	config, _ := os.ReadFile("/config.json")
+	fmt.Println(string(config))
+	time.Sleep(time.Second)
+
+	listener, err := net.Listen("tcp", ":1080")
+	if err != nil {
+		panic(err)
+	}
+	defer listener.Close()
+
+	for {
+		conn, err := listener.Accept()
+		if err != nil {
+			panic(err)
+		}
+		conn.Close()
+	}
+}
 
 -- Dockerfile --
 FROM ubuntu:22.04
 
-RUN useradd dependabot
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends netcat-openbsd \
+    && rm -rf /var/lib/apt/lists/* \
+    && useradd dependabot
 
 COPY --chown=dependabot --chmod=755 update-ca-certificates /usr/bin/update-ca-certificates
 COPY --chown=dependabot --chmod=755 run bin/run
@@ -55,4 +99,8 @@ grep crash /usr/local/share/ca-certificates/custom-ca-cert.crt && exit 1 || true
 -- run --
 #!/usr/bin/env bash
 
+set -e
+
+proxy="${HTTP_PROXY#http://}"
+nc -w 1 "${proxy%:*}" "${proxy##*:}"
 echo "Updater is running"