Add lambda janitor plugin (#1569)

Signed-off-by: Alex Demicev <alex.demicev@lambdal.com>

Author: alexander-demicev

commons/pkg/lambda/client.go

@@ -17,8 +17,10 @@
 package lambda
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
@@ -41,22 +43,32 @@ const (
 	defaultInitialBackoff = 500 * time.Millisecond
 	defaultBackoffFactor  = 2.0
 	defaultBackoffJitter  = 0.2
+
+	// maxRedirects mirrors net/http's own cap, which is dropped the moment
+	// CheckRedirect is set.
+	maxRedirects = 10
 )
 
+// errRedirect marks a redirect this client refused to follow. A redirect chain
+// is deterministic, so retrying only replays it.
+var errRedirect = errors.New("redirect refused")
+
 // Client is an authenticated HTTP client for the Lambda Cloud API.
 // The API key is read from LAMBDA_API_KEY on every request so credential
 // rotation works without a process restart.
 //
 // Requests are retried with exponential backoff on transient failures
 // (network errors, 5xx responses, and 429 Too Many Requests). Permanent
 // failures (4xx other than 429, malformed responses, missing API key)
-// short-circuit the retry loop.
+// short-circuit the retry loop. Post retries on less, see retryRateLimitOnly.
 type Client struct {
 	endpoint string
 	http     *http.Client
 	retry    retryPolicy
 }
 
+// retryPolicy is the exponential-backoff schedule applied to a request.
+// maxAttempts counts the initial attempt.
 type retryPolicy struct {
 	maxAttempts    int
 	initialBackoff time.Duration
@@ -114,13 +126,70 @@ func NewClient(endpoint string, opts ...Option) *Client {
 		o(c)
 	}
 
+	// Wrapped rather than assigned outright, and after the options so an
+	// injected client cannot opt out of the scheme check.
+	c.http.CheckRedirect = refuseInsecureRedirect(c.http.CheckRedirect)
+
 	return c
 }
 
+// redirectPolicy is net/http's CheckRedirect signature.
+type redirectPolicy func(req *http.Request, via []*http.Request) error
+
+// refuseInsecureRedirect wraps next with a scheme check. net/http copies
+// Authorization to the same host or a subdomain without looking at the scheme,
+// so a downgrade would hand the API key to a plaintext listener.
+//
+// Setting CheckRedirect at all replaces net/http's ten-redirect cap, so with no
+// next to defer to this reimposes it. Without that an https to https loop runs
+// until the client timeout.
+func refuseInsecureRedirect(next redirectPolicy) redirectPolicy {
+	return func(req *http.Request, via []*http.Request) error {
+		if req.URL.Scheme != "https" {
+			return fmt.Errorf("%w: %s would send the API key over %s",
+				errRedirect, req.URL.Redacted(), req.URL.Scheme)
+		}
+
+		if next != nil {
+			return next(req, via)
+		}
+
+		if len(via) >= maxRedirects {
+			return fmt.Errorf("%w: stopped after %d redirects", errRedirect, maxRedirects)
+		}
+
+		return nil
+	}
+}
+
 // Get performs an authenticated GET against endpoint+path with optional query
 // params and decodes the JSON response body into out. If out is nil, the body
 // is discarded. Retries transient failures with exponential backoff.
 func (c *Client) Get(ctx context.Context, path string, query url.Values, out any) error {
+	return c.do(ctx, http.MethodGet, path, query, nil, out, retryTransient)
+}
+
+// Post performs an authenticated POST against endpoint+path with in marshalled
+// as the JSON request body, and decodes the JSON response into out.
+//
+// It serves the instance-operations endpoints, which are not idempotent and
+// take no idempotency key, so it retries on less than Get does. See
+// retryRateLimitOnly.
+func (c *Client) Post(ctx context.Context, path string, in, out any) error {
+	// Marshalled once so every retry resends identical bytes.
+	payload, err := json.Marshal(in)
+	if err != nil {
+		return fmt.Errorf("marshal request body: %w", err)
+	}
+
+	return c.do(ctx, http.MethodPost, path, nil, payload, out, retryRateLimitOnly)
+}
+
+// do performs an authenticated request and decodes the JSON response into out.
+// payload is nil for methods without a request body.
+func (c *Client) do(
+	ctx context.Context, method, path string, query url.Values, payload []byte, out any, retry retryScope,
+) error {
 	apiKey := os.Getenv(APIKeyEnvVar)
 	if apiKey == "" {
 		return fmt.Errorf("env var %s is not set", APIKeyEnvVar)
@@ -146,27 +215,32 @@ func (c *Client) Get(ctx context.Context, path string, query url.Values, out any
 	err := wait.ExponentialBackoffWithContext(ctx, backoff, func(ctx context.Context) (bool, error) {
 		attempts++
 
-		body, statusCode, doErr := c.doOnce(ctx, u, apiKey)
+		body, statusCode, doErr := c.doOnce(ctx, method, u, apiKey, payload)
 		if doErr != nil {
-			// Transport-level failures (dial, TLS, i/o) are transient.
+			// Transport-level failures (dial, TLS, i/o).
+			if !retry(statusCode, doErr) {
+				return false, doErr
+			}
+
 			lastErr = doErr
 			slog.Debug("Lambda API request failed, will retry", "url", u, "error", doErr)
 
 			return false, nil
 		}
 
-		if statusCode == http.StatusTooManyRequests || statusCode >= 500 {
-			lastErr = fmt.Errorf("GET %s: status %d: %s", u, statusCode, body)
+		if statusCode != http.StatusOK {
+			statusErr := fmt.Errorf("%s %s: status %d: %s", method, u, statusCode, body)
+			if !retry(statusCode, nil) {
+				return false, statusErr
+			}
+
+			lastErr = statusErr
+
 			slog.Debug("Lambda API returned retryable status", "url", u, "status", statusCode)
 
 			return false, nil
 		}
 
-		if statusCode != http.StatusOK {
-			// Permanent client error (401, 403, 404, ...) — do not retry.
-			return false, fmt.Errorf("GET %s: status %d: %s", u, statusCode, body)
-		}
-
 		return true, decodeBody(body, out)
 	})
 	if err == nil {
@@ -183,8 +257,32 @@ func (c *Client) Get(ctx context.Context, path string, query url.Values, out any
 	return err
 }
 
+// retryScope reports whether a failed attempt can be repeated. transportErr is
+// nil when the request completed, in which case statusCode holds the response
+// status.
+type retryScope func(statusCode int, transportErr error) bool
+
+// retryTransient retries network errors, 429, and 5xx. Safe for a request that
+// can be repeated without side effects. A refused redirect is excluded: the
+// chain is deterministic, so a repeat earns the same refusal.
+func retryTransient(statusCode int, transportErr error) bool {
+	if errors.Is(transportErr, errRedirect) {
+		return false
+	}
+
+	return transportErr != nil || statusCode == http.StatusTooManyRequests || statusCode >= 500
+}
+
+// retryRateLimitOnly retries only 429, where the rate limiter rejected the
+// request without acting on it. A transport error or 5xx is ambiguous, the
+// request may already have taken effect so it surfaces on the first attempt
+// rather than resubmitting an operation that cannot be undone.
+func retryRateLimitOnly(statusCode int, transportErr error) bool {
+	return transportErr == nil && statusCode == http.StatusTooManyRequests
+}
+
 // decodeBody unmarshals body into out when out is non-nil, otherwise discards
-// body. Extracted from Get so the retry closure stays under the cyclomatic
+// body. Extracted from do so the retry closure stays under the cyclomatic
 // complexity limit.
 func decodeBody(body []byte, out any) error {
 	if out == nil {
@@ -201,18 +299,27 @@ func decodeBody(body []byte, out any) error {
 // doOnce performs a single request. It returns the response body and status
 // code separately so the retry loop can classify the outcome without having
 // to keep the *http.Response around.
-func (c *Client) doOnce(ctx context.Context, u, apiKey string) ([]byte, int, error) {
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
+func (c *Client) doOnce(ctx context.Context, method, u, apiKey string, payload []byte) ([]byte, int, error) {
+	var reqBody io.Reader
+	if payload != nil {
+		reqBody = bytes.NewReader(payload)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, method, u, reqBody)
 	if err != nil {
 		return nil, 0, fmt.Errorf("build request: %w", err)
 	}
 
 	req.Header.Set("Accept", "application/json")
 	req.Header.Set("Authorization", "Bearer "+apiKey)
 
+	if payload != nil {
+		req.Header.Set("Content-Type", "application/json")
+	}
+
 	resp, err := c.http.Do(req)
 	if err != nil {
-		return nil, 0, fmt.Errorf("GET %s: %w", u, err)
+		return nil, 0, fmt.Errorf("%s %s: %w", method, u, err)
 	}
 	defer resp.Body.Close()