[codex] refactor gateway paid-request verification flow

gateway/cache.go

@@ -12,7 +12,6 @@ import (
 	"log"
 	"net/http"
 	"os"
-	"strconv"
 	"sync"
 	"time"
 
@@ -123,58 +122,18 @@ func CacheMiddleware() gin.HandlerFunc {
 			}
 			cacheHits.WithLabelValues(routePath).Inc()
 
-			// Cache HIT! -> Verify Payment *BEFORE* serving
-			// verifyPayment creates its own timeout context, so pass request context directly
-			timestampStr := c.GetHeader("X-402-Timestamp")
-			if timestampStr == "" {
-				respondError(c, 400, "invalid_timestamp", fmt.Errorf("missing X-402-Timestamp header"))
+			payment, ok := verifyPaidRequest(c)
+			if !ok {
 				c.Abort()
 				return
 			}
-			timestamp, err := strconv.ParseUint(timestampStr, 10, 64)
-			if err != nil || timestamp == 0 {
-				respondError(c, 400, "invalid_timestamp", fmt.Errorf("invalid X-402-Timestamp header"))
-				c.Abort()
-				return
-			}
-			verifyResp, paymentCtx, err := verifyPayment(c.Request.Context(), signature, nonce, timestamp)
-			if err != nil {
-
-				verificationTotal.WithLabelValues("error").Inc()
-				log.Printf("Verification error on cache hit: %v", err)
-
-				if errors.Is(err, context.DeadlineExceeded) {
-					respondError(c, 504, "verifier_timeout", err)
-				} else {
-					respondError(c, 502, "verification_unavailable", err)
-				}
-				c.Abort()
-				return
-			}
-
-			if !verifyResp.IsValid {
-				verificationTotal.WithLabelValues("invalid").Inc()
-				respondVerificationFailure(c, verifyResp)
-				c.Abort()
-				return
-			}
-			if verifyResp.RecoveredAddress == "" {
-				verificationTotal.WithLabelValues("error").Inc()
-				respondError(c, 502, "verification_unavailable", fmt.Errorf("verifier success missing recovered_address"))
-				c.Abort()
-				return
-			}
-			verificationTotal.WithLabelValues("success").Inc()
-			// Payment Verified. Store verification for downstream if needed (though we abort)
-			c.Set("payment_verification", verifyResp)
-			c.Set("payment_context", paymentCtx)
 
 			// Generate Receipt and Respond
 			// We treat the cached result as the AI result
 			// Generate receipt for cache hit using current request and cached result.
 			// Note: request_hash matches current request, response is from cache,
 			// but both are cryptographically valid since cache key ensures identical text.
-			if err := generateAndSendReceipt(c, *paymentCtx, verifyResp.RecoveredAddress, requestBody, cached.Result); err != nil {
+			if err := sendPaidResult(c, payment, requestBody, cached.Result); err != nil {
 				log.Printf("Failed to send cached response receipt: %v", err)
 				// generateAndSendReceipt already sent an error response (500)
 			}

gateway/main.go

@@ -427,33 +427,11 @@ func main() {
 // applied by middleware and returns appropriate HTTP errors (402, 403, 504,
 // 500) to the client.
 func handleSummarize(c *gin.Context) {
-	// 1. Payment Verification
-	// Note: CacheMiddleware aborts on cache HIT, so this handler only runs on cache MISS or when caching is disabled
 	var requestBody []byte
 	var err error
 
-	signature := c.GetHeader("X-402-Signature")
-	nonce := c.GetHeader("X-402-Nonce")
-	timestampHeader := c.GetHeader("X-402-Timestamp")
-
-	// Basic check
-	if signature == "" || nonce == "" {
-		c.JSON(402, gin.H{
-			"error":          "Payment Required",
-			"message":        "Please sign the payment context",
-			"paymentContext": createPaymentContext(),
-		})
-		return
-	}
-
-	if timestampHeader == "" {
-		respondError(c, 400, "invalid_timestamp", fmt.Errorf("missing X-402-Timestamp header"))
-		return
-	}
-
-	timestampValue, err := strconv.ParseUint(timestampHeader, 10, 64)
-	if err != nil || timestampValue == 0 {
-		respondError(c, 400, "invalid_timestamp", fmt.Errorf("invalid X-402-Timestamp header"))
+	payment, ok := verifyPaidRequest(c)
+	if !ok {
 		return
 	}
 
@@ -480,33 +458,6 @@ func handleSummarize(c *gin.Context) {
 		}
 	}
 
-	// Verify
-	verifyResp, paymentCtx, err := verifyPayment(c.Request.Context(), signature, nonce, uint64(timestampValue))
-	if err != nil {
-		verificationTotal.WithLabelValues("error").Inc()
-
-		if errors.Is(err, context.DeadlineExceeded) {
-			respondError(c, 504, "verifier_timeout", err)
-		} else {
-			respondError(c, 502, "verification_unavailable", err)
-		}
-		return
-	}
-
-	if !verifyResp.IsValid {
-		verificationTotal.WithLabelValues("invalid").Inc()
-
-		respondVerificationFailure(c, verifyResp)
-		return
-	}
-	if verifyResp.RecoveredAddress == "" {
-		verificationTotal.WithLabelValues("error").Inc()
-		respondError(c, 502, "verification_unavailable", fmt.Errorf("verifier success missing recovered_address"))
-		return
-	}
-
-	verificationTotal.WithLabelValues("success").Inc()
-
 	// 2. Parse Request
 	var req SummarizeRequest
 	if err := json.Unmarshal(requestBody, &req); err != nil {
@@ -532,11 +483,8 @@ func handleSummarize(c *gin.Context) {
 	}
 
 	// 4. Generate & Send Receipt
-	if err := generateAndSendReceipt(c, *paymentCtx, verifyResp.RecoveredAddress, requestBody, summary); err != nil {
+	if err := sendPaidResult(c, payment, requestBody, summary); err != nil {
 		log.Printf("Failed to generate receipt: %v", err)
-		// generateAndSendReceipt sends error response if it fails?
-		// No, it returns error, we might have already written status if we aren't careful.
-		// Let's implement generateAndSendReceipt to handle sending response.
 		return
 	}
 }

gateway/payment_flow.go

@@ -0,0 +1,91 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"strconv"
+
+	"github.com/gin-gonic/gin"
+)
+
+type verifiedPayment struct {
+	PaymentContext   PaymentContext
+	RecoveredAddress string
+}
+
+func verifyPaidRequest(c *gin.Context) (*verifiedPayment, bool) {
+	signature := c.GetHeader("X-402-Signature")
+	nonce := c.GetHeader("X-402-Nonce")
+
+	if signature == "" || nonce == "" {
+		c.JSON(http.StatusPaymentRequired, gin.H{
+			"error":          "Payment Required",
+			"message":        "Please sign the payment context",
+			"paymentContext": createPaymentContext(),
+		})
+		return nil, false
+	}
+
+	timestampValue, ok := paymentTimestamp(c)
+	if !ok {
+		return nil, false
+	}
+
+	verifyResp, paymentCtx, err := verifyPayment(c.Request.Context(), signature, nonce, timestampValue)
+	if err != nil {
+		verificationTotal.WithLabelValues("error").Inc()
+
+		if errors.Is(err, context.DeadlineExceeded) {
+			respondError(c, http.StatusGatewayTimeout, "verifier_timeout", err)
+		} else {
+			respondError(c, http.StatusBadGateway, "verification_unavailable", err)
+		}
+		return nil, false
+	}
+
+	if !verifyResp.IsValid {
+		verificationTotal.WithLabelValues("invalid").Inc()
+		respondVerificationFailure(c, verifyResp)
+		return nil, false
+	}
+	if verifyResp.RecoveredAddress == "" {
+		verificationTotal.WithLabelValues("error").Inc()
+		respondError(c, http.StatusBadGateway, "verification_unavailable", fmt.Errorf("verifier success missing recovered_address"))
+		return nil, false
+	}
+
+	verificationTotal.WithLabelValues("success").Inc()
+
+	return &verifiedPayment{
+		PaymentContext:   *paymentCtx,
+		RecoveredAddress: verifyResp.RecoveredAddress,
+	}, true
+}
+
+func paymentTimestamp(c *gin.Context) (uint64, bool) {
+	timestampHeader := c.GetHeader("X-402-Timestamp")
+	if timestampHeader == "" {
+		respondError(c, http.StatusBadRequest, "invalid_timestamp", fmt.Errorf("missing X-402-Timestamp header"))
+		return 0, false
+	}
+
+	timestampValue, err := strconv.ParseUint(timestampHeader, 10, 64)
+	if err != nil || timestampValue == 0 {
+		respondError(c, http.StatusBadRequest, "invalid_timestamp", fmt.Errorf("invalid X-402-Timestamp header"))
+		return 0, false
+	}
+
+	return timestampValue, true
+}
+
+func sendPaidResult(c *gin.Context, payment *verifiedPayment, requestBody []byte, result string) error {
+	if payment == nil {
+		err := fmt.Errorf("missing verified payment")
+		respondError(c, http.StatusInternalServerError, "receipt_generation_failed", err)
+		return err
+	}
+
+	return generateAndSendReceipt(c, payment.PaymentContext, payment.RecoveredAddress, requestBody, result)
+}

gateway/payment_flow_test.go

@@ -0,0 +1,106 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+)
+
+func newPaymentFlowTestContext(req *http.Request) (*gin.Context, *httptest.ResponseRecorder) {
+	gin.SetMode(gin.TestMode)
+	recorder := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(recorder)
+	c.Request = req
+	if correlationID := req.Header.Get("X-Correlation-ID"); correlationID != "" {
+		c.Set("correlation_id", correlationID)
+		c.Request = c.Request.WithContext(context.WithValue(c.Request.Context(), CorrelationIDKey, correlationID))
+	}
+	return c, recorder
+}
+
+func TestVerifyPaidRequestWritesPaymentChallengeForMissingHeaders(t *testing.T) {
+	req := httptest.NewRequest(http.MethodPost, "/api/ai/summarize", strings.NewReader(`{"text":"hello"}`))
+	c, recorder := newPaymentFlowTestContext(req)
+
+	verified, ok := verifyPaidRequest(c)
+
+	require.False(t, ok)
+	require.Nil(t, verified)
+	require.Equal(t, http.StatusPaymentRequired, recorder.Code)
+
+	var response struct {
+		Error          string         `json:"error"`
+		Message        string         `json:"message"`
+		PaymentContext PaymentContext `json:"paymentContext"`
+	}
+	require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &response))
+	require.Equal(t, "Payment Required", response.Error)
+	require.Equal(t, "Please sign the payment context", response.Message)
+	require.NotEmpty(t, response.PaymentContext.Recipient)
+	require.NotEmpty(t, response.PaymentContext.Token)
+	require.NotEmpty(t, response.PaymentContext.Amount)
+	require.NotEmpty(t, response.PaymentContext.Nonce)
+	require.Positive(t, response.PaymentContext.ChainID)
+	require.Positive(t, response.PaymentContext.Timestamp)
+}
+
+func TestVerifyPaidRequestReturnsVerifiedPayment(t *testing.T) {
+	withVerifierResponse(t, http.StatusOK, `{"is_valid":true,"recovered_address":"0xabc","error":""}`)
+	req := signedSummarizeRequest(`{"text":"hello"}`)
+	c, recorder := newPaymentFlowTestContext(req)
+
+	verified, ok := verifyPaidRequest(c)
+
+	require.True(t, ok)
+	require.Equal(t, http.StatusOK, recorder.Code)
+	require.Equal(t, "0xabc", verified.RecoveredAddress)
+	require.Equal(t, "nonce-1", verified.PaymentContext.Nonce)
+	require.Equal(t, uint64(1700000000), verified.PaymentContext.Timestamp)
+}
+
+func TestVerifyPaidRequestMapsVerifierTimeout(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(1500 * time.Millisecond)
+	}))
+	t.Cleanup(server.Close)
+	t.Setenv("VERIFIER_URL", server.URL)
+	t.Setenv("VERIFIER_TIMEOUT_SECONDS", "1")
+
+	req := signedSummarizeRequest(`{"text":"hello"}`)
+	c, recorder := newPaymentFlowTestContext(req)
+
+	verified, ok := verifyPaidRequest(c)
+
+	require.False(t, ok)
+	require.Nil(t, verified)
+	require.Equal(t, http.StatusGatewayTimeout, recorder.Code)
+
+	var response map[string]string
+	require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &response))
+	require.Equal(t, "verifier_timeout", response["error"])
+	require.Equal(t, "test-correlation-id", response["correlation_id"])
+}
+
+func TestVerifyPaidRequestRequiresRecoveredAddress(t *testing.T) {
+	withVerifierResponse(t, http.StatusOK, `{"is_valid":true,"recovered_address":"","error":""}`)
+	req := signedSummarizeRequest(`{"text":"hello"}`)
+	c, recorder := newPaymentFlowTestContext(req)
+
+	verified, ok := verifyPaidRequest(c)
+
+	require.False(t, ok)
+	require.Nil(t, verified)
+	require.Equal(t, http.StatusBadGateway, recorder.Code)
+
+	var response map[string]string
+	require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &response))
+	require.Equal(t, "verification_unavailable", response["error"])
+	require.Equal(t, "test-correlation-id", response["correlation_id"])
+}