diff --git a/.github/workflows/goreleaser.yml b/.github/workflows/goreleaser.yml index 6187a18..c90114e 100644 --- a/.github/workflows/goreleaser.yml +++ b/.github/workflows/goreleaser.yml @@ -7,28 +7,49 @@ on: - "*" permissions: - contents: write + contents: read jobs: + test: + runs-on: ubuntu-24.04 + steps: + - name: Checkout + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + + - name: Set up Go + uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 + with: + go-version-file: go.mod + + - name: Verify and test the tagged source + run: | + go mod download + go mod verify + go mod tidy -diff + go test -race ./... + goreleaser: - runs-on: ubuntu-latest + needs: [test] + runs-on: ubuntu-24.04 + permissions: + contents: write steps: - name: Checkout - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 with: - go-version: '>=1.22.2' + go-version-file: go.mod - name: Run GoReleaser uses: goreleaser/goreleaser-action@f06c13b6b1a9625abc9e6e439d9c05a8f2190e94 # v7 with: distribution: goreleaser - version: latest + version: v2.17.0 args: release --clean env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/lint-test-build.yml b/.github/workflows/lint-test-build.yml index 6fe6a5f..c6747dd 100644 --- a/.github/workflows/lint-test-build.yml +++ b/.github/workflows/lint-test-build.yml @@ -1,13 +1,21 @@ name: build-push on: + pull_request: push: + branches: + - main + tags: + - "*" workflow_dispatch: +permissions: + contents: read + jobs: lint-test: runs-on: ubuntu-24.04 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -17,10 +25,12 @@ jobs: - name: golangci-lint uses: golangci/golangci-lint-action@ba0d7d2ec06a0ea1cb5fa41b2e4a3ab91d21278a # v9 with: - version: latest + version: v2.12.2 - name: Install dependencies - run: go get . + run: | + go mod download + go mod verify - run: go build @@ -38,10 +48,21 @@ jobs: env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} - run: + publish-ghcr: + if: github.event_name != 'pull_request' needs: [lint-test] - uses: libops/.github/.github/workflows/build-push.yaml@3ef8469c16c40842cc1b0b9195df2850ad8e616f # main + uses: libops/.github/.github/workflows/build-push.yaml@3bdb895c3fd05490b0c6a1341e630ac554842468 # main permissions: contents: read packages: write + + publish-gar: + if: github.event_name != 'pull_request' + needs: [lint-test] + uses: libops/.github/.github/workflows/build-push.yaml@3bdb895c3fd05490b0c6a1341e630ac554842468 # main + with: + docker-registry: us-docker.pkg.dev/libops-images/public + permissions: + contents: read + id-token: write secrets: inherit diff --git a/.goreleaser.yml b/.goreleaser.yml index 7cf9229..af81a19 100644 --- a/.goreleaser.yml +++ b/.goreleaser.yml @@ -1,6 +1,6 @@ before: hooks: - - go mod tidy + - go mod tidy -diff builds: - binary: ppb env: diff --git a/README.md b/README.md index 04d241a..0a48777 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,8 @@ PPB runs as a reverse proxy with intelligent rate limiting: 1. **Request handling**: Each authorized request attempts to power on the target machine if needed 2. **Rate limiting**: Power-on attempts are rate limited with a configurable cooldown period (default 30s) 3. **API protection**: Prevents hammering Google Cloud APIs during high traffic periods -4. **IP Security**: Only requests from allowed IP ranges can trigger machine power-on attempts +4. **Client filter**: Only explicitly allowed original-client CIDRs can trigger power-on after a configured, trusted proxy suffix is removed +5. **Bounded readiness retry**: Retries connection establishment while a VM or Direct VPC path becomes reachable, without adding application-level request retries This approach balances responsiveness with API efficiency - machines get powered on when traffic arrives, but multiple requests don't spam the Google Cloud APIs. @@ -53,13 +54,11 @@ type: google_compute_engine port: 80 scheme: http allowedIps: - - 127.0.0.1/32 - - 10.0.0.0/8 - - 172.16.0.0/12 - - 192.168.0.0/16 -ipForwardedHeader: X-Forwarded-For # header to check for real client IP -ipDepth: 0 # depth in X-Forwarded-For chain + - 127.0.0.1/32 # replace with an original-client CIDR for deployment +ipForwardedHeader: "" # direct peer address; configure only for a proven proxy chain +ipDepth: 0 # trusted proxy hops after the client in X-Forwarded-For powerOnCooldown: 30 # seconds between power-on attempts (default: 30) +powerOnTimeout: 360 # total queued/API/instance-ready budget (default: 360) # metadata about the machine that will be turned on machineMetadata: project_id: foo @@ -68,7 +67,9 @@ machineMetadata: usePrivateIp: false # true for VPC-native setups # Optional proxy timeout configuration (defaults shown) proxyTimeouts: - dialTimeout: 120 # Connection dial timeout in seconds + dialTimeout: 120 # Total connection retry window in seconds + dialAttemptTimeout: 5 # Timeout for one connection attempt in seconds + dialRetryInterval: 1 # Delay between connection attempts in seconds keepAlive: 120 # TCP keep-alive timeout in seconds idleConnTimeout: 90 # Idle connection timeout in seconds tlsHandshakeTimeout: 10 # TLS handshake timeout in seconds @@ -86,9 +87,12 @@ proxyTimeouts: | `scheme` | string | ✅ | - | Protocol scheme (`http` or `https`) | | `allowedIps` | []string | ✅ | - | CIDR ranges of IPs allowed to access the proxy | | `ipForwardedHeader` | string | ❌ | `""` | Header to check for real client IP (e.g., `X-Forwarded-For`) | -| `ipDepth` | int | ❌ | `0` | Depth in forwarded header chain (0 = rightmost IP) | +| `ipDepth` | int | ❌ | `0` | Trusted proxy hops after the client (0 selects rightmost IP) | | `powerOnCooldown` | int | ❌ | `30` | Seconds between power-on attempts (rate limiting) | -| `proxyTimeouts.dialTimeout` | int | ❌ | `120` | Connection dial timeout in seconds | +| `powerOnTimeout` | int | ❌ | `360` | Total queued, API, and instance-ready budget in seconds | +| `proxyTimeouts.dialTimeout` | int | ❌ | `120` | Total connection retry window in seconds | +| `proxyTimeouts.dialAttemptTimeout` | int | ❌ | `5` | Timeout for one connection attempt in seconds | +| `proxyTimeouts.dialRetryInterval` | int | ❌ | `1` | Delay between connection attempts in seconds | | `proxyTimeouts.keepAlive` | int | ❌ | `120` | TCP keep-alive timeout in seconds | | `proxyTimeouts.idleConnTimeout` | int | ❌ | `90` | Idle connection timeout in seconds | | `proxyTimeouts.tlsHandshakeTimeout` | int | ❌ | `10` | TLS handshake timeout in seconds | @@ -101,6 +105,24 @@ proxyTimeouts: Deploy this service on **Google Cloud Run** as the public endpoint for your application. Configure the `machineMetadata` to point to your GCE VM running the actual application stack. Only requests from allowed IPs will power on the VM and be proxied through. Set to `0.0.0.0/0` to allow any request to power on the machine. +Treat `ipForwardedHeader` and `ipDepth` as a proxy-chain contract, not an IP +parser convenience. PPB selects from the right edge so an attacker-supplied +prefix cannot replace the client hop. A missing header, malformed address, or +chain shorter than `ipDepth + 1` is denied; PPB never falls back to the proxy's +`RemoteAddr` when a trusted header is configured. Use depth `0` only for a +direct Cloud Run path whose hosted test proves Google appends the original +client last. Duplicate field lines are folded before suffix selection. An added +external Application Load Balancer normally adds a hop and requires a different +depth. Do not enable a forwarded header for an ingress path that has not proven +its appended suffix. Prefer Cloud Run IAM or a controlled load balancer/Cloud +Armor policy when caller authentication, rather than a wake-up filter, is +required. Before proxying, PPB replaces forwarding identity headers with the +single validated client address. + +`proxyTimeouts.dialTimeout` bounds the complete TCP connection-establishment retry window. Each attempt is bounded by `dialAttemptTimeout`, with `dialRetryInterval` between failures. The readiness loop runs in the transport dialer before an HTTP connection exists; PPB does not add application-level request or status retries. Go's standard transport can retry requests it defines as replayable when a pooled connection is found stale. If the connection window expires, PPB returns `503 Service Unavailable` with `Retry-After: 5` so the client can make a deliberate retry. Request cancellation stops GCE polling, queued power-on work, and connection retry. HTTPS handshake readiness is outside the TCP retry loop and fails with the same retryable 503 response. + +For Direct VPC egress, use a supported `/26` or larger subnet with sufficient free addresses, grant the Cloud Run service agent subnet use, and authorize the whole Cloud Run subnet CIDR at the VM firewall. Cloud Run addresses are ephemeral; never build the firewall around one revision address. PPB tolerates initial connection refusal and timeout within the configured retry window, but clients must still tolerate occasional connection resets after a connection has been established. + ### Environment Variables PPB also supports these environment variables for runtime configuration: @@ -138,10 +160,12 @@ resource "google_project_iam_custom_role" "compute-start" { ### Assign Role to Service Account ```hcl -resource "google_project_iam_member" "ppb_service_account" { - project = var.project - role = google_project_iam_custom_role.compute-start.name - member = "serviceAccount:${var.service_account_email}" +resource "google_compute_instance_iam_member" "ppb_service_account" { + project = var.project + zone = var.zone + instance_name = var.instance_name + role = google_project_iam_custom_role.compute-start.name + member = "serviceAccount:${var.service_account_email}" } ``` @@ -154,7 +178,9 @@ gcloud iam roles create startVM \ --description="Minimal permissions for PPB to control compute instances" \ --permissions="compute.instances.start,compute.instances.resume,compute.instances.get" -gcloud projects add-iam-policy-binding PROJECT_ID \ +gcloud compute instances add-iam-policy-binding INSTANCE_NAME \ + --zone=ZONE \ + --project=PROJECT_ID \ --member="serviceAccount:SERVICE_ACCOUNT_EMAIL" \ --role="projects/PROJECT_ID/roles/startVM" ``` @@ -166,30 +192,51 @@ gcloud projects add-iam-policy-binding PROJECT_ID \ ```bash docker build -t us-docker.pkg.dev/your/gar/ppb:development . docker run \ - -p 8080:8080 \ + --network=host \ -v $GOOGLE_APPLICATION_CREDENTIALS:/app/svc.json:ro \ -v ./ppb.yaml:/app/ppb.yaml:ro \ --env PPB_CONFIG_PATH=/app/ppb.yaml \ --env GOOGLE_APPLICATION_CREDENTIALS=/app/svc.json \ - --env PRIVATE_IP=false \ --env LOG_LEVEL=DEBUG \ us-docker.pkg.dev/your/gar/ppb:development ``` PPB will act as a local ingress proxy to your GCE VM. Visit http://localhost:8080 and PPB will power on your target VM (if needed) then proxy your requests through to the application running on the VM. +The host-network example is for Linux and makes the checked-in loopback +allowlist match the direct browser peer. On Docker Desktop, set `allowedIps` to +the exact host-gateway address reported inside the container instead of adding +a broad private range. + ### Google Cloud Run Deployment Deploy PPB to Cloud Run and configure it as your application's public endpoint: ```bash gcloud run deploy ppb \ - --image=us-docker.pkg.dev/your/gar/ppb:latest \ + --image=us-docker.pkg.dev/your/gar/ppb:VERSION@sha256:DIGEST \ + --set-secrets="/app/ppb.yaml=ppb-config:latest" \ --set-env-vars="PPB_CONFIG_PATH=/app/ppb.yaml" \ --service-account=your-ppb-service-account@project.iam.gserviceaccount.com \ --allow-unauthenticated \ --port=8080 \ + --timeout=600 \ + --max-instances=1 \ + --network=NETWORK \ + --subnet=SUBNET \ + --vpc-egress=private-ranges-only \ --region=us-central1 ``` +Keep the Cloud Run request timeout greater than `powerOnTimeout + +dialTimeout`, with enough margin for the proxied application response. The +defaults use a 360-second power budget and 120-second connection window, so the +example configures 600 seconds. + +Keep one active PPB instance so concurrent requests share its in-process power +coordinator. Revision overlap can briefly run an old and new instance at once; +PPB reconciles a conflicting start or resume by joining the observed VM +transition. Grant the service account the custom start role on the target +instance, not across the project. + Your users will access `https://ppb-xxx-uc.a.run.app` which automatically powers on your GCE VM and proxies traffic through. diff --git a/main.go b/main.go index a49f5c5..342677d 100644 --- a/main.go +++ b/main.go @@ -67,7 +67,9 @@ func startPingRoutine(ctx context.Context, wg *sync.WaitGroup, c *config.Config, slog.Debug("Ping failed", "url", pingURL, "error", err) continue } - resp.Body.Close() + if err := resp.Body.Close(); err != nil { + slog.Debug("Unable to close ping response body", "error", err) + } slog.Debug("Ping successful", "url", pingURL, "status", resp.StatusCode) } @@ -83,11 +85,6 @@ func main() { slog.Debug("Loaded config", "config", c) - // Set default cooldown interval if not specified - if c.PowerOnCooldown <= 0 { - c.PowerOnCooldown = 30 - } - ctx, cancel := context.WithCancel(context.Background()) defer cancel() sigChan := make(chan os.Signal, 1) @@ -96,33 +93,12 @@ func main() { wg.Add(1) go startPingRoutine(ctx, &wg, c, 30*time.Second) - http.HandleFunc("/healthcheck", func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusOK) - _, _ = fmt.Fprintln(w, "OK") - }) - p := proxy.New(c) - http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { - if !c.IpIsAllowed(r) { - http.Error(w, "Forbidden", http.StatusForbidden) - return - } - - // Attempt to power on machine with cooldown protection - reqCtx := context.Background() - err := c.Machine.PowerOnWithCooldown(reqCtx, c.PowerOnCooldown) - if err != nil { - slog.Error("Power-on attempt failed", "err", err) - http.Error(w, "Backend not available", http.StatusServiceUnavailable) - return - } - - p.SetRequestHeaders(r) - p.SetHost() - p.ServeHTTP(w, r) - }) - - server := &http.Server{Addr: ":8080"} + server := &http.Server{ + Addr: ":8080", + Handler: newHandler(c, p), + ReadHeaderTimeout: 10 * time.Second, + } go func() { slog.Info("Server listening on :8080") if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { @@ -143,3 +119,44 @@ func main() { wg.Wait() slog.Info("Shutdown complete") } + +func newHandler(c *config.Config, backend http.Handler) http.Handler { + mux := http.NewServeMux() + mux.HandleFunc("/healthcheck", func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = fmt.Fprintln(w, "OK") + }) + + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + clientIP, err := c.AllowedClientIP(r) + if err != nil { + http.Error(w, "Forbidden", http.StatusForbidden) + return + } + + // Do not pass attacker-controlled forwarding identities to the backend. + // Preserve one canonical, already validated original-client address. + r.Header.Del("Forwarded") + r.Header.Del("X-Forwarded-For") + r.Header.Del("X-Real-IP") + if c.IpForwardedHeader != "" { + r.Header.Del(c.IpForwardedHeader) + } + r.Header.Set("X-Forwarded-For", clientIP.String()) + + // Attempt to power on the machine within the request lifetime. Waiting + // requests can then be cancelled cleanly during disconnect or shutdown. + powerCtx, powerCancel := context.WithTimeout(r.Context(), time.Duration(c.PowerOnTimeout)*time.Second) + err = c.Machine.PowerOnWithCooldown(powerCtx, c.PowerOnCooldown) + powerCancel() + if err != nil { + slog.Error("Power-on attempt failed", "err", err) + w.Header().Set("Retry-After", "5") + http.Error(w, "Backend not available", http.StatusServiceUnavailable) + return + } + + backend.ServeHTTP(w, r) + }) + return mux +} diff --git a/main_test.go b/main_test.go index d9c9f69..48dffd5 100644 --- a/main_test.go +++ b/main_test.go @@ -3,7 +3,9 @@ package main import ( "context" "log/slog" + "net" "net/http" + "net/http/httptest" "os" "strings" "sync" @@ -14,6 +16,86 @@ import ( "github.com/libops/ppb/pkg/machine" ) +func TestHandlerPowerFailureReturnsRetryableUnavailable(t *testing.T) { + t.Parallel() + + _, allowed, err := net.ParseCIDR("127.0.0.1/32") + if err != nil { + t.Fatalf("net.ParseCIDR() error = %v", err) + } + backendCalled := false + handler := newHandler(&config.Config{ + AllowedIps: []config.IPNet{{IPNet: allowed}}, + PowerOnCooldown: 30, + PowerOnTimeout: 1, + Machine: &machine.GoogleComputeEngine{}, + }, http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + backendCalled = true + })) + request := httptest.NewRequest(http.MethodGet, "http://example.test/", nil) + request.RemoteAddr = "127.0.0.1:12345" + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusServiceUnavailable) + } + if got := recorder.Header().Get("Retry-After"); got != "5" { + t.Fatalf("Retry-After = %q, want 5", got) + } + if backendCalled { + t.Fatal("backend handler ran after power-on failure") + } +} + +func TestHandlerCanonicalizesValidatedClientIdentity(t *testing.T) { + t.Parallel() + + _, allowed, err := net.ParseCIDR("203.0.113.9/32") + if err != nil { + t.Fatalf("net.ParseCIDR() error = %v", err) + } + machine := machine.NewGceMachine() + machine.SetHostForTesting("10.42.0.8") + machine.LastPowerOnAttempt = time.Now() + + backendCalled := false + handler := newHandler(&config.Config{ + AllowedIps: []config.IPNet{{IPNet: allowed}}, + IpForwardedHeader: "X-Forwarded-For", + PowerOnCooldown: 30, + PowerOnTimeout: 1, + Machine: machine, + }, http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) { + backendCalled = true + if got := request.Header.Values("X-Forwarded-For"); len(got) != 1 || got[0] != "203.0.113.9" { + t.Errorf("X-Forwarded-For = %#v, want one validated client address", got) + } + if got := request.Header.Get("X-Real-IP"); got != "" { + t.Errorf("X-Real-IP = %q, want stripped", got) + } + if got := request.Header.Get("Forwarded"); got != "" { + t.Errorf("Forwarded = %q, want stripped", got) + } + })) + request := httptest.NewRequest(http.MethodGet, "http://example.test/", nil) + request.Header.Add("X-Forwarded-For", "10.0.0.8") + request.Header.Add("X-Forwarded-For", "203.0.113.9") + request.Header.Set("X-Real-IP", "10.0.0.8") + request.Header.Set("Forwarded", "for=10.0.0.8") + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusOK) + } + if !backendCalled { + t.Fatal("backend handler was not called") + } +} + func TestStartPingRoutine_Integration(t *testing.T) { // Track ping requests var pingCount int diff --git a/pkg/config/config.go b/pkg/config/config.go index 860c924..b385c3a 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -5,6 +5,7 @@ import ( "log/slog" "net" "os" + "strings" "github.com/libops/ppb/pkg/machine" yaml "gopkg.in/yaml.v3" @@ -18,6 +19,7 @@ type Config struct { IpForwardedHeader string `yaml:"ipForwardedHeader"` IpDepth int `yaml:"ipDepth"` PowerOnCooldown int `yaml:"powerOnCooldown"` // seconds + PowerOnTimeout int `yaml:"powerOnTimeout"` // seconds, default: 360 ProxyTimeouts ProxyTimeouts `yaml:"proxyTimeouts"` MachineMetadata map[string]any `yaml:"machineMetadata"` ProxyTarget *ProxyTarget `yaml:"proxyTarget"` @@ -36,7 +38,9 @@ type ProxyTarget struct { } type ProxyTimeouts struct { - DialTimeout int `yaml:"dialTimeout"` // seconds, default: 120 + DialTimeout int `yaml:"dialTimeout"` // total connection retry window in seconds, default: 120 + DialAttemptTimeout int `yaml:"dialAttemptTimeout"` // timeout for one connection attempt in seconds, default: 5 + DialRetryInterval int `yaml:"dialRetryInterval"` // delay between connection attempts in seconds, default: 1 KeepAlive int `yaml:"keepAlive"` // seconds, default: 120 IdleConnTimeout int `yaml:"idleConnTimeout"` // seconds, default: 90 TLSHandshakeTimeout int `yaml:"tlsHandshakeTimeout"` // seconds, default: 10 @@ -79,7 +83,10 @@ func LoadConfig() (*Config, error) { filename = "ppb.yaml" } slog.Debug("Loading config", "filename", filename) - data, err = os.ReadFile(filename) + // PPB_CONFIG_PATH is an intentional local operator interface, not a + // request-derived path. The unprivileged container may read any mounted + // configuration path selected by its deployer. + data, err = os.ReadFile(filename) // #nosec G304,G703 -- trusted deployment configuration if err != nil { return nil, err } @@ -92,6 +99,13 @@ func LoadConfig() (*Config, error) { if err != nil { return nil, err } + config.IpForwardedHeader = strings.TrimSpace(config.IpForwardedHeader) + if config.IpDepth < 0 { + return nil, fmt.Errorf("ipDepth must not be negative") + } + if config.IpDepth > 0 && config.IpForwardedHeader == "" { + return nil, fmt.Errorf("ipForwardedHeader is required when ipDepth is greater than zero") + } switch config.Type { case "google_compute_engine": @@ -107,16 +121,32 @@ func LoadConfig() (*Config, error) { } // Set default proxy timeouts if not specified + config.setPowerDefaults() config.setProxyTimeoutDefaults() return &config, nil } +func (c *Config) setPowerDefaults() { + if c.PowerOnCooldown <= 0 { + c.PowerOnCooldown = 30 + } + if c.PowerOnTimeout <= 0 { + c.PowerOnTimeout = 360 + } +} + // setProxyTimeoutDefaults sets default values for proxy timeouts if not configured func (c *Config) setProxyTimeoutDefaults() { if c.ProxyTimeouts.DialTimeout <= 0 { c.ProxyTimeouts.DialTimeout = 120 } + if c.ProxyTimeouts.DialAttemptTimeout <= 0 { + c.ProxyTimeouts.DialAttemptTimeout = 5 + } + if c.ProxyTimeouts.DialRetryInterval <= 0 { + c.ProxyTimeouts.DialRetryInterval = 1 + } if c.ProxyTimeouts.KeepAlive <= 0 { c.ProxyTimeouts.KeepAlive = 120 } diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index f5c8dc0..de72f63 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -4,6 +4,8 @@ import ( "net" "net/http" "os" + "path/filepath" + "runtime" "testing" yaml "gopkg.in/yaml.v3" @@ -136,6 +138,10 @@ func TestConfig_IpIsAllowed(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + testConfig := *config + if tt.forwardedFor == "" { + testConfig.IpForwardedHeader = "" + } req := &http.Request{ RemoteAddr: tt.remoteAddr, Header: make(http.Header), @@ -145,7 +151,7 @@ func TestConfig_IpIsAllowed(t *testing.T) { req.Header.Set("X-Forwarded-For", tt.forwardedFor) } - result := config.IpIsAllowed(req) + result := testConfig.IpIsAllowed(req) if result != tt.expectedResult { t.Errorf("Config.IpIsAllowed() = %v, want %v for IP %s (forwarded: %s)", result, tt.expectedResult, tt.remoteAddr, tt.forwardedFor) @@ -161,6 +167,8 @@ func TestConfig_getClientIP(t *testing.T) { forwardedFor string ipDepth int expectedIP string + useHeader bool + wantErr bool }{ { name: "no forwarded header", @@ -172,6 +180,7 @@ func TestConfig_getClientIP(t *testing.T) { remoteAddr: "10.0.0.1:8080", forwardedFor: "203.0.113.1", expectedIP: "203.0.113.1", + useHeader: true, }, { name: "multiple forwarded IPs, depth 0", @@ -179,6 +188,7 @@ func TestConfig_getClientIP(t *testing.T) { forwardedFor: "203.0.113.1, 198.51.100.1, 192.168.1.1", ipDepth: 0, expectedIP: "192.168.1.1", + useHeader: true, }, { name: "multiple forwarded IPs, depth 1", @@ -186,6 +196,30 @@ func TestConfig_getClientIP(t *testing.T) { forwardedFor: "203.0.113.1, 198.51.100.1, 192.168.1.1", ipDepth: 1, expectedIP: "198.51.100.1", + useHeader: true, + }, + { + name: "attacker prefix cannot replace client behind one trusted proxy", + remoteAddr: "10.0.0.1:8080", + forwardedFor: "10.0.0.8, 203.0.113.9, 192.0.2.10", + ipDepth: 1, + expectedIP: "203.0.113.9", + useHeader: true, + }, + { + name: "missing trusted header fails closed", + remoteAddr: "10.0.0.1:8080", + ipDepth: 0, + useHeader: true, + wantErr: true, + }, + { + name: "insufficient trusted hops fail closed", + remoteAddr: "10.0.0.1:8080", + forwardedFor: "203.0.113.9", + ipDepth: 1, + useHeader: true, + wantErr: true, }, { name: "IPv6 remote addr", @@ -196,9 +230,9 @@ func TestConfig_getClientIP(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - config := &Config{ - IpForwardedHeader: "X-Forwarded-For", - IpDepth: tt.ipDepth, + config := &Config{IpDepth: tt.ipDepth} + if tt.useHeader { + config.IpForwardedHeader = "X-Forwarded-For" } req := &http.Request{ @@ -210,7 +244,13 @@ func TestConfig_getClientIP(t *testing.T) { req.Header.Set("X-Forwarded-For", tt.forwardedFor) } - result := config.getClientIP(req) + result, err := config.getClientIP(req) + if (err != nil) != tt.wantErr { + t.Fatalf("Config.getClientIP() error = %v, wantErr %v", err, tt.wantErr) + } + if tt.wantErr { + return + } if result.String() != tt.expectedIP { t.Errorf("Config.getClientIP() = %v, want %v", result.String(), tt.expectedIP) } @@ -218,20 +258,38 @@ func TestConfig_getClientIP(t *testing.T) { } } -func TestConfig_PowerOnCooldownDefaultValue(t *testing.T) { - config := &Config{} +func TestConfigGetClientIPUsesTheRightmostDuplicateHeaderValue(t *testing.T) { + t.Parallel() + + config := &Config{ + IpForwardedHeader: "X-Forwarded-For", + IpDepth: 0, + } + request := &http.Request{Header: make(http.Header)} + request.Header.Add("X-Forwarded-For", "10.0.0.8") + request.Header.Add("X-Forwarded-For", "203.0.113.9") + + clientIP, err := config.getClientIP(request) + if err != nil { + t.Fatalf("getClientIP() error = %v", err) + } + if got := clientIP.String(); got != "203.0.113.9" { + t.Fatalf("getClientIP() = %s, want proxy-appended duplicate value", got) + } +} - // Test that default is applied in main.go logic - // (We can't easily test the main.go logic here, but we can test the field exists) - if config.PowerOnCooldown != 0 { - // Should start at 0, then main.go sets to 30 if <= 0 - t.Errorf("PowerOnCooldown should start at 0, got %v", config.PowerOnCooldown) +func TestConfig_SetPowerDefaults(t *testing.T) { + config := &Config{} + config.setPowerDefaults() + if config.PowerOnCooldown != 30 || config.PowerOnTimeout != 360 { + t.Fatalf("power defaults = cooldown %d timeout %d, want 30 and 360", config.PowerOnCooldown, config.PowerOnTimeout) } - // Test setting a custom value config.PowerOnCooldown = 60 - if config.PowerOnCooldown != 60 { - t.Errorf("PowerOnCooldown should be 60, got %v", config.PowerOnCooldown) + config.PowerOnTimeout = 480 + config.setPowerDefaults() + if config.PowerOnCooldown != 60 || config.PowerOnTimeout != 480 { + t.Fatalf("configured power values = cooldown %d timeout %d, want 60 and 480", config.PowerOnCooldown, config.PowerOnTimeout) } } @@ -248,6 +306,8 @@ func TestConfig_setProxyTimeoutDefaults(t *testing.T) { }, expected: ProxyTimeouts{ DialTimeout: 120, + DialAttemptTimeout: 5, + DialRetryInterval: 1, KeepAlive: 120, IdleConnTimeout: 90, TLSHandshakeTimeout: 10, @@ -265,6 +325,8 @@ func TestConfig_setProxyTimeoutDefaults(t *testing.T) { }, expected: ProxyTimeouts{ DialTimeout: 60, // custom + DialAttemptTimeout: 5, // default + DialRetryInterval: 1, // default KeepAlive: 120, // default IdleConnTimeout: 90, // default TLSHandshakeTimeout: 10, // default @@ -277,6 +339,8 @@ func TestConfig_setProxyTimeoutDefaults(t *testing.T) { config: Config{ ProxyTimeouts: ProxyTimeouts{ DialTimeout: 30, + DialAttemptTimeout: 4, + DialRetryInterval: 2, KeepAlive: 40, IdleConnTimeout: 50, TLSHandshakeTimeout: 5, @@ -286,6 +350,8 @@ func TestConfig_setProxyTimeoutDefaults(t *testing.T) { }, expected: ProxyTimeouts{ DialTimeout: 30, + DialAttemptTimeout: 4, + DialRetryInterval: 2, KeepAlive: 40, IdleConnTimeout: 50, TLSHandshakeTimeout: 5, @@ -302,6 +368,12 @@ func TestConfig_setProxyTimeoutDefaults(t *testing.T) { if tt.config.ProxyTimeouts.DialTimeout != tt.expected.DialTimeout { t.Errorf("DialTimeout = %v, want %v", tt.config.ProxyTimeouts.DialTimeout, tt.expected.DialTimeout) } + if tt.config.ProxyTimeouts.DialAttemptTimeout != tt.expected.DialAttemptTimeout { + t.Errorf("DialAttemptTimeout = %v, want %v", tt.config.ProxyTimeouts.DialAttemptTimeout, tt.expected.DialAttemptTimeout) + } + if tt.config.ProxyTimeouts.DialRetryInterval != tt.expected.DialRetryInterval { + t.Errorf("DialRetryInterval = %v, want %v", tt.config.ProxyTimeouts.DialRetryInterval, tt.expected.DialRetryInterval) + } if tt.config.ProxyTimeouts.KeepAlive != tt.expected.KeepAlive { t.Errorf("KeepAlive = %v, want %v", tt.config.ProxyTimeouts.KeepAlive, tt.expected.KeepAlive) } @@ -378,6 +450,29 @@ machineMetadata: } } +func TestLoadConfigCheckedInExample(t *testing.T) { + _, sourceFile, _, ok := runtime.Caller(0) + if !ok { + t.Fatal("runtime.Caller() could not locate the test source") + } + t.Setenv("PPB_YAML", "") + t.Setenv("PPB_CONFIG_PATH", filepath.Join(filepath.Dir(sourceFile), "..", "..", "ppb.example.yaml")) + + loaded, err := LoadConfig() + if err != nil { + t.Fatalf("LoadConfig() could not load ppb.example.yaml: %v", err) + } + if len(loaded.AllowedIps) != 2 { + t.Fatalf("ppb.example.yaml allowedIps = %+v, want two loopback networks", loaded.AllowedIps) + } + if loaded.IpForwardedHeader != "" { + t.Fatalf("ppb.example.yaml ipForwardedHeader = %q, want direct-client mode", loaded.IpForwardedHeader) + } + if loaded.PowerOnCooldown != 30 || loaded.PowerOnTimeout != 360 { + t.Fatalf("ppb.example.yaml power values = %d/%d, want 30/360", loaded.PowerOnCooldown, loaded.PowerOnTimeout) + } +} + func TestLoadConfig_PriorityOrder(t *testing.T) { // Save original env vars originalYAML := os.Getenv("PPB_YAML") diff --git a/pkg/config/ipAuth.go b/pkg/config/ipAuth.go index e3fd542..dfbaf6b 100644 --- a/pkg/config/ipAuth.go +++ b/pkg/config/ipAuth.go @@ -1,6 +1,7 @@ package config import ( + "fmt" "log/slog" "net" "net/http" @@ -8,48 +9,59 @@ import ( ) func (c *Config) IpIsAllowed(r *http.Request) bool { - ip := c.getClientIP(r) + _, err := c.AllowedClientIP(r) + return err == nil +} + +// AllowedClientIP returns the parsed original-client address only when it is +// inside an explicitly configured allowlist block. +func (c *Config) AllowedClientIP(r *http.Request) (net.IP, error) { + ip, err := c.getClientIP(r) + if err != nil { + slog.Warn("Unable to determine client IP; denying request", "error", err) + return nil, err + } for _, block := range c.AllowedIps { if block.Contains(ip) { - return true + return ip, nil } } - - return false + slog.Debug("Client IP is not allowed", "ip", ip) + return nil, fmt.Errorf("client IP %s is not allowed", ip) } -func (c *Config) getClientIP(r *http.Request) net.IP { - ip := r.Header.Get(c.IpForwardedHeader) +func (c *Config) getClientIP(r *http.Request) (net.IP, error) { + value := "" if c.IpForwardedHeader != "" { - components := strings.Split(ip, ",") - depth := c.IpDepth - ip = "" - for i := len(components) - 1; i >= 0; i-- { - _ip := strings.TrimSpace(components[i]) - if depth == 0 { - ip = _ip - break - } - depth-- + // Flatten every field line before selecting from the trusted suffix. + // Header.Get returns only the first value, which could otherwise let an + // attacker-controlled duplicate hide a proxy-appended client hop. + headerValue := strings.TrimSpace(strings.Join(r.Header.Values(c.IpForwardedHeader), ",")) + if headerValue == "" { + return nil, fmt.Errorf("trusted client IP header %q is missing", c.IpForwardedHeader) } - } - - if ip == "" { - if c.IpForwardedHeader != "" { - slog.Debug("No IPs in header. Using r.RemoteAddr", "ipDepth", c.IpDepth, c.IpForwardedHeader, r.Header.Get(c.IpForwardedHeader)) + components := strings.Split(headerValue, ",") + index := len(components) - 1 - c.IpDepth + if c.IpDepth < 0 || index < 0 { + return nil, fmt.Errorf("trusted client IP header %q has %d hops, fewer than configured depth %d", c.IpForwardedHeader, len(components), c.IpDepth) } - ip = r.RemoteAddr - } - - if strings.Contains(ip, ":") { - host, _, err := net.SplitHostPort(ip) + value = strings.TrimSpace(components[index]) + if strings.ContainsAny(value, " \t") { + return nil, fmt.Errorf("trusted client IP header %q contains an invalid hop", c.IpForwardedHeader) + } + } else { + host, _, err := net.SplitHostPort(r.RemoteAddr) if err != nil { - slog.Error("Error splitting host and port from client IP", "ip", ip, "err", err) + value = strings.TrimSpace(r.RemoteAddr) + } else { + value = host } - ip = host } + ip := net.ParseIP(value) + if ip == nil { + return nil, fmt.Errorf("client IP value %q is not an IP address", value) + } slog.Debug("Got client IP", "ip", ip) - - return net.ParseIP(ip) + return ip, nil } diff --git a/pkg/machine/gce.go b/pkg/machine/gce.go index 7d50798..e09ac6f 100644 --- a/pkg/machine/gce.go +++ b/pkg/machine/gce.go @@ -21,9 +21,21 @@ type GoogleComputeEngine struct { host string hostMutex sync.RWMutex LastPowerOnAttempt time.Time - powerOnMutex sync.Mutex + getInstanceHook func(context.Context) (*compute.Instance, error) + powerOnHook func(context.Context, string) error + pollInterval time.Duration + joinTimeout time.Duration + now func() time.Time } +type instanceStatusAction int + +const ( + instanceReady instanceStatusAction = iota + instanceStart + instanceWait +) + func NewGceMachine() *GoogleComputeEngine { return &GoogleComputeEngine{ Lock: semaphore.NewWeighted(1), @@ -50,25 +62,33 @@ func (m *GoogleComputeEngine) PowerOn(ctx context.Context) error { return fmt.Errorf("could not fetch instance metadata: %v", err) } - if vm.Status == "RUNNING" { - slog.Debug("VM is running") - return m.setIp(vm) - } - slog.Debug("Instance status", "status", vm.Status, "instance", m.Name) - if vm.Status != "TERMINATED" && vm.Status != "SUSPENDED" { - return nil - } - - err = m.powerOn(ctx, vm.Status) + action, err := classifyInstanceStatus(vm.Status) if err != nil { - return fmt.Errorf("could not power on: %v", err) + return err + } + startRequested := false + switch action { + case instanceReady: + return m.setIp(vm) + case instanceStart: + if err := m.powerOn(ctx, vm.Status); err != nil { + return m.joinAfterPowerOnError(ctx, fmt.Errorf("could not power on: %w", err)) + } + startRequested = true + case instanceWait: + // Another request or operator has already initiated a state change. + // Continue through the bounded state loop instead of returning without + // a usable target IP. } - return m.waitForInstanceRunning(ctx) + return m.waitForInstanceRunning(ctx, startRequested) } func (m *GoogleComputeEngine) getInstanceMetadata(ctx context.Context) (*compute.Instance, error) { + if m.getInstanceHook != nil { + return m.getInstanceHook(ctx) + } computeService, err := compute.NewService(ctx, option.WithScopes(compute.CloudPlatformScope)) if err != nil { return nil, fmt.Errorf("failed to create compute service: %v", err) @@ -84,6 +104,9 @@ func (m *GoogleComputeEngine) getInstanceMetadata(ctx context.Context) (*compute } func (m *GoogleComputeEngine) powerOn(ctx context.Context, status string) error { + if m.powerOnHook != nil { + return m.powerOnHook(ctx, status) + } computeService, err := compute.NewService(ctx, option.WithScopes(compute.CloudPlatformScope)) if err != nil { return fmt.Errorf("failed to create compute service: %v", err) @@ -109,18 +132,16 @@ func (m *GoogleComputeEngine) powerOn(ctx context.Context, status string) error } // Polls instance metadata until the VM is in the RUNNING state -func (m *GoogleComputeEngine) waitForInstanceRunning(ctx context.Context) error { - ticker := time.NewTicker(2 * time.Second) +func (m *GoogleComputeEngine) waitForInstanceRunning(ctx context.Context, startRequested bool) error { + ticker := time.NewTicker(m.effectivePollInterval()) defer ticker.Stop() - timeout := time.After(5 * time.Minute) + seenTransition := false for { select { case <-ctx.Done(): return ctx.Err() - case <-timeout: - return fmt.Errorf("timeout waiting for instance to start") case <-ticker.C: vm, err := m.getInstanceMetadata(ctx) if err != nil { @@ -130,10 +151,101 @@ func (m *GoogleComputeEngine) waitForInstanceRunning(ctx context.Context) error slog.Debug("Polling instance status", "status", vm.Status, "instance", m.Name) - if vm.Status == "RUNNING" { + action, err := classifyInstanceStatus(vm.Status) + if err != nil { + return err + } + switch action { + case instanceReady: + return m.setIp(vm) + case instanceStart: + if startRequested && !seenTransition { + // The start/resume API can become visible before the instance + // status changes. Avoid issuing the same mutation twice. + continue + } + if err := m.powerOn(ctx, vm.Status); err != nil { + return m.joinAfterPowerOnError(ctx, fmt.Errorf("could not power on after transitional state: %w", err)) + } + startRequested = true + seenTransition = false + case instanceWait: + seenTransition = true + continue + } + } + } +} + +// joinAfterPowerOnError handles the cross-process race where another Cloud +// Run revision starts or resumes the same VM after both observed a terminal +// state. A conflicting mutation is successful from PPB's perspective once the +// instance is observed transitioning or running. Permanent failures still +// return within a short bounded window. +func (m *GoogleComputeEngine) joinAfterPowerOnError(ctx context.Context, powerErr error) error { + timer := time.NewTimer(m.effectiveJoinTimeout()) + defer timer.Stop() + ticker := time.NewTicker(m.effectivePollInterval()) + defer ticker.Stop() + + for { + vm, err := m.getInstanceMetadata(ctx) + if err == nil { + action, classifyErr := classifyInstanceStatus(vm.Status) + if classifyErr != nil { + return classifyErr + } + switch action { + case instanceReady: return m.setIp(vm) + case instanceWait: + return m.waitForInstanceRunning(ctx, false) + case instanceStart: + // The competing operation might not be visible yet. } } + + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return powerErr + case <-ticker.C: + } + } +} + +func (m *GoogleComputeEngine) effectivePollInterval() time.Duration { + if m.pollInterval > 0 { + return m.pollInterval + } + return 2 * time.Second +} + +func (m *GoogleComputeEngine) effectiveJoinTimeout() time.Duration { + if m.joinTimeout > 0 { + return m.joinTimeout + } + return 10 * time.Second +} + +func (m *GoogleComputeEngine) currentTime() time.Time { + if m.now != nil { + return m.now() + } + return time.Now() +} + +func classifyInstanceStatus(status string) (instanceStatusAction, error) { + switch status { + case "RUNNING": + return instanceReady, nil + case "TERMINATED", "SUSPENDED": + return instanceStart, nil + case "PROVISIONING", "STAGING", "STOPPING", "SUSPENDING", "REPAIRING": + return instanceWait, nil + default: + return instanceWait, fmt.Errorf("unsupported instance status %q", status) } } @@ -168,22 +280,65 @@ func (m *GoogleComputeEngine) setIp(vm *compute.Instance) error { // PowerOnWithCooldown attempts to power on the machine if enough time has elapsed since the last attempt func (m *GoogleComputeEngine) PowerOnWithCooldown(ctx context.Context, cooldownSeconds int) error { - m.powerOnMutex.Lock() - defer m.powerOnMutex.Unlock() + if m.Lock == nil { + return fmt.Errorf("machine power-on lock is not initialized") + } + if err := m.Lock.Acquire(ctx, 1); err != nil { + return fmt.Errorf("wait for concurrent power-on attempt: %w", err) + } + defer m.Lock.Release(1) - // Check if we're still in cooldown period - if time.Since(m.LastPowerOnAttempt) < time.Duration(cooldownSeconds)*time.Second { + now := m.currentTime() + retryAt := m.LastPowerOnAttempt.Add(time.Duration(cooldownSeconds) * time.Second) + if !m.LastPowerOnAttempt.IsZero() && now.Before(retryAt) { slog.Debug("Power-on attempt skipped due to cooldown", "instance", m.Name, "cooldown", cooldownSeconds) - // Still check if we have a host IP set from previous attempts if m.Host() == "" { - return fmt.Errorf("backend not available and still in cooldown period") + return m.waitForCooldownTransition(ctx, retryAt) } return nil } // Update the last attempt time before making the API call - m.LastPowerOnAttempt = time.Now() + m.LastPowerOnAttempt = now slog.Debug("Attempting power-on check", "instance", m.Name) return m.PowerOn(ctx) } + +// waitForCooldownTransition lets a caller join an accepted start whose state +// change was not visible before the previous request was cancelled. It polls +// without issuing another mutation until the cooldown expires, then makes one +// fresh power-on attempt if the VM is still terminal. +func (m *GoogleComputeEngine) waitForCooldownTransition(ctx context.Context, retryAt time.Time) error { + ticker := time.NewTicker(m.effectivePollInterval()) + defer ticker.Stop() + + for { + vm, err := m.getInstanceMetadata(ctx) + if err == nil { + action, classifyErr := classifyInstanceStatus(vm.Status) + if classifyErr != nil { + return classifyErr + } + switch action { + case instanceReady: + return m.setIp(vm) + case instanceWait: + return m.waitForInstanceRunning(ctx, false) + case instanceStart: + // Wait out the mutation cooldown before retrying. + } + } + + if !m.currentTime().Before(retryAt) { + m.LastPowerOnAttempt = m.currentTime() + return m.PowerOn(ctx) + } + + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + } + } +} diff --git a/pkg/machine/gce_test.go b/pkg/machine/gce_test.go index 0846963..d48d5e5 100644 --- a/pkg/machine/gce_test.go +++ b/pkg/machine/gce_test.go @@ -2,6 +2,10 @@ package machine import ( "context" + "errors" + "strings" + "sync" + "sync/atomic" "testing" "testing/synctest" "time" @@ -9,170 +13,217 @@ import ( compute "google.golang.org/api/compute/v1" ) -func TestGoogleComputeEngine_PowerOnWithCooldown(t *testing.T) { - tests := []struct { - name string - cooldownSeconds int - initialHost string - timeBetweenCalls time.Duration - expectAPICall bool - expectError bool - needsSynctest bool - }{ - { - name: "first call should always make API call", - cooldownSeconds: 30, - timeBetweenCalls: 0, - expectAPICall: true, - expectError: false, - needsSynctest: false, - }, - { - name: "second call within cooldown should skip API call if host is set", - cooldownSeconds: 30, - initialHost: "10.0.0.1", - timeBetweenCalls: 10 * time.Second, - expectAPICall: false, - expectError: false, - needsSynctest: true, - }, - { - name: "second call within cooldown without host should return error", - cooldownSeconds: 30, - timeBetweenCalls: 10 * time.Second, - expectAPICall: false, - expectError: true, - needsSynctest: true, - }, - { - name: "call after cooldown period should make API call", - cooldownSeconds: 5, - timeBetweenCalls: 6 * time.Second, - expectAPICall: true, - expectError: false, - needsSynctest: true, - }, +func TestGoogleComputeEnginePowerOnFollowsStateSequence(t *testing.T) { + t.Parallel() + + statuses := []string{"TERMINATED", "STAGING", "RUNNING"} + var statusMu sync.Mutex + statusIndex := 0 + mutations := 0 + m := NewGceMachine() + m.UsePrivateIp = true + m.pollInterval = time.Millisecond + m.getInstanceHook = func(context.Context) (*compute.Instance, error) { + statusMu.Lock() + defer statusMu.Unlock() + index := statusIndex + if index < len(statuses)-1 { + statusIndex++ + } + return testInstance(statuses[index]), nil + } + m.powerOnHook = func(_ context.Context, status string) error { + mutations++ + if status != "TERMINATED" { + t.Fatalf("powerOn status = %q, want TERMINATED", status) + } + return nil } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - testFunc := func(t *testing.T) { - // Create machine instance - m := NewGceMachine() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := m.PowerOn(ctx); err != nil { + t.Fatalf("PowerOn() error = %v", err) + } + if mutations != 1 { + t.Fatalf("power mutations = %d, want 1", mutations) + } + if host := m.Host(); host != "10.42.0.8" { + t.Fatalf("Host() = %q, want 10.42.0.8", host) + } +} - // Set initial host if provided - if tt.initialHost != "" { - m.SetHostForTesting(tt.initialHost) - } +func TestGoogleComputeEngineJoinsConflictingMutation(t *testing.T) { + t.Parallel() - ctx := context.Background() + statuses := []string{"TERMINATED", "STAGING", "RUNNING"} + statusIndex := 0 + mutations := 0 + m := NewGceMachine() + m.UsePrivateIp = true + m.pollInterval = time.Millisecond + m.joinTimeout = 50 * time.Millisecond + m.getInstanceHook = func(context.Context) (*compute.Instance, error) { + index := statusIndex + if index < len(statuses)-1 { + statusIndex++ + } + return testInstance(statuses[index]), nil + } + m.powerOnHook = func(context.Context, string) error { + mutations++ + return errors.New("instance is already starting") + } - // Make first call (this will always attempt API call) - err1 := m.PowerOnWithCooldown(ctx, tt.cooldownSeconds) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := m.PowerOn(ctx); err != nil { + t.Fatalf("PowerOn() should join the competing transition: %v", err) + } + if mutations != 1 { + t.Fatalf("power mutations = %d, want 1", mutations) + } +} - if tt.timeBetweenCalls > 0 { - time.Sleep(tt.timeBetweenCalls) - } +func TestGoogleComputeEngineReturnsPermanentMutationError(t *testing.T) { + t.Parallel() - // Make second call to test cooldown behavior - err2 := m.PowerOnWithCooldown(ctx, tt.cooldownSeconds) + m := NewGceMachine() + m.pollInterval = time.Millisecond + m.joinTimeout = 5 * time.Millisecond + m.getInstanceHook = func(context.Context) (*compute.Instance, error) { + return testInstance("TERMINATED"), nil + } + m.powerOnHook = func(context.Context, string) error { + return errors.New("permission denied") + } - if tt.expectError && err2 == nil { - t.Errorf("Expected error on second call but got none") - } - if !tt.expectError && err2 != nil && tt.initialHost != "" { - t.Errorf("Unexpected error on second call with host set: %v", err2) - } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := m.PowerOn(ctx); err == nil || !strings.Contains(err.Error(), "permission denied") { + t.Fatalf("PowerOn() error = %v, want permanent mutation failure", err) + } +} - _ = err1 // Avoid unused variable warning - } +func TestGoogleComputeEngineCooldownJoinsAcceptedTransition(t *testing.T) { + t.Parallel() - if tt.needsSynctest { - synctest.Test(t, testFunc) - } else { - testFunc(t) - } - }) + statuses := []string{"STAGING", "RUNNING"} + statusIndex := 0 + mutations := 0 + m := NewGceMachine() + m.UsePrivateIp = true + m.LastPowerOnAttempt = time.Now() + m.pollInterval = time.Millisecond + m.getInstanceHook = func(context.Context) (*compute.Instance, error) { + index := statusIndex + if index < len(statuses)-1 { + statusIndex++ + } + return testInstance(statuses[index]), nil + } + m.powerOnHook = func(context.Context, string) error { + mutations++ + return nil + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := m.PowerOnWithCooldown(ctx, 30); err != nil { + t.Fatalf("PowerOnWithCooldown() error = %v", err) + } + if mutations != 0 { + t.Fatalf("power mutations = %d, want no duplicate mutation", mutations) } } -func TestGoogleComputeEngine_CooldownTimingLogic(t *testing.T) { +func TestGoogleComputeEngineCooldownRetriesTerminalStateAfterWindow(t *testing.T) { synctest.Test(t, func(t *testing.T) { + started := false + mutations := 0 m := NewGceMachine() - cooldownSeconds := 2 - - // Set a mock host so we don't get "backend not available" errors - m.SetHostForTesting("10.0.0.1") - - ctx := context.Background() - - // First call - should set LastPowerOnAttempt - err1 := m.PowerOnWithCooldown(ctx, cooldownSeconds) - firstAttemptTime := m.LastPowerOnAttempt - - // Immediate second call - should be in cooldown - err2 := m.PowerOnWithCooldown(ctx, cooldownSeconds) - secondAttemptTime := m.LastPowerOnAttempt - - // The attempt time should not have changed (still in cooldown) - if !firstAttemptTime.Equal(secondAttemptTime) { - t.Errorf("Second call should not update LastPowerOnAttempt during cooldown") + m.UsePrivateIp = true + m.LastPowerOnAttempt = time.Now() + m.getInstanceHook = func(context.Context) (*compute.Instance, error) { + if started { + return testInstance("RUNNING"), nil + } + return testInstance("TERMINATED"), nil } - - // Advance time beyond cooldown period using fake clock - time.Sleep(time.Duration(cooldownSeconds+1) * time.Second) - - // Third call - should update LastPowerOnAttempt - err3 := m.PowerOnWithCooldown(ctx, cooldownSeconds) - thirdAttemptTime := m.LastPowerOnAttempt - - // The attempt time should have been updated - if !thirdAttemptTime.After(firstAttemptTime) { - t.Errorf("Third call should update LastPowerOnAttempt after cooldown expires") + m.powerOnHook = func(context.Context, string) error { + mutations++ + started = true + return nil } - // Note: These calls will likely fail due to GCP API authentication in test environment - // but that's okay - we're testing the timing logic, not the actual GCP integration - _ = err1 - _ = err2 - _ = err3 + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if err := m.PowerOnWithCooldown(ctx, 2); err != nil { + t.Fatalf("PowerOnWithCooldown() error = %v", err) + } + if mutations != 1 { + t.Fatalf("power mutations = %d, want one retry after cooldown", mutations) + } }) } -func TestGoogleComputeEngine_ConcurrentCooldownCalls(t *testing.T) { +func TestGoogleComputeEngineWaitUsesCallerDeadline(t *testing.T) { synctest.Test(t, func(t *testing.T) { m := NewGceMachine() + m.getInstanceHook = func(context.Context) (*compute.Instance, error) { + return testInstance("STAGING"), nil + } + ctx, cancel := context.WithTimeout(context.Background(), 6*time.Minute) + defer cancel() - // Set a mock host so we don't get "backend not available" errors - m.SetHostForTesting("10.0.0.1") - - ctx := context.Background() - cooldownSeconds := 5 - - // Launch multiple concurrent calls - done := make(chan bool, 3) - - for i := 0; i < 3; i++ { - go func() { - _ = m.PowerOnWithCooldown(ctx, cooldownSeconds) - done <- true - }() + err := m.PowerOn(ctx) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("PowerOn() error = %v, want caller deadline", err) } + }) +} - // Wait for all goroutines to start and complete - synctest.Wait() +func TestGoogleComputeEngineConcurrentCooldownCallsUseCachedHost(t *testing.T) { + t.Parallel() - // Collect results - for i := 0; i < 3; i++ { - <-done - } + m := NewGceMachine() + m.SetHostForTesting("10.42.0.8") + m.LastPowerOnAttempt = time.Now() + var metadataCalls atomic.Int64 + m.getInstanceHook = func(context.Context) (*compute.Instance, error) { + metadataCalls.Add(1) + return testInstance("RUNNING"), nil + } - // Verify that the mutex protected the shared state properly - // (No specific assertion here, but the test would fail with race conditions if mutex wasn't working) - if m.LastPowerOnAttempt.IsZero() { - t.Error("Expected LastPowerOnAttempt to be set after concurrent calls") + var wg sync.WaitGroup + errorsCh := make(chan error, 8) + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + errorsCh <- m.PowerOnWithCooldown(context.Background(), 30) + }() + } + wg.Wait() + close(errorsCh) + for err := range errorsCh { + if err != nil { + t.Errorf("PowerOnWithCooldown() error = %v", err) } - }) + } + if calls := metadataCalls.Load(); calls != 0 { + t.Fatalf("metadata calls = %d, want cached host fast path", calls) + } +} + +func testInstance(status string) *compute.Instance { + return &compute.Instance{ + Name: "test-instance", + Status: status, + NetworkInterfaces: []*compute.NetworkInterface{{ + NetworkIP: "10.42.0.8", + }}, + } } func TestGoogleComputeEngine_setIp(t *testing.T) { @@ -332,3 +383,37 @@ func TestNewGceMachine(t *testing.T) { t.Error("NewGceMachine() semaphore should not allow acquiring more than 1") } } + +func TestClassifyInstanceStatus(t *testing.T) { + t.Parallel() + + tests := []struct { + status string + action instanceStatusAction + wantErr bool + }{ + {status: "RUNNING", action: instanceReady}, + {status: "TERMINATED", action: instanceStart}, + {status: "SUSPENDED", action: instanceStart}, + {status: "PROVISIONING", action: instanceWait}, + {status: "STAGING", action: instanceWait}, + {status: "STOPPING", action: instanceWait}, + {status: "SUSPENDING", action: instanceWait}, + {status: "REPAIRING", action: instanceWait}, + {status: "UNKNOWN", action: instanceWait, wantErr: true}, + } + + for _, test := range tests { + test := test + t.Run(test.status, func(t *testing.T) { + t.Parallel() + action, err := classifyInstanceStatus(test.status) + if (err != nil) != test.wantErr { + t.Fatalf("classifyInstanceStatus(%q) error = %v, wantErr %v", test.status, err, test.wantErr) + } + if action != test.action { + t.Fatalf("classifyInstanceStatus(%q) = %v, want %v", test.status, action, test.action) + } + }) + } +} diff --git a/pkg/proxy/proxy_test.go b/pkg/proxy/proxy_test.go index 7af21a7..47c1813 100644 --- a/pkg/proxy/proxy_test.go +++ b/pkg/proxy/proxy_test.go @@ -1,6 +1,18 @@ package proxy import ( + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strconv" + "strings" + "sync" + "syscall" "testing" "time" @@ -93,10 +105,6 @@ func TestNew_UsesConfiguredTimeouts(t *testing.T) { t.Fatal("New() returned nil") } - if proxy.Target.Scheme != tt.config.Scheme { - t.Errorf("Target.Scheme = %v, want %v", proxy.Target.Scheme, tt.config.Scheme) - } - transport := proxy.Transport if transport == nil { t.Fatal("Transport is nil") @@ -122,16 +130,14 @@ func TestNew_UsesConfiguredTimeouts(t *testing.T) { t.Error("DialContext is nil") } - if transport.DialTLSContext == nil { - t.Error("DialTLSContext is nil") - } }) } } -func TestReverseProxy_SetHost(t *testing.T) { +func TestReverseProxy_TargetURLUsesMachineHost(t *testing.T) { // Create a mock machine with a known host machine := machine.NewGceMachine() + machine.SetHostForTesting("10.0.0.8") config := &config.Config{ Scheme: "http", @@ -148,16 +154,16 @@ func TestReverseProxy_SetHost(t *testing.T) { } proxy := New(config) - proxy.SetHost() - - // Verify the target host format is correct when machine has a host - // This test is limited without being able to mock the machine host - if proxy.Target == nil { - t.Error("Target is nil after SetHost()") + target, err := proxy.targetURL() + if err != nil { + t.Fatalf("targetURL() error = %v", err) + } + if target.String() != "http://10.0.0.8:8080" { + t.Fatalf("targetURL() = %q, want http://10.0.0.8:8080", target) } } -func TestReverseProxy_SetHost_ProxyTargetOverride(t *testing.T) { +func TestReverseProxy_TargetURLUsesProxyTargetOverride(t *testing.T) { config := &config.Config{ Scheme: "http", Port: 8080, @@ -170,12 +176,383 @@ func TestReverseProxy_SetHost_ProxyTargetOverride(t *testing.T) { } proxy := New(config) - proxy.SetHost() + target, err := proxy.targetURL() + if err != nil { + t.Fatalf("targetURL() error = %v", err) + } + + if target.Host != "localhost:9000" { + t.Errorf("Target.Host = %q, want localhost:9000", target.Host) + } + if target.Scheme != "http" { + t.Errorf("Target.Scheme = %q, want http", target.Scheme) + } +} + +func TestRetryingDialerRetriesConnectionRefusalWithoutReplayingHTTP(t *testing.T) { + t.Parallel() + + clientConnection, serverConnection := net.Pipe() + t.Cleanup(func() { + _ = clientConnection.Close() + _ = serverConnection.Close() + }) + + var mu sync.Mutex + attempts := 0 + dialer := &retryingDialer{ + totalTimeout: time.Second, + attemptTimeout: 100 * time.Millisecond, + retryInterval: time.Millisecond, + dial: func(context.Context, string, string) (net.Conn, error) { + mu.Lock() + defer mu.Unlock() + attempts++ + if attempts < 3 { + return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED} + } + return clientConnection, nil + }, + } + + connection, err := dialer.DialContext(context.Background(), "tcp", "10.0.0.8:8080") + if err != nil { + t.Fatalf("DialContext() error = %v", err) + } + if connection != clientConnection { + t.Fatal("DialContext() did not return the successful connection") + } + mu.Lock() + defer mu.Unlock() + if attempts != 3 { + t.Fatalf("dial attempts = %d, want 3", attempts) + } +} + +func TestRetryingDialerFailsFastForPermanentDNSFailure(t *testing.T) { + t.Parallel() + + attempts := 0 + dialer := &retryingDialer{ + totalTimeout: time.Second, + attemptTimeout: 100 * time.Millisecond, + retryInterval: time.Millisecond, + dial: func(context.Context, string, string) (net.Conn, error) { + attempts++ + return nil, &net.DNSError{Err: "no such host", Name: "invalid", IsNotFound: true} + }, + } + + _, err := dialer.DialContext(context.Background(), "tcp", "invalid:8080") + if err == nil { + t.Fatal("DialContext() unexpectedly succeeded") + } + if attempts != 1 { + t.Fatalf("dial attempts = %d, want a fail-fast single attempt", attempts) + } +} - if proxy.Target.Host != "localhost:9000" { - t.Errorf("Target.Host = %q, want localhost:9000", proxy.Target.Host) +func TestRetryingDialerFailsFastForPermissionError(t *testing.T) { + t.Parallel() + + attempts := 0 + dialer := &retryingDialer{ + totalTimeout: time.Second, + attemptTimeout: 100 * time.Millisecond, + retryInterval: time.Millisecond, + dial: func(context.Context, string, string) (net.Conn, error) { + attempts++ + return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.EACCES} + }, + } + + _, err := dialer.DialContext(context.Background(), "tcp", "10.0.0.8:8080") + if err == nil { + t.Fatal("DialContext() unexpectedly succeeded") + } + if attempts != 1 { + t.Fatalf("dial attempts = %d, want a fail-fast single attempt", attempts) + } +} + +func TestRetryingDialerBoundsConnectionWindow(t *testing.T) { + t.Parallel() + + attempts := 0 + dialer := &retryingDialer{ + totalTimeout: 30 * time.Millisecond, + attemptTimeout: 5 * time.Millisecond, + retryInterval: time.Millisecond, + jitter: func(delay time.Duration) time.Duration { return delay }, + dial: func(context.Context, string, string) (net.Conn, error) { + attempts++ + return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED} + }, + } + started := time.Now() + _, err := dialer.DialContext(context.Background(), "tcp", "10.0.0.8:8080") + if err == nil { + t.Fatal("DialContext() unexpectedly succeeded") + } + if elapsed := time.Since(started); elapsed > 250*time.Millisecond { + t.Fatalf("DialContext() elapsed = %s, want a bounded retry window", elapsed) + } + if attempts < 2 { + t.Fatalf("dial attempts = %d, want retries within the bound", attempts) + } +} + +func TestRetryingDialerHonorsRequestCancellation(t *testing.T) { + t.Parallel() + + dialer := &retryingDialer{ + totalTimeout: time.Second, + attemptTimeout: time.Second, + retryInterval: time.Millisecond, + dial: func(ctx context.Context, _, _ string) (net.Conn, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := dialer.DialContext(ctx, "tcp", "10.0.0.8:8080") + if !errors.Is(err, context.Canceled) { + t.Fatalf("DialContext() error = %v, want context cancellation", err) + } +} + +func TestReverseProxyKeepsForwardedHeadersRequestLocal(t *testing.T) { + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = fmt.Fprintf(w, "%s|%s|%s", r.Header.Get("X-Forwarded-Host"), r.Header.Get("X-Forwarded-For"), r.Header.Get("X-Cloud-Trace-Context")) + })) + t.Cleanup(backend.Close) + + backendURL, err := url.Parse(backend.URL) + if err != nil { + t.Fatalf("url.Parse() error = %v", err) + } + backendHost, backendPortText, err := net.SplitHostPort(backendURL.Host) + if err != nil { + t.Fatalf("net.SplitHostPort() error = %v", err) + } + backendPort, err := strconv.Atoi(backendPortText) + if err != nil { + t.Fatalf("strconv.Atoi() error = %v", err) + } + + proxyHandler := New(&config.Config{ + Scheme: "http", + Port: backendPort, + ProxyTarget: &config.ProxyTarget{ + Scheme: "http", + Host: backendHost, + Port: backendPort, + }, + ProxyTimeouts: config.ProxyTimeouts{ + DialTimeout: 2, + DialAttemptTimeout: 1, + DialRetryInterval: 1, + KeepAlive: 1, + IdleConnTimeout: 1, + TLSHandshakeTimeout: 1, + ExpectContinueTimeout: 1, + MaxIdleConns: 100, + }, + Machine: machine.NewGceMachine(), + }) + frontend := httptest.NewServer(proxyHandler) + t.Cleanup(frontend.Close) + + const requestCount = 32 + errorsCh := make(chan error, requestCount) + var wg sync.WaitGroup + for index := 0; index < requestCount; index++ { + index := index + wg.Add(1) + go func() { + defer wg.Done() + want := fmt.Sprintf("tenant-%d.example|192.0.2.%d|trace-%d", index, index+1, index) + request, err := http.NewRequest(http.MethodGet, frontend.URL, nil) + if err != nil { + errorsCh <- err + return + } + request.Host = fmt.Sprintf("tenant-%d.example", index) + request.Header.Set("X-Forwarded-For", fmt.Sprintf("192.0.2.%d", index+1)) + request.Header.Set("X-Cloud-Trace-Context", fmt.Sprintf("trace-%d", index)) + response, err := frontend.Client().Do(request) + if err != nil { + errorsCh <- err + return + } + defer response.Body.Close() + body, err := io.ReadAll(response.Body) + if err != nil { + errorsCh <- err + return + } + if got := strings.TrimSpace(string(body)); got != want { + errorsCh <- fmt.Errorf("forwarded headers = %q, want %q", got, want) + } + }() + } + wg.Wait() + close(errorsCh) + for err := range errorsCh { + if err != nil { + t.Error(err) + } + } +} + +func TestReverseProxyDialRetrySendsPostBodyExactlyOnce(t *testing.T) { + var mu sync.Mutex + requestCount := 0 + requestBody := "" + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + mu.Lock() + requestCount++ + requestBody += string(body) + mu.Unlock() + w.WriteHeader(http.StatusNoContent) + })) + t.Cleanup(backend.Close) + + backendURL, err := url.Parse(backend.URL) + if err != nil { + t.Fatalf("url.Parse() error = %v", err) + } + backendHost, backendPortText, err := net.SplitHostPort(backendURL.Host) + if err != nil { + t.Fatalf("net.SplitHostPort() error = %v", err) + } + backendPort, err := strconv.Atoi(backendPortText) + if err != nil { + t.Fatalf("strconv.Atoi() error = %v", err) + } + + proxyHandler := New(&config.Config{ + Scheme: "http", + Port: backendPort, + ProxyTarget: &config.ProxyTarget{ + Scheme: "http", + Host: backendHost, + Port: backendPort, + }, + ProxyTimeouts: config.ProxyTimeouts{ + DialTimeout: 2, + DialAttemptTimeout: 1, + DialRetryInterval: 1, + KeepAlive: 1, + IdleConnTimeout: 1, + TLSHandshakeTimeout: 1, + ExpectContinueTimeout: 1, + MaxIdleConns: 10, + }, + Machine: machine.NewGceMachine(), + }) + realDialer := &net.Dialer{} + dialAttempts := 0 + proxyHandler.Transport.DialContext = (&retryingDialer{ + totalTimeout: time.Second, + attemptTimeout: 100 * time.Millisecond, + retryInterval: time.Millisecond, + jitter: func(delay time.Duration) time.Duration { return delay }, + dial: func(ctx context.Context, network, address string) (net.Conn, error) { + dialAttempts++ + if dialAttempts < 3 { + return nil, &net.OpError{Op: "dial", Net: network, Err: syscall.ECONNREFUSED} + } + return realDialer.DialContext(ctx, network, address) + }, + }).DialContext + + frontend := httptest.NewServer(proxyHandler) + t.Cleanup(frontend.Close) + response, err := http.Post(frontend.URL, "text/plain", strings.NewReader("one-copy")) + if err != nil { + t.Fatalf("http.Post() error = %v", err) + } + response.Body.Close() + if response.StatusCode != http.StatusNoContent { + t.Fatalf("response status = %d, want %d", response.StatusCode, http.StatusNoContent) + } + + mu.Lock() + defer mu.Unlock() + if requestCount != 1 || requestBody != "one-copy" { + t.Fatalf("backend received count=%d body=%q, want one exact request", requestCount, requestBody) + } + if dialAttempts != 3 { + t.Fatalf("dial attempts = %d, want 3", dialAttempts) + } +} + +func TestReverseProxyDoesNotRetryHTTPStatus(t *testing.T) { + var mu sync.Mutex + requestCount := 0 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + mu.Lock() + requestCount++ + mu.Unlock() + http.Error(w, "try later", http.StatusServiceUnavailable) + })) + t.Cleanup(backend.Close) + + backendURL, err := url.Parse(backend.URL) + if err != nil { + t.Fatalf("url.Parse() error = %v", err) + } + backendHost, backendPortText, err := net.SplitHostPort(backendURL.Host) + if err != nil { + t.Fatalf("net.SplitHostPort() error = %v", err) + } + backendPort, err := strconv.Atoi(backendPortText) + if err != nil { + t.Fatalf("strconv.Atoi() error = %v", err) + } + + proxyHandler := New(&config.Config{ + Scheme: "http", + Port: backendPort, + ProxyTarget: &config.ProxyTarget{ + Scheme: "http", + Host: backendHost, + Port: backendPort, + }, + ProxyTimeouts: config.ProxyTimeouts{ + DialTimeout: 2, + DialAttemptTimeout: 1, + DialRetryInterval: 1, + KeepAlive: 1, + IdleConnTimeout: 1, + TLSHandshakeTimeout: 1, + ExpectContinueTimeout: 1, + MaxIdleConns: 10, + }, + Machine: machine.NewGceMachine(), + }) + frontend := httptest.NewServer(proxyHandler) + t.Cleanup(frontend.Close) + + response, err := frontend.Client().Get(frontend.URL) + if err != nil { + t.Fatalf("Get() error = %v", err) + } + response.Body.Close() + if response.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("response status = %d, want %d", response.StatusCode, http.StatusServiceUnavailable) } - if proxy.Target.Scheme != "http" { - t.Errorf("Target.Scheme = %q, want http", proxy.Target.Scheme) + mu.Lock() + defer mu.Unlock() + if requestCount != 1 { + t.Fatalf("backend request count = %d, want 1", requestCount) } } diff --git a/pkg/proxy/tls.go b/pkg/proxy/tls.go index 4d2b05f..72c6316 100644 --- a/pkg/proxy/tls.go +++ b/pkg/proxy/tls.go @@ -2,57 +2,56 @@ package proxy import ( "context" - "crypto/tls" + cryptorand "crypto/rand" + "errors" + "fmt" "log/slog" "net" "net/http" "net/http/httputil" "net/url" "strconv" + "syscall" "time" "github.com/libops/ppb/pkg/config" ) type ReverseProxy struct { - Target *url.URL Transport *http.Transport Config *config.Config - Host []string - Ip []string - Trace []string +} + +type retryingDialer struct { + totalTimeout time.Duration + attemptTimeout time.Duration + retryInterval time.Duration + keepAlive time.Duration + dial func(context.Context, string, string) (net.Conn, error) + jitter func(time.Duration) time.Duration } func New(c *config.Config) *ReverseProxy { // Use configured timeout values (defaults already set in config loading) dialTimeout := time.Duration(c.ProxyTimeouts.DialTimeout) * time.Second + dialAttemptTimeout := time.Duration(c.ProxyTimeouts.DialAttemptTimeout) * time.Second + dialRetryInterval := time.Duration(c.ProxyTimeouts.DialRetryInterval) * time.Second keepAlive := time.Duration(c.ProxyTimeouts.KeepAlive) * time.Second idleConnTimeout := time.Duration(c.ProxyTimeouts.IdleConnTimeout) * time.Second tlsHandshakeTimeout := time.Duration(c.ProxyTimeouts.TLSHandshakeTimeout) * time.Second expectContinueTimeout := time.Duration(c.ProxyTimeouts.ExpectContinueTimeout) * time.Second - scheme := c.Scheme - if c.ProxyTarget != nil && c.ProxyTarget.Scheme != "" { - scheme = c.ProxyTarget.Scheme + dialer := &retryingDialer{ + totalTimeout: dialTimeout, + attemptTimeout: dialAttemptTimeout, + retryInterval: dialRetryInterval, + keepAlive: keepAlive, } return &ReverseProxy{ - Target: &url.URL{ - Scheme: scheme, - }, Config: c, Transport: &http.Transport{ - Proxy: http.ProxyFromEnvironment, - DialContext: (&net.Dialer{ - Timeout: dialTimeout, - KeepAlive: keepAlive, - }).DialContext, - DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) { - dialer := &net.Dialer{ - Timeout: dialTimeout, - KeepAlive: keepAlive, - } - return tls.DialWithDialer(dialer, network, addr, nil) - }, + Proxy: http.ProxyFromEnvironment, + DialContext: dialer.DialContext, ForceAttemptHTTP2: true, MaxIdleConns: c.ProxyTimeouts.MaxIdleConns, IdleConnTimeout: idleConnTimeout, @@ -62,47 +61,180 @@ func New(c *config.Config) *ReverseProxy { } } -func (p *ReverseProxy) SetHost() { +func (p *ReverseProxy) targetURL() (*url.URL, error) { + scheme := p.Config.Scheme + if p.Config.ProxyTarget != nil && p.Config.ProxyTarget.Scheme != "" { + scheme = p.Config.ProxyTarget.Scheme + } + if scheme != "http" && scheme != "https" { + return nil, fmt.Errorf("unsupported proxy target scheme %q", scheme) + } + if p.Config.ProxyTarget != nil && p.Config.ProxyTarget.Host != "" { port := p.Config.ProxyTarget.Port if port == 0 { port = p.Config.Port } - p.Target.Host = net.JoinHostPort(p.Config.ProxyTarget.Host, strconv.Itoa(port)) - slog.Debug("Set proxy target host", "host", p.Target.Host) - return + return &url.URL{ + Scheme: scheme, + Host: net.JoinHostPort(p.Config.ProxyTarget.Host, strconv.Itoa(port)), + }, nil + } + + host := p.Config.Machine.Host() + if host == "" { + return nil, errors.New("machine does not have a proxy target IP") } - p.Target.Host = net.JoinHostPort( - p.Config.Machine.Host(), - strconv.Itoa(p.Config.Port), - ) - slog.Debug("Set machine host", "host", p.Target.Host) + return &url.URL{ + Scheme: scheme, + Host: net.JoinHostPort(host, strconv.Itoa(p.Config.Port)), + }, nil } func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { + target, err := p.targetURL() + if err != nil { + slog.Warn("Proxy target is unavailable", "error", err) + w.Header().Set("Retry-After", "5") + http.Error(w, "Backend not available", http.StatusServiceUnavailable) + return + } + + forwardedFor := r.Header.Get("X-Forwarded-For") + forwardedHost := r.Host + trace := r.Header.Get("X-Cloud-Trace-Context") + rp := &httputil.ReverseProxy{ Transport: p.Transport, Rewrite: func(pr *httputil.ProxyRequest) { - pr.SetURL(p.Target) - pr.Out.Header["X-Cloud-Trace-Context"] = p.Trace - pr.Out.Header["X-Forwarded-For"] = p.Ip - pr.Out.Header["X-Forwarded-Host"] = p.Host - pr.Out.Header["X-Forwarded-Proto"] = []string{"https"} + pr.SetURL(target) + setOrDeleteHeader(pr.Out.Header, "X-Cloud-Trace-Context", trace) + setOrDeleteHeader(pr.Out.Header, "X-Forwarded-For", forwardedFor) + setOrDeleteHeader(pr.Out.Header, "X-Forwarded-Host", forwardedHost) + pr.Out.Header.Set("X-Forwarded-Proto", "https") + }, + ErrorHandler: func(w http.ResponseWriter, _ *http.Request, err error) { + slog.Warn("Backend proxy request failed", "target", target.Redacted(), "error", err) + w.Header().Set("Retry-After", "5") + http.Error(w, "Backend not available", http.StatusServiceUnavailable) }, } rp.ServeHTTP(w, r) } -func (p *ReverseProxy) SetRequestHeaders(r *http.Request) { - p.Ip = []string{ - r.Header.Get("X-Forwarded-For"), +func setOrDeleteHeader(header http.Header, name, value string) { + if value == "" { + header.Del(name) + return } - p.Host = []string{ - r.Host, + header.Set(name, value) +} + +func (d *retryingDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) { + retryCtx, cancel := context.WithTimeout(ctx, d.totalTimeout) + defer cancel() + + dial := d.dial + if dial == nil { + networkDialer := &net.Dialer{KeepAlive: d.keepAlive} + dial = networkDialer.DialContext } - p.Trace = []string{ - r.Header.Get("X-Cloud-Trace-Context"), + + var lastErr error + retryDelay := d.retryInterval + for attempt := 1; ; attempt++ { + attemptTimeout := d.attemptTimeout + if deadline, ok := retryCtx.Deadline(); ok && time.Until(deadline) < attemptTimeout { + attemptTimeout = time.Until(deadline) + } + if attemptTimeout <= 0 { + return nil, d.exhaustedError(ctx, address, lastErr) + } + + attemptCtx, attemptCancel := context.WithTimeout(retryCtx, attemptTimeout) + connection, err := dial(attemptCtx, network, address) + attemptCancel() + if err == nil { + return connection, nil + } + lastErr = err + if retryCtx.Err() != nil { + return nil, d.exhaustedError(ctx, address, lastErr) + } + if !isRetryableDialError(err) { + return nil, err + } + + slog.Debug("Backend connection attempt failed; retrying", "address", address, "attempt", attempt, "error", err) + jitter := d.jitter + if jitter == nil { + jitter = jitteredDelay + } + timer := time.NewTimer(jitter(retryDelay)) + select { + case <-retryCtx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + return nil, d.exhaustedError(ctx, address, lastErr) + case <-timer.C: + } + if retryDelay < 5*time.Second { + retryDelay *= 2 + if retryDelay > 5*time.Second { + retryDelay = 5 * time.Second + } + } + } +} + +func (d *retryingDialer) exhaustedError(parent context.Context, address string, lastErr error) error { + if err := parent.Err(); err != nil { + return err + } + if lastErr == nil { + return fmt.Errorf("backend %s did not accept a connection within %s", address, d.totalTimeout) + } + return fmt.Errorf("backend %s did not accept a connection within %s: %w", address, d.totalTimeout, lastErr) +} + +func isRetryableDialError(err error) bool { + if errors.Is(err, context.Canceled) { + return false + } + if errors.Is(err, context.DeadlineExceeded) { + return true + } + var addressError *net.AddrError + if errors.As(err, &addressError) { + return false + } + var dnsError *net.DNSError + if errors.As(err, &dnsError) && !dnsError.IsTimeout && !dnsError.IsTemporary { + return false + } + var networkError net.Error + if errors.As(err, &networkError) && networkError.Timeout() { + return true + } + return errors.Is(err, syscall.ECONNREFUSED) || + errors.Is(err, syscall.ECONNRESET) || + errors.Is(err, syscall.ETIMEDOUT) || + errors.Is(err, syscall.EHOSTUNREACH) || + errors.Is(err, syscall.ENETUNREACH) +} + +func jitteredDelay(delay time.Duration) time.Duration { + // Spread simultaneous cold-start requests without making the configured + // interval an unbounded delay. The backoff remains within 75%-125%. + var randomByte [1]byte + if _, err := cryptorand.Read(randomByte[:]); err != nil { + return delay } - slog.Debug("Request headers", "p.Ip", p.Ip, "p.Host", p.Host) + factor := 0.75 + (float64(randomByte[0])/255.0)*0.5 + return time.Duration(float64(delay) * factor) } diff --git a/ppb.example.yaml b/ppb.example.yaml index 46afde0..1794673 100644 --- a/ppb.example.yaml +++ b/ppb.example.yaml @@ -1,13 +1,18 @@ type: google_compute_engine port: 80 scheme: http -allowedIps: ${ALLOWED_IPS:-127.0.0.1/32,10.0.0.0/8,172.16.0.0/12,192.168.0.0/16,::1/128} -ipForwardedHeader: X-Forwarded-For +allowedIps: + - 127.0.0.1/32 + - ::1/128 +ipForwardedHeader: "" ipDepth: 0 -powerOnCooldown: ${POWER_ON_COOLDOWN:-30} # seconds between power-on attempts +powerOnCooldown: 30 # seconds between power-on attempts +powerOnTimeout: 360 # total queued/API/instance-ready budget # Optional proxy timeout configuration (defaults shown) proxyTimeouts: - dialTimeout: 120 # Connection dial timeout in seconds + dialTimeout: 120 # Total connection retry window in seconds + dialAttemptTimeout: 5 # Timeout for one connection attempt in seconds + dialRetryInterval: 1 # Delay between connection attempts in seconds keepAlive: 120 # TCP keep-alive timeout in seconds idleConnTimeout: 90 # Idle connection timeout in seconds tlsHandshakeTimeout: 10 # TLS handshake timeout in seconds @@ -18,4 +23,4 @@ machineMetadata: zone: us-central1-f name: compose-librechat # usePrivateIp=true means cloud run is using direct VPC egress to the machine - usePrivateIp: ${PRIVATE_IP:-false} + usePrivateIp: false