From 6f7c12b06206577efcc1cc5446d1ba828e245277 Mon Sep 17 00:00:00 2001 From: AnkanMisra Date: Thu, 18 Jun 2026 17:04:34 +0530 Subject: [PATCH] refactor: centralize gateway paid request flow Co-authored-by: codex --- gateway/cache.go | 47 +--------------- gateway/main.go | 58 +------------------ gateway/payment_flow.go | 91 ++++++++++++++++++++++++++++++ gateway/payment_flow_test.go | 106 +++++++++++++++++++++++++++++++++++ 4 files changed, 203 insertions(+), 99 deletions(-) create mode 100644 gateway/payment_flow.go create mode 100644 gateway/payment_flow_test.go diff --git a/gateway/cache.go b/gateway/cache.go index 031d1723..aeb44ed3 100644 --- a/gateway/cache.go +++ b/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) } diff --git a/gateway/main.go b/gateway/main.go index 953c94d1..9de84d06 100644 --- a/gateway/main.go +++ b/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 } } diff --git a/gateway/payment_flow.go b/gateway/payment_flow.go new file mode 100644 index 00000000..47f48560 --- /dev/null +++ b/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) +} diff --git a/gateway/payment_flow_test.go b/gateway/payment_flow_test.go new file mode 100644 index 00000000..ffc79d23 --- /dev/null +++ b/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"]) +}