From 8306ef593c14c5a4b113257f336378867179add4 Mon Sep 17 00:00:00 2001 From: pisum-sativum Date: Sun, 5 Jul 2026 04:31:54 +0530 Subject: [PATCH 1/4] feat: stream AI response tokens via SSE This prevents holding up the HTTP handler goroutine for the duration of inference. Co-authored-by: codex --- gateway/cache.go | 12 +-- gateway/cache_integration_test.go | 35 +++++-- gateway/errors_test.go | 5 + gateway/internal/ai/ollama.go | 55 +++++++++++ gateway/internal/ai/openrouter.go | 86 +++++++++++++++++ gateway/internal/ai/provider.go | 10 ++ gateway/main.go | 112 ++++++++++++++++------ gateway/receipt_store_integration_test.go | 21 +++- gateway/timeout_test.go | 2 +- tests/e2e.test.ts | 5 +- web/src/hooks/use-x402.ts | 8 +- web/src/lib/x402-client.ts | 40 +++++++- 12 files changed, 333 insertions(+), 58 deletions(-) diff --git a/gateway/cache.go b/gateway/cache.go index 031d1723..3f43209d 100644 --- a/gateway/cache.go +++ b/gateway/cache.go @@ -174,9 +174,9 @@ func CacheMiddleware() gin.HandlerFunc { // 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 := streamResultAndReceipt(c, *paymentCtx, verifyResp.RecoveredAddress, requestBody, cached.Result); err != nil { log.Printf("Failed to send cached response receipt: %v", err) - // generateAndSendReceipt already sent an error response (500) + // streamResultAndReceipt already sent an error response } c.Abort() return @@ -208,10 +208,10 @@ func CacheMiddleware() gin.HandlerFunc { writer.mu.RUnlock() if statusCode == 200 { - // Response format: {"result": "...", "receipt": ...} - var resp map[string]interface{} - if err := json.Unmarshal(bodyBytes, &resp); err == nil { - if result, ok := resp["result"].(string); ok { + // Instead of parsing the SSE stream from bodyBytes, we retrieve the full summary + // which handleSummarize saves in the gin context. + if fullSummaryVal, exists := c.Get("full_summary"); exists { + if result, ok := fullSummaryVal.(string); ok { // Store asynchronously with a deadline to prevent indefinite goroutines go func(k, v string) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) diff --git a/gateway/cache_integration_test.go b/gateway/cache_integration_test.go index a0bf11a4..b0e75829 100644 --- a/gateway/cache_integration_test.go +++ b/gateway/cache_integration_test.go @@ -198,18 +198,33 @@ func TestCacheIntegration_FullFlow(t *testing.T) { } // Verify Body - var resp1, resp2 map[string]interface{} - if err := json.Unmarshal(w1.Body.Bytes(), &resp1); err != nil { - t.Fatalf("Failed to unmarshal response 1: %v", err) - } - if err := json.Unmarshal(w2.Body.Bytes(), &resp2); err != nil { - t.Fatalf("Failed to unmarshal response 2: %v", err) + parseSSEResultForTest := func(body []byte) string { + lines := strings.Split(string(body), "\n") + var result string + for _, line := range lines { + if strings.HasPrefix(line, "data: ") { + data := strings.TrimSpace(strings.TrimPrefix(line, "data: ")) + if data == "[DONE]" || data == "" { + continue + } + var chunk map[string]interface{} + if err := json.Unmarshal([]byte(data), &chunk); err == nil { + if text, ok := chunk["text"].(string); ok { + result += text + } + } + } + } + return result } - if val, ok := resp1["result"].(string); !ok || val != "AI Summary Result" { - t.Errorf("Unexpected result 1: %v", resp1["result"]) + result1 := parseSSEResultForTest(w1.Body.Bytes()) + result2 := parseSSEResultForTest(w2.Body.Bytes()) + + if result1 != "AI Summary Result" { + t.Errorf("Unexpected result 1: %v", result1) } - if val, ok := resp2["result"].(string); !ok || val != "AI Summary Result" { - t.Errorf("Unexpected result 2: %v", resp2["result"]) + if result2 != "AI Summary Result" { + t.Errorf("Unexpected result 2: %v", result2) } } diff --git a/gateway/errors_test.go b/gateway/errors_test.go index d1f17039..85afff5e 100644 --- a/gateway/errors_test.go +++ b/gateway/errors_test.go @@ -10,6 +10,7 @@ import ( "testing" "time" + "gateway/internal/ai" "github.com/alicebob/miniredis/v2" "github.com/gin-gonic/gin" "github.com/redis/go-redis/v9" @@ -24,6 +25,10 @@ func (p failingProvider) Generate(context.Context, string) (string, error) { return "", p.err } +func (p failingProvider) GenerateStream(context.Context, string) (ai.Stream, error) { + return nil, p.err +} + func newSummarizeTestRouter() *gin.Engine { gin.SetMode(gin.TestMode) r := gin.New() diff --git a/gateway/internal/ai/ollama.go b/gateway/internal/ai/ollama.go index 57847cba..92b837fc 100644 --- a/gateway/internal/ai/ollama.go +++ b/gateway/internal/ai/ollama.go @@ -77,3 +77,58 @@ func (p *OllamaProvider) Generate(ctx context.Context, text string) (string, err return response, nil } + +type ollamaStream struct { + resp *http.Response + decoder *json.Decoder +} + +func (s *ollamaStream) Recv() (string, error) { + var chunk map[string]interface{} + if err := s.decoder.Decode(&chunk); err != nil { + return "", err // will return io.EOF when stream ends + } + if response, ok := chunk["response"].(string); ok { + return response, nil + } + return "", fmt.Errorf("invalid response from Ollama: missing response field") +} + +func (s *ollamaStream) Close() error { + return s.resp.Body.Close() +} + +func (p *OllamaProvider) GenerateStream(ctx context.Context, text string) (Stream, error) { + prompt := fmt.Sprintf("Summarize this text in 2 sentences: %s", text) + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": p.model, + "prompt": prompt, + "stream": true, + }) + + req, err := http.NewRequestWithContext(ctx, "POST", p.url+"/api/generate", bytes.NewBuffer(reqBody)) + if err != nil { + return nil, fmt.Errorf("failed to create Ollama request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) || ctx.Err() == context.DeadlineExceeded { + return nil, context.DeadlineExceeded + } + return nil, fmt.Errorf("failed to connect to Ollama: %w", err) + } + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + return nil, fmt.Errorf("ollama returned status %d: %s", resp.StatusCode, string(body)) + } + + return &ollamaStream{ + resp: resp, + decoder: json.NewDecoder(resp.Body), + }, nil +} diff --git a/gateway/internal/ai/openrouter.go b/gateway/internal/ai/openrouter.go index 36bc48ab..2418008a 100644 --- a/gateway/internal/ai/openrouter.go +++ b/gateway/internal/ai/openrouter.go @@ -1,6 +1,7 @@ package ai import ( + "bufio" "bytes" "context" "encoding/json" @@ -100,3 +101,88 @@ func (p *OpenRouterProvider) Generate(ctx context.Context, text string) (string, return content, nil } + +type openRouterStream struct { + resp *http.Response + reader *bufio.Reader +} + +func (s *openRouterStream) Recv() (string, error) { + for { + line, err := s.reader.ReadBytes('\n') + if err != nil { + return "", err + } + line = bytes.TrimSpace(line) + if len(line) == 0 { + continue + } + if bytes.HasPrefix(line, []byte("data: ")) { + data := bytes.TrimPrefix(line, []byte("data: ")) + if string(data) == "[DONE]" { + return "", io.EOF + } + var chunk map[string]interface{} + if err := json.Unmarshal(data, &chunk); err != nil { + continue // skip invalid json + } + choices, ok := chunk["choices"].([]interface{}) + if !ok || len(choices) == 0 { + continue + } + choice, ok := choices[0].(map[string]interface{}) + if !ok { + continue + } + delta, ok := choice["delta"].(map[string]interface{}) + if !ok { + continue + } + if content, ok := delta["content"].(string); ok { + return content, nil + } + } + } +} + +func (s *openRouterStream) Close() error { + return s.resp.Body.Close() +} + +func (p *OpenRouterProvider) GenerateStream(ctx context.Context, text string) (Stream, error) { + prompt := fmt.Sprintf("Summarize this text in 2 sentences: %s", text) + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": p.model, + "messages": []map[string]string{ + {"role": "user", "content": prompt}, + }, + "stream": true, + }) + + req, err := http.NewRequestWithContext(ctx, "POST", p.url, bytes.NewBuffer(reqBody)) + if err != nil { + return nil, fmt.Errorf("failed to create OpenRouter request: %w", err) + } + req.Header.Set("Authorization", "Bearer "+p.apiKey) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) || ctx.Err() == context.DeadlineExceeded { + return nil, context.DeadlineExceeded + } + return nil, err + } + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + return nil, fmt.Errorf("openrouter returned status %d: %s", resp.StatusCode, string(body)) + } + + return &openRouterStream{ + resp: resp, + reader: bufio.NewReader(resp.Body), + }, nil +} diff --git a/gateway/internal/ai/provider.go b/gateway/internal/ai/provider.go index ec027973..cc0314d6 100644 --- a/gateway/internal/ai/provider.go +++ b/gateway/internal/ai/provider.go @@ -6,10 +6,20 @@ import ( "os" ) +// Stream defines the interface for streaming AI responses +type Stream interface { + // Recv returns the next token or an error (io.EOF when done) + Recv() (string, error) + // Close releases any resources associated with the stream + Close() error +} + // Provider defines the interface for AI service providers type Provider interface { // Generate takes a context and prompt text, returns the AI-generated response Generate(ctx context.Context, prompt string) (string, error) + // GenerateStream takes a context and prompt text, returns a Stream of generated tokens + GenerateStream(ctx context.Context, prompt string) (Stream, error) } // NewProvider creates an AI provider based on the AI_PROVIDER environment variable diff --git a/gateway/main.go b/gateway/main.go index 953c94d1..1c817a62 100644 --- a/gateway/main.go +++ b/gateway/main.go @@ -521,7 +521,7 @@ func handleSummarize(c *gin.Context) { } // 3. Call AI Service - summary, err := aiProvider.Generate(c.Request.Context(), req.Text) + stream, err := aiProvider.GenerateStream(c.Request.Context(), req.Text) if err != nil { if errors.Is(err, context.DeadlineExceeded) || errors.Is(c.Request.Context().Err(), context.DeadlineExceeded) { respondError(c, 504, "upstream_timeout", err) @@ -530,15 +530,65 @@ func handleSummarize(c *gin.Context) { } return } + defer stream.Close() - // 4. Generate & Send Receipt - if err := generateAndSendReceipt(c, *paymentCtx, verifyResp.RecoveredAddress, 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. + c.Writer.Header().Set("Content-Type", "text/event-stream") + c.Writer.Header().Set("Cache-Control", "no-cache") + c.Writer.Header().Set("Connection", "keep-alive") + + flusher, ok := c.Writer.(http.Flusher) + if !ok { + respondError(c, 500, "streaming_unsupported", fmt.Errorf("streaming unsupported")) + return + } + + var fullSummary string + for { + text, err := stream.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + fmt.Fprintf(c.Writer, "data: {\"error\": %q}\n\n", err.Error()) + flusher.Flush() + return + } + + fullSummary += text + chunkBytes, _ := json.Marshal(map[string]string{"text": text}) + fmt.Fprintf(c.Writer, "data: %s\n\n", chunkBytes) + flusher.Flush() + } + + // 4. Generate Receipt + responseMap := map[string]interface{}{ + "result": fullSummary, + } + responseBody, _ := json.Marshal(responseMap) + + receipt, err := GenerateReceipt(*paymentCtx, verifyResp.RecoveredAddress, c.Request.URL.Path, requestBody, responseBody) + if err != nil { + fmt.Fprintf(c.Writer, "data: {\"error\": \"receipt_generation_failed\"}\n\n") + flusher.Flush() + return + } + + if err := storeReceiptWithContext(c.Request.Context(), receipt, getReceiptTTL()); err != nil { + fmt.Fprintf(c.Writer, "data: {\"error\": \"receipt_store_failed\"}\n\n") + flusher.Flush() return } + + receiptJSON, _ := json.Marshal(receipt) + receiptBase64 := base64.StdEncoding.EncodeToString(receiptJSON) + + // Send receipt as final event + fmt.Fprintf(c.Writer, "data: {\"receipt\": %q}\n\n", receiptBase64) + fmt.Fprintf(c.Writer, "data: [DONE]\n\n") + flusher.Flush() + + // Store for cache middleware + c.Set("full_summary", fullSummary) } // verifyPayment calls the verification service. @@ -609,42 +659,48 @@ func verifyPayment(ctx context.Context, signature, nonce string, timestamp uint6 return &verifyResp, &paymentCtx, nil } -// generateAndSendReceipt handles receipt generation, storage, and sending the final JSON response. -// The receipt is sent ONLY in the X-402-Receipt header, not in the response body, -// to ensure the ResponseHash in the receipt matches the actual JSON body clients receive. -func generateAndSendReceipt(c *gin.Context, paymentCtx PaymentContext, recoveredAddr string, requestBody []byte, aiResult string) error { - // Construct the response body that will be sent to client (without receipt) +// streamResultAndReceipt is used by CacheMiddleware for cache hits to stream the response identically to a live request. +func streamResultAndReceipt(c *gin.Context, paymentCtx PaymentContext, recoveredAddr string, requestBody []byte, aiResult string) error { + c.Writer.Header().Set("Content-Type", "text/event-stream") + c.Writer.Header().Set("Cache-Control", "no-cache") + c.Writer.Header().Set("Connection", "keep-alive") + + flusher, ok := c.Writer.(http.Flusher) + if !ok { + respondError(c, 500, "streaming_unsupported", fmt.Errorf("streaming unsupported")) + return fmt.Errorf("streaming unsupported") + } + + chunkBytes, _ := json.Marshal(map[string]string{"text": aiResult}) + fmt.Fprintf(c.Writer, "data: %s\n\n", chunkBytes) + flusher.Flush() + + // Generate Receipt responseMap := map[string]interface{}{ "result": aiResult, } - responseBody, err := json.Marshal(responseMap) - if err != nil { - respondError(c, 500, "response_encoding_failed", err) - return err - } + responseBody, _ := json.Marshal(responseMap) - // Generate receipt with the actual response body hash receipt, err := GenerateReceipt(paymentCtx, recoveredAddr, c.Request.URL.Path, requestBody, responseBody) if err != nil { - respondError(c, 500, "receipt_generation_failed", err) + fmt.Fprintf(c.Writer, "data: {\"error\": \"receipt_generation_failed\"}\n\n") + flusher.Flush() return err } if err := storeReceiptWithContext(c.Request.Context(), receipt, getReceiptTTL()); err != nil { - respondError(c, 500, "receipt_store_failed", err) + fmt.Fprintf(c.Writer, "data: {\"error\": \"receipt_store_failed\"}\n\n") + flusher.Flush() return err } - receiptJSON, err := json.Marshal(receipt) - if err != nil { - respondError(c, 500, "receipt_encoding_failed", err) - return err - } + receiptJSON, _ := json.Marshal(receipt) receiptBase64 := base64.StdEncoding.EncodeToString(receiptJSON) - // Send receipt in header only (not in body) so ResponseHash matches body - c.Header("X-402-Receipt", receiptBase64) - c.JSON(200, responseMap) + // Send receipt as final event + fmt.Fprintf(c.Writer, "data: {\"receipt\": %q}\n\n", receiptBase64) + fmt.Fprintf(c.Writer, "data: [DONE]\n\n") + flusher.Flush() return nil } diff --git a/gateway/receipt_store_integration_test.go b/gateway/receipt_store_integration_test.go index c7a2e400..2175f0d8 100644 --- a/gateway/receipt_store_integration_test.go +++ b/gateway/receipt_store_integration_test.go @@ -83,11 +83,24 @@ func TestRedisReceiptStore_PersistsAcrossGatewayRestart(t *testing.T) { t.Fatalf("create receipt status=%d body=%s", createResp.Code, createResp.Body.String()) } - receiptHeader := createResp.Header().Get("X-402-Receipt") - if receiptHeader == "" { - t.Fatal("missing X-402-Receipt header") + var receiptBase64 string + lines := strings.Split(createResp.Body.String(), "\n") + for _, line := range lines { + if strings.HasPrefix(line, "data: ") { + dataStr := strings.TrimSpace(strings.TrimPrefix(line, "data: ")) + var chunk map[string]interface{} + if err := json.Unmarshal([]byte(dataStr), &chunk); err == nil { + if r, ok := chunk["receipt"].(string); ok { + receiptBase64 = r + } + } + } + } + + if receiptBase64 == "" { + t.Fatal("missing receipt event in SSE stream") } - receiptJSON, err := base64.StdEncoding.DecodeString(receiptHeader) + receiptJSON, err := base64.StdEncoding.DecodeString(receiptBase64) if err != nil { t.Fatalf("decode receipt header: %v", err) } diff --git a/gateway/timeout_test.go b/gateway/timeout_test.go index b0e2e03e..7581c3f0 100644 --- a/gateway/timeout_test.go +++ b/gateway/timeout_test.go @@ -63,7 +63,7 @@ func TestCallOpenRouter_RespectsContextTimeout(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) defer cancel() - _, err = provider.Generate(ctx, "hello") + _, err = provider.GenerateStream(ctx, "hello") if err == nil { t.Fatalf("Expected timeout error from provider, got nil") } diff --git a/tests/e2e.test.ts b/tests/e2e.test.ts index 3f5ce3b7..902b7156 100644 --- a/tests/e2e.test.ts +++ b/tests/e2e.test.ts @@ -86,8 +86,9 @@ describe("MicroAI Paygate E2E Flow", () => { } expect(res.status).toBe(200); - const data = await res.json() as any; - expect(data.result).toBeDefined(); + const textStr = await res.text(); + expect(textStr).toContain("data: {"); + expect(textStr).toContain("[DONE]"); }, 30000); it("should reject replayed signed payment context", async () => { diff --git a/web/src/hooks/use-x402.ts b/web/src/hooks/use-x402.ts index 7f5933fe..dea0eb33 100644 --- a/web/src/hooks/use-x402.ts +++ b/web/src/hooks/use-x402.ts @@ -111,7 +111,9 @@ export function useX402() { if (first.status === 200) { update({ step: "receipt" }); - const { summary, receipt } = await readSummarizeSuccess(first); + const { summary, receipt } = await readSummarizeSuccess(first, (text) => { + update({ summary: text }); + }); if (receipt) saveReceipt(receipt, text); track(AnalyticsEvent.SummaryCompleted, { ...flowProps, @@ -280,7 +282,9 @@ export function useX402() { } update({ step: "receipt" }); - const { summary, receipt } = await readSummarizeSuccess(retry); + const { summary, receipt } = await readSummarizeSuccess(retry, (text) => { + update({ summary: text }); + }); if (receipt) saveReceipt(receipt, text); track(AnalyticsEvent.SummaryCompleted, { ...flowProps, diff --git a/web/src/lib/x402-client.ts b/web/src/lib/x402-client.ts index 2a24c42a..70cfca06 100644 --- a/web/src/lib/x402-client.ts +++ b/web/src/lib/x402-client.ts @@ -68,12 +68,42 @@ export type SummarizeSuccess = { receipt: SignedReceipt | null; }; -export async function readSummarizeSuccess(res: Response): Promise { - const data = (await res.json()) as { result?: string }; - const summary = data.result ?? ""; +export async function readSummarizeSuccess( + res: Response, + onChunk?: (text: string) => void +): Promise { + if (!res.body) throw new Error("No response body"); + const reader = res.body.getReader(); + const decoder = new TextDecoder("utf-8"); + let summary = ""; + let receipt: SignedReceipt | null = null; + let buffer = ""; - const headerVal = res.headers.get("x-402-receipt") ?? res.headers.get("X-402-Receipt"); - const receipt = headerVal ? safeDecodeReceiptHeader(headerVal) : null; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + buffer += decoder.decode(value, { stream: true }); + + const lines = buffer.split("\n"); + buffer = lines.pop() ?? ""; + for (const line of lines) { + if (line.startsWith("data: ")) { + const dataStr = line.slice("data: ".length).trim(); + if (dataStr === "[DONE]") continue; + try { + const parsed = JSON.parse(dataStr); + if (parsed.text) { + summary += parsed.text; + if (onChunk) onChunk(summary); + } else if (parsed.receipt) { + receipt = safeDecodeReceiptHeader(parsed.receipt); + } + } catch (e) { + // ignore + } + } + } + } return { summary, receipt }; } From 924fd70cd55fa3d74710cfaa8ebb83bacaa4f858 Mon Sep 17 00:00:00 2001 From: pisum-sativum Date: Sun, 5 Jul 2026 04:43:09 +0530 Subject: [PATCH 2/4] feat: fix transaction hash replay attack Fixes issue #241 where the same payment transaction hash could be replayed indefinitely. This adds an atomic check and set in Redis to record used transaction hashes (EIP-712 signatures) immediately after first successful verification. A TTL of 30 days is applied to the keys in Redis to match the on-chain finality + dispute window. Additionally, it implements an in-memory fallback to avoid breaking tests and adds a load test simulating concurrent replay submissions. Co-authored-by: codex --- gateway/cache.go | 19 ++++ gateway/main.go | 16 ++++ gateway/receipt_store_integration_test.go | 106 ++++++++++++++++++++++ gateway/redis.go | 27 ++++++ 4 files changed, 168 insertions(+) diff --git a/gateway/cache.go b/gateway/cache.go index 3f43209d..f3ecec3f 100644 --- a/gateway/cache.go +++ b/gateway/cache.go @@ -164,6 +164,25 @@ func CacheMiddleware() gin.HandlerFunc { c.Abort() return } + + // === Fix for #241: Prevent replay attacks on cache hits === + used, err := markTransactionUsed(c.Request.Context(), signature) + if err != nil { + respondError(c, 500, "internal_error", fmt.Errorf("failed to check transaction replay: %w", err)) + c.Abort() + return + } + if used { + c.JSON(http.StatusPaymentRequired, gin.H{ + "error": "Payment Required", + "message": "Transaction already used", + "paymentContext": createPaymentContext(), + }) + c.Abort() + return + } + // ========================================================== + verificationTotal.WithLabelValues("success").Inc() // Payment Verified. Store verification for downstream if needed (though we abort) c.Set("payment_verification", verifyResp) diff --git a/gateway/main.go b/gateway/main.go index 1c817a62..6886ad54 100644 --- a/gateway/main.go +++ b/gateway/main.go @@ -505,6 +505,22 @@ func handleSummarize(c *gin.Context) { return } + // === Fix for #241: Prevent replay attacks === + used, err := markTransactionUsed(c.Request.Context(), signature) + if err != nil { + respondError(c, 500, "internal_error", fmt.Errorf("failed to check transaction replay: %w", err)) + return + } + if used { + c.JSON(http.StatusPaymentRequired, gin.H{ + "error": "Payment Required", + "message": "Transaction already used", + "paymentContext": createPaymentContext(), + }) + return + } + // ============================================ + verificationTotal.WithLabelValues("success").Inc() // 2. Parse Request diff --git a/gateway/receipt_store_integration_test.go b/gateway/receipt_store_integration_test.go index 2175f0d8..4c79bdeb 100644 --- a/gateway/receipt_store_integration_test.go +++ b/gateway/receipt_store_integration_test.go @@ -194,3 +194,109 @@ func replaceReceiptGlobalsForTest(t *testing.T) func() { receiptGlobalsTestMu.Unlock() } } + +func TestHandleSummarize_ConcurrentReplayAttack(t *testing.T) { + ctx := t.Context() + redisServer := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: redisServer.Addr()}) + defer rdb.Close() + + verifier := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Mock verifier that always approves + _ = json.NewEncoder(w).Encode(VerifyResponse{ + IsValid: true, + RecoveredAddress: "0x742d35Cc6634C0532925a3b844Bc9e7595f8fE21", + }) + })) + defer verifier.Close() + + aiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"AI response"}}]}`)) + })) + defer aiServer.Close() + + t.Setenv("CACHE_ENABLED", "false") + t.Setenv("RECEIPT_STORE", "redis") + t.Setenv("REDIS_URL", redisServer.Addr()) + t.Setenv("AI_PROVIDER", "openrouter") + t.Setenv("OPENROUTER_URL", aiServer.URL) + t.Setenv("OPENROUTER_API_KEY", "test-key") + t.Setenv("VERIFIER_URL", verifier.URL) + t.Setenv("SERVER_WALLET_PRIVATE_KEY", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef") + t.Setenv("RECIPIENT_ADDRESS", "0x2cAF48b4BA1C58721a85dFADa5aC01C2DFa62219") + t.Setenv("RECEIPT_TTL", "86400") + + resetServerPrivateKeyForTest(t) + restoreReceiptGlobals := replaceReceiptGlobalsForTest(t) + defer restoreReceiptGlobals() + + if err := initRedis(); err != nil { + t.Fatalf("init redis: %v", err) + } + if err := initReceiptStore(); err != nil { + t.Fatalf("init redis receipt store: %v", err) + } + + var err error + aiProvider, err = ai.NewProvider() + if err != nil { + t.Fatalf("new AI provider: %v", err) + } + + gateway := newReceiptPersistenceTestRouter() + + const numConcurrentRequests = 20 + var wg sync.WaitGroup + wg.Add(numConcurrentRequests) + + successCount := 0 + conflictCount := 0 + var mu sync.Mutex + + // Create requests beforehand so they start at the exact same time + var requests []*http.Request + for i := 0; i < numConcurrentRequests; i++ { + req := httptest.NewRequest(http.MethodPost, "/api/ai/summarize", bytes.NewBufferString(`{"text":"summarize me"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-402-Signature", "0xSameSignatureReplayTest") // Same signature for all requests + req.Header.Set("X-402-Nonce", "replay-test-nonce") + req.Header.Set("X-402-Timestamp", strconv.FormatInt(time.Now().Unix(), 10)) + requests = append(requests, req) + } + + start := make(chan struct{}) + + for i := 0; i < numConcurrentRequests; i++ { + go func(req *http.Request) { + defer wg.Done() + <-start // wait for signal to start simultaneously + + resp := httptest.NewRecorder() + gateway.ServeHTTP(resp, req) + + mu.Lock() + if resp.Code == http.StatusOK { + successCount++ + } else if resp.Code == http.StatusPaymentRequired { + var body map[string]interface{} + _ = json.Unmarshal(resp.Body.Bytes(), &body) + if body["message"] == "Transaction already used" { + conflictCount++ + } + } + mu.Unlock() + }(requests[i]) + } + + close(start) // Signal all goroutines to start + wg.Wait() + + if successCount != 1 { + t.Errorf("Expected exactly 1 successful request, got %d", successCount) + } + if conflictCount != numConcurrentRequests - 1 { + t.Errorf("Expected exactly %d blocked replays, got %d", numConcurrentRequests - 1, conflictCount) + } +} + diff --git a/gateway/redis.go b/gateway/redis.go index 53d7a5b2..c7eb96fc 100644 --- a/gateway/redis.go +++ b/gateway/redis.go @@ -6,6 +6,7 @@ import ( "log" "os" "strings" + "sync" "time" "github.com/redis/go-redis/v9" @@ -76,3 +77,29 @@ func getEnv(key, fallback string) string { } return fallback } + +var ( + memoryUsedTx = make(map[string]bool) + memoryUsedTxMu sync.Mutex +) + +// markTransactionUsed checks and records a transaction signature as used atomically. +// It returns true if the transaction was already used, or false if it was successfully marked. +func markTransactionUsed(ctx context.Context, txHash string) (bool, error) { + if redisClient == nil { + memoryUsedTxMu.Lock() + defer memoryUsedTxMu.Unlock() + if memoryUsedTx[txHash] { + return true, nil + } + memoryUsedTx[txHash] = true + return false, nil + } + key := "used_tx:" + txHash + // SetNX returns true if the key was set, meaning it wasn't used before. + added, err := redisClient.SetNX(ctx, key, "1", 30*24*time.Hour).Result() + if err != nil { + return false, err + } + return !added, nil +} From f713c3da059a6688486f44a69c61f5ac2fde6b7f Mon Sep 17 00:00:00 2001 From: pisum-sativum Date: Sun, 5 Jul 2026 22:05:10 +0530 Subject: [PATCH 3/4] fix: resolve build and vet errors in gateway - Remove unused bodyBytes in cache.go - Add missing strings package imports - Remove unused ctx in receipt_store_integration_test.go Co-authored-by: codex --- gateway/cache.go | 1 - gateway/cache_integration_test.go | 1 + gateway/receipt_store_integration_test.go | 2 +- 3 files changed, 2 insertions(+), 2 deletions(-) diff --git a/gateway/cache.go b/gateway/cache.go index f3ecec3f..55b00e91 100644 --- a/gateway/cache.go +++ b/gateway/cache.go @@ -223,7 +223,6 @@ func CacheMiddleware() gin.HandlerFunc { // Handler finished. Check status and extract result with proper locking writer.mu.RLock() statusCode := writer.ResponseWriter.Status() - bodyBytes := writer.body.Bytes() writer.mu.RUnlock() if statusCode == 200 { diff --git a/gateway/cache_integration_test.go b/gateway/cache_integration_test.go index b0e75829..373d8cb8 100644 --- a/gateway/cache_integration_test.go +++ b/gateway/cache_integration_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "strings" "gateway/internal/ai" "net/http" "net/http/httptest" diff --git a/gateway/receipt_store_integration_test.go b/gateway/receipt_store_integration_test.go index 4c79bdeb..b86920d6 100644 --- a/gateway/receipt_store_integration_test.go +++ b/gateway/receipt_store_integration_test.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/base64" "encoding/json" + "strings" "gateway/internal/ai" "net/http" "net/http/httptest" @@ -196,7 +197,6 @@ func replaceReceiptGlobalsForTest(t *testing.T) func() { } func TestHandleSummarize_ConcurrentReplayAttack(t *testing.T) { - ctx := t.Context() redisServer := miniredis.RunT(t) rdb := redis.NewClient(&redis.Options{Addr: redisServer.Addr()}) defer rdb.Close() From b875dcbc682ca7582b1c4912f04f952e092562e1 Mon Sep 17 00:00:00 2001 From: pisum-sativum Date: Sun, 5 Jul 2026 22:20:48 +0530 Subject: [PATCH 4/4] fix: address P1/P2 review findings on fix-tx-replay-attack-241 - gateway/internal/ai/openrouter.go: add ErrIncompleteStream sentinel and track sawDone so bare EOF (dropped TCP) is distinguished from a clean [DONE] termination; prevents partial/empty summaries getting receipts - gateway/redis.go: replace no-TTL bool map with expiry-time map in memoryUsedTx; entries now age out after 30 days matching Redis path; add on-call pruning to prevent unbounded growth on standalone instances; normalize signatures to lowercase before keying (hex-casing bypass fix) - gateway/main.go: canonicalize X-402-Signature to lowercase before markTransactionUsed; change gateway-side replay response from 402 Payment Required to 409 Conflict / nonce_already_used matching the verifier and web client contract; treat ErrIncompleteStream as upstream error; guard against empty-summary receipt issuance - gateway/cache.go: same signature canonicalization and 409 replay fix in CacheMiddleware; add missing strings import - gateway/cache_integration_test.go: use 0xValidSig2 for cache-hit request so replay guard does not block it; update mock verifier to accept both signatures - gateway/receipt_store_integration_test.go: update concurrent-replay test to assert 409 Conflict / nonce_already_used instead of 402 - sdk/typescript/src/protocol/microai.ts: add extractReceiptFromSseBody and readSseSuccessBody helpers that parse the SSE event stream; keep readReceipt as a no-op stub for interface compatibility - sdk/typescript/src/client.ts: replace JSON-body readSuccess with SSE- aware implementation using readSseSuccessBody; fix response_hash binding to hash canonical JSON.stringify({result}) matching gateway GenerateReceipt; propagate stream error events as PaygateSdkError instead of silently ignoring them Co-authored-by: codex --- gateway/cache.go | 11 ++-- gateway/cache_integration_test.go | 9 ++- gateway/internal/ai/openrouter.go | 17 ++++- gateway/main.go | 27 ++++++-- gateway/receipt_store_integration_test.go | 4 +- gateway/redis.go | 27 ++++++-- sdk/typescript/src/client.ts | 31 +++++---- sdk/typescript/src/protocol/microai.ts | 79 ++++++++++++++++++++++- 8 files changed, 163 insertions(+), 42 deletions(-) diff --git a/gateway/cache.go b/gateway/cache.go index 55b00e91..de5f34cd 100644 --- a/gateway/cache.go +++ b/gateway/cache.go @@ -13,6 +13,7 @@ import ( "net/http" "os" "strconv" + "strings" "sync" "time" @@ -166,17 +167,17 @@ func CacheMiddleware() gin.HandlerFunc { } // === Fix for #241: Prevent replay attacks on cache hits === - used, err := markTransactionUsed(c.Request.Context(), signature) + // Normalize to lowercase to prevent hex-casing bypass attacks. + used, err := markTransactionUsed(c.Request.Context(), strings.ToLower(signature)) if err != nil { respondError(c, 500, "internal_error", fmt.Errorf("failed to check transaction replay: %w", err)) c.Abort() return } if used { - c.JSON(http.StatusPaymentRequired, gin.H{ - "error": "Payment Required", - "message": "Transaction already used", - "paymentContext": createPaymentContext(), + c.JSON(http.StatusConflict, gin.H{ + "error": "nonce_already_used", + "message": "Transaction already used", }) c.Abort() return diff --git a/gateway/cache_integration_test.go b/gateway/cache_integration_test.go index 373d8cb8..2a85c29c 100644 --- a/gateway/cache_integration_test.go +++ b/gateway/cache_integration_test.go @@ -37,7 +37,9 @@ func TestCacheIntegration_FullFlow(t *testing.T) { return } - isValid := req.Signature == "0xValidSig" + // Accept two distinct valid signatures so the cache-hit request + // uses a fresh signature that isn't blocked by the replay guard. + isValid := req.Signature == "0xValidSig" || req.Signature == "0xValidSig2" resp := VerifyResponse{ IsValid: isValid, RecoveredAddress: "0xTestUser", @@ -158,9 +160,10 @@ func TestCacheIntegration_FullFlow(t *testing.T) { } assertCachePopulated() - // Request 2: Cache Hit (Valid Sig) + // Request 2: Cache Hit — must use a DIFFERENT valid signature so the + // gateway-side replay guard (markTransactionUsed) doesn't block it. start = time.Now() - w2 := makeRequest("0xValidSig") + w2 := makeRequest("0xValidSig2") duration2 := time.Since(start) if w2.Code != 200 { diff --git a/gateway/internal/ai/openrouter.go b/gateway/internal/ai/openrouter.go index 2418008a..43271129 100644 --- a/gateway/internal/ai/openrouter.go +++ b/gateway/internal/ai/openrouter.go @@ -102,15 +102,27 @@ func (p *OpenRouterProvider) Generate(ctx context.Context, text string) (string, return content, nil } +// ErrIncompleteStream is returned by Recv when the upstream TCP connection +// closes before the [DONE] sentinel is received. This distinguishes a clean +// completion from a dropped or truncated stream. +var ErrIncompleteStream = errors.New("upstream stream ended without [DONE]") + type openRouterStream struct { - resp *http.Response - reader *bufio.Reader + resp *http.Response + reader *bufio.Reader + sawDone bool } func (s *openRouterStream) Recv() (string, error) { for { line, err := s.reader.ReadBytes('\n') if err != nil { + if errors.Is(err, io.EOF) { + if s.sawDone { + return "", io.EOF + } + return "", ErrIncompleteStream + } return "", err } line = bytes.TrimSpace(line) @@ -120,6 +132,7 @@ func (s *openRouterStream) Recv() (string, error) { if bytes.HasPrefix(line, []byte("data: ")) { data := bytes.TrimPrefix(line, []byte("data: ")) if string(data) == "[DONE]" { + s.sawDone = true return "", io.EOF } var chunk map[string]interface{} diff --git a/gateway/main.go b/gateway/main.go index 6886ad54..e7c7a93f 100644 --- a/gateway/main.go +++ b/gateway/main.go @@ -506,16 +506,20 @@ func handleSummarize(c *gin.Context) { } // === Fix for #241: Prevent replay attacks === - used, err := markTransactionUsed(c.Request.Context(), signature) + // Normalize signature to lowercase before keying so the same 65-byte + // value with different hex casing maps to the same replay entry. + used, err := markTransactionUsed(c.Request.Context(), strings.ToLower(signature)) if err != nil { respondError(c, 500, "internal_error", fmt.Errorf("failed to check transaction replay: %w", err)) return } if used { - c.JSON(http.StatusPaymentRequired, gin.H{ - "error": "Payment Required", - "message": "Transaction already used", - "paymentContext": createPaymentContext(), + // Return the established replay-rejection contract (409 Conflict / + // nonce_already_used) that the web client and verifier both use, + // rather than a fresh 402 payment challenge. + c.JSON(http.StatusConflict, gin.H{ + "error": "nonce_already_used", + "message": "Transaction already used", }) return } @@ -565,17 +569,28 @@ func handleSummarize(c *gin.Context) { break } if err != nil { + // ErrIncompleteStream means the upstream TCP connection dropped + // before [DONE] arrived — treat as an upstream error so a partial + // or empty summary never gets a signed receipt. fmt.Fprintf(c.Writer, "data: {\"error\": %q}\n\n", err.Error()) flusher.Flush() return } - + fullSummary += text chunkBytes, _ := json.Marshal(map[string]string{"text": text}) fmt.Fprintf(c.Writer, "data: %s\n\n", chunkBytes) flusher.Flush() } + // Guard: if the stream was empty (upstream returned nothing), do not + // issue a receipt for a zero-content response. + if fullSummary == "" { + fmt.Fprintf(c.Writer, "data: {\"error\": \"upstream_empty_response\"}\n\n") + flusher.Flush() + return + } + // 4. Generate Receipt responseMap := map[string]interface{}{ "result": fullSummary, diff --git a/gateway/receipt_store_integration_test.go b/gateway/receipt_store_integration_test.go index b86920d6..13032910 100644 --- a/gateway/receipt_store_integration_test.go +++ b/gateway/receipt_store_integration_test.go @@ -278,10 +278,10 @@ func TestHandleSummarize_ConcurrentReplayAttack(t *testing.T) { mu.Lock() if resp.Code == http.StatusOK { successCount++ - } else if resp.Code == http.StatusPaymentRequired { + } else if resp.Code == http.StatusConflict { var body map[string]interface{} _ = json.Unmarshal(resp.Body.Bytes(), &body) - if body["message"] == "Transaction already used" { + if body["error"] == "nonce_already_used" { conflictCount++ } } diff --git a/gateway/redis.go b/gateway/redis.go index c7eb96fc..3ef5c8d8 100644 --- a/gateway/redis.go +++ b/gateway/redis.go @@ -78,26 +78,43 @@ func getEnv(key, fallback string) string { return fallback } +const memoryUsedTxTTL = 30 * 24 * time.Hour + var ( - memoryUsedTx = make(map[string]bool) + memoryUsedTx = make(map[string]time.Time) // signature -> expiry time memoryUsedTxMu sync.Mutex ) // markTransactionUsed checks and records a transaction signature as used atomically. // It returns true if the transaction was already used, or false if it was successfully marked. +// Signatures are normalized to lowercase before keying to prevent hex-casing bypass attacks. func markTransactionUsed(ctx context.Context, txHash string) (bool, error) { + // Normalize to lowercase so the same 65-byte signature in different + // hex casing maps to the same replay key, preventing casing bypass. + key := strings.ToLower(txHash) + if redisClient == nil { + now := time.Now() memoryUsedTxMu.Lock() defer memoryUsedTxMu.Unlock() - if memoryUsedTx[txHash] { + + // Prune expired entries to prevent unbounded map growth. + for k, exp := range memoryUsedTx { + if now.After(exp) { + delete(memoryUsedTx, k) + } + } + + if exp, exists := memoryUsedTx[key]; exists && now.Before(exp) { return true, nil } - memoryUsedTx[txHash] = true + memoryUsedTx[key] = now.Add(memoryUsedTxTTL) return false, nil } - key := "used_tx:" + txHash + + redisKey := "used_tx:" + key // SetNX returns true if the key was set, meaning it wasn't used before. - added, err := redisClient.SetNX(ctx, key, "1", 30*24*time.Hour).Result() + added, err := redisClient.SetNX(ctx, redisKey, "1", memoryUsedTxTTL).Result() if err != nil { return false, err } diff --git a/sdk/typescript/src/client.ts b/sdk/typescript/src/client.ts index 23765f28..3a7722ca 100644 --- a/sdk/typescript/src/client.ts +++ b/sdk/typescript/src/client.ts @@ -1,6 +1,6 @@ import { ethers } from "ethers"; import { PaygateSdkError } from "./errors"; -import { MicroAIPaygateProtocol } from "./protocol/microai"; +import { MicroAIPaygateProtocol, readSseSuccessBody } from "./protocol/microai"; import type { FetchLike, PaygateProtocolAdapter, @@ -153,12 +153,14 @@ export class PaygateClient { private receiptMatchesPayload( receipt: SignedReceipt, context: { endpoint: string; requestBodyText?: string }, - responseBodyText: string, + canonicalResponseBodyText: string, ): boolean { return ( receipt.receipt.service.endpoint === context.endpoint && receipt.receipt.service.request_hash === this.hashBody(context.requestBodyText) && - receipt.receipt.service.response_hash === this.hashBody(responseBodyText) + // The gateway hashes JSON.stringify({"result": fullSummary}) — we must + // hash the same canonical form, not the raw SSE event-stream bytes. + receipt.receipt.service.response_hash === this.hashBody(canonicalResponseBodyText) ); } @@ -166,18 +168,15 @@ export class PaygateClient { response: Response, context: { endpoint: string; requestBodyText?: string }, ): Promise> { - const bodyText = await response.text(); - let data: TData; - try { - data = JSON.parse(bodyText) as TData; - } catch (error) { - throw new PaygateSdkError("network_error", "Gateway returned invalid JSON", { - status: response.status, - bodyText, - cause: error, - }); - } - const receipt = this.protocol.readReceipt(response); + // The gateway streams the response as SSE events. Consume the stream, + // accumulate text chunks, extract the embedded receipt event, and derive + // the canonical response body string for receipt hash verification. + const { fullText, receipt, canonicalResponseBody } = await readSseSuccessBody(response); + + // Expose the result in the same shape the summarize() caller expects: + // { result: string } + const data = { result: fullText } as unknown as TData; + if (receipt === null) { return { data, @@ -187,7 +186,7 @@ export class PaygateClient { }; } - if (!this.receiptMatchesPayload(receipt, context, bodyText)) { + if (!this.receiptMatchesPayload(receipt, context, canonicalResponseBody)) { throw new PaygateSdkError( "receipt_verification_failed", "Gateway receipt does not match the request and response payload", diff --git a/sdk/typescript/src/protocol/microai.ts b/sdk/typescript/src/protocol/microai.ts index b2e2bd23..d2b8ff45 100644 --- a/sdk/typescript/src/protocol/microai.ts +++ b/sdk/typescript/src/protocol/microai.ts @@ -11,6 +11,8 @@ import type { export const MICROAI_SIGNATURE_HEADER = "X-402-Signature"; export const MICROAI_NONCE_HEADER = "X-402-Nonce"; export const MICROAI_TIMESTAMP_HEADER = "X-402-Timestamp"; +// Kept for backwards compatibility; the gateway now embeds the receipt in the +// SSE stream body as `data: {"receipt": ""}` rather than this header. export const MICROAI_RECEIPT_HEADER = "X-402-Receipt"; function isRecord(value: unknown): value is Record { @@ -33,6 +35,75 @@ function isPaymentContext(value: unknown): value is PaymentContext { ); } +/** + * Extracts the base64-encoded receipt from an SSE stream body string. + * The gateway emits: `data: {"receipt": ""}\n\n` + * Returns null if no receipt event is found. + */ +export function extractReceiptFromSseBody(bodyText: string): SignedReceipt | null { + const lines = bodyText.split("\n"); + for (const line of lines) { + if (!line.startsWith("data: ")) continue; + const dataStr = line.slice("data: ".length).trim(); + if (dataStr === "[DONE]" || dataStr === "") continue; + try { + const parsed: unknown = JSON.parse(dataStr); + if (isRecord(parsed) && typeof parsed.receipt === "string") { + return decodeReceiptHeader(parsed.receipt); + } + } catch { + // skip malformed lines + } + } + return null; +} + +/** + * Consumes an SSE stream response, accumulating text chunks and extracting + * the embedded receipt event. Returns the full summary text, the receipt, + * and the canonical response body JSON string used for receipt hash binding. + */ +export async function readSseSuccessBody(response: Response): Promise<{ + fullText: string; + receipt: SignedReceipt | null; + canonicalResponseBody: string; +}> { + const bodyText = await response.text(); + let fullText = ""; + let receipt: SignedReceipt | null = null; + + const lines = bodyText.split("\n"); + for (const line of lines) { + if (!line.startsWith("data: ")) continue; + const dataStr = line.slice("data: ".length).trim(); + if (dataStr === "[DONE]" || dataStr === "") continue; + try { + const parsed: unknown = JSON.parse(dataStr); + if (!isRecord(parsed)) continue; + if (typeof parsed.text === "string") { + fullText += parsed.text; + } else if (typeof parsed.receipt === "string") { + receipt = decodeReceiptHeader(parsed.receipt); + } else if (parsed.error !== undefined) { + throw new PaygateSdkError( + "network_error", + `Gateway stream error: ${String(parsed.error)}`, + {}, + ); + } + } catch (e) { + if (e instanceof PaygateSdkError) throw e; + // ignore unparseable lines + } + } + + // Build the canonical response body that the gateway hashes for the receipt. + // GenerateReceipt in the gateway uses: json.Marshal({"result": fullSummary}) + const canonicalResponseBody = JSON.stringify({ result: fullText }); + + return { fullText, receipt, canonicalResponseBody }; +} + export class MicroAIPaygateProtocol implements PaygateProtocolAdapter { async readPaymentContext(response: Response): Promise { const bodyText = await response.text(); @@ -67,8 +138,10 @@ export class MicroAIPaygateProtocol implements PaygateProtocolAdapter { return buildSignedHeaders(ctx, signature); } - readReceipt(response: Response): SignedReceipt | null { - const header = response.headers.get(MICROAI_RECEIPT_HEADER); - return header ? decodeReceiptHeader(header) : null; + // readReceipt is no longer used for SSE responses; receipt extraction is + // handled by readSseSuccessBody in client.ts. This method is retained for + // protocol-adapter interface compatibility. + readReceipt(_response: Response): SignedReceipt | null { + return null; } }