fix(claude): prevent compressed SSE streams and add magic-byte decompression fallback
- Set Accept-Encoding: identity for SSE streams; upstream must not compress line-delimited SSE bodies that bufio.Scanner reads directly - Re-enforce identity after ApplyCustomHeadersFromAttrs to prevent auth attribute injection from re-enabling compression on the stream path - Add peekableBody type wrapping bufio.Reader for non-consuming magic-byte inspection of the first 4 bytes without affecting downstream readers - Detect gzip (0x1f 0x8b) and zstd (0x28 0xb5 0x2f 0xfd) by magic bytes when Content-Encoding header is absent, covering misbehaving upstreams - Remove if-Content-Encoding guard on all three error paths (Execute, ExecuteStream, CountTokens); unconditionally delegate to decodeResponseBody so magic-byte detection applies consistently to all response paths - Add 10 tests covering stream identity enforcement, compressed success bodies, magic-byte detection without headers, error path decoding, and auth attribute override prevention
This commit is contained in:
@@ -187,17 +187,15 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
|
|||||||
}
|
}
|
||||||
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||||
// Decompress error responses (e.g. gzip-compressed 400 errors from Anthropic API).
|
// Decompress error responses — pass the Content-Encoding value (may be empty)
|
||||||
errBody := httpResp.Body
|
// and let decodeResponseBody handle both header-declared and magic-byte-detected
|
||||||
if ce := httpResp.Header.Get("Content-Encoding"); ce != "" {
|
// compression. This keeps error-path behaviour consistent with the success path.
|
||||||
var decErr error
|
errBody, decErr := decodeResponseBody(httpResp.Body, httpResp.Header.Get("Content-Encoding"))
|
||||||
errBody, decErr = decodeResponseBody(httpResp.Body, ce)
|
if decErr != nil {
|
||||||
if decErr != nil {
|
recordAPIResponseError(ctx, e.cfg, decErr)
|
||||||
recordAPIResponseError(ctx, e.cfg, decErr)
|
msg := fmt.Sprintf("failed to decode error response body: %v", decErr)
|
||||||
msg := fmt.Sprintf("failed to decode error response body (encoding=%s): %v", ce, decErr)
|
logWithRequestID(ctx).Warn(msg)
|
||||||
logWithRequestID(ctx).Warn(msg)
|
return resp, statusErr{code: httpResp.StatusCode, msg: msg}
|
||||||
return resp, statusErr{code: httpResp.StatusCode, msg: msg}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
b, readErr := io.ReadAll(errBody)
|
b, readErr := io.ReadAll(errBody)
|
||||||
if readErr != nil {
|
if readErr != nil {
|
||||||
@@ -352,17 +350,15 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
|
|||||||
}
|
}
|
||||||
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||||
// Decompress error responses (e.g. gzip-compressed 400 errors from Anthropic API).
|
// Decompress error responses — pass the Content-Encoding value (may be empty)
|
||||||
errBody := httpResp.Body
|
// and let decodeResponseBody handle both header-declared and magic-byte-detected
|
||||||
if ce := httpResp.Header.Get("Content-Encoding"); ce != "" {
|
// compression. This keeps error-path behaviour consistent with the success path.
|
||||||
var decErr error
|
errBody, decErr := decodeResponseBody(httpResp.Body, httpResp.Header.Get("Content-Encoding"))
|
||||||
errBody, decErr = decodeResponseBody(httpResp.Body, ce)
|
if decErr != nil {
|
||||||
if decErr != nil {
|
recordAPIResponseError(ctx, e.cfg, decErr)
|
||||||
recordAPIResponseError(ctx, e.cfg, decErr)
|
msg := fmt.Sprintf("failed to decode error response body: %v", decErr)
|
||||||
msg := fmt.Sprintf("failed to decode error response body (encoding=%s): %v", ce, decErr)
|
logWithRequestID(ctx).Warn(msg)
|
||||||
logWithRequestID(ctx).Warn(msg)
|
return nil, statusErr{code: httpResp.StatusCode, msg: msg}
|
||||||
return nil, statusErr{code: httpResp.StatusCode, msg: msg}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
b, readErr := io.ReadAll(errBody)
|
b, readErr := io.ReadAll(errBody)
|
||||||
if readErr != nil {
|
if readErr != nil {
|
||||||
@@ -521,17 +517,15 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut
|
|||||||
}
|
}
|
||||||
recordAPIResponseMetadata(ctx, e.cfg, resp.StatusCode, resp.Header.Clone())
|
recordAPIResponseMetadata(ctx, e.cfg, resp.StatusCode, resp.Header.Clone())
|
||||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
// Decompress error responses (e.g. gzip-compressed 400 errors from Anthropic API).
|
// Decompress error responses — pass the Content-Encoding value (may be empty)
|
||||||
errBody := resp.Body
|
// and let decodeResponseBody handle both header-declared and magic-byte-detected
|
||||||
if ce := resp.Header.Get("Content-Encoding"); ce != "" {
|
// compression. This keeps error-path behaviour consistent with the success path.
|
||||||
var decErr error
|
errBody, decErr := decodeResponseBody(resp.Body, resp.Header.Get("Content-Encoding"))
|
||||||
errBody, decErr = decodeResponseBody(resp.Body, ce)
|
if decErr != nil {
|
||||||
if decErr != nil {
|
recordAPIResponseError(ctx, e.cfg, decErr)
|
||||||
recordAPIResponseError(ctx, e.cfg, decErr)
|
msg := fmt.Sprintf("failed to decode error response body: %v", decErr)
|
||||||
msg := fmt.Sprintf("failed to decode error response body (encoding=%s): %v", ce, decErr)
|
logWithRequestID(ctx).Warn(msg)
|
||||||
logWithRequestID(ctx).Warn(msg)
|
return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: msg}
|
||||||
return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: msg}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
b, readErr := io.ReadAll(errBody)
|
b, readErr := io.ReadAll(errBody)
|
||||||
if readErr != nil {
|
if readErr != nil {
|
||||||
@@ -662,12 +656,61 @@ func (c *compositeReadCloser) Close() error {
|
|||||||
return firstErr
|
return firstErr
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// peekableBody wraps a bufio.Reader around the original ReadCloser so that
|
||||||
|
// magic bytes can be inspected without consuming them from the stream.
|
||||||
|
type peekableBody struct {
|
||||||
|
*bufio.Reader
|
||||||
|
closer io.Closer
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *peekableBody) Close() error {
|
||||||
|
return p.closer.Close()
|
||||||
|
}
|
||||||
|
|
||||||
func decodeResponseBody(body io.ReadCloser, contentEncoding string) (io.ReadCloser, error) {
|
func decodeResponseBody(body io.ReadCloser, contentEncoding string) (io.ReadCloser, error) {
|
||||||
if body == nil {
|
if body == nil {
|
||||||
return nil, fmt.Errorf("response body is nil")
|
return nil, fmt.Errorf("response body is nil")
|
||||||
}
|
}
|
||||||
if contentEncoding == "" {
|
if contentEncoding == "" {
|
||||||
return body, nil
|
// No Content-Encoding header. Attempt best-effort magic-byte detection to
|
||||||
|
// handle misbehaving upstreams that compress without setting the header.
|
||||||
|
// Only gzip (1f 8b) and zstd (28 b5 2f fd) have reliable magic sequences;
|
||||||
|
// br and deflate have none and are left as-is.
|
||||||
|
// The bufio wrapper preserves unread bytes so callers always see the full
|
||||||
|
// stream regardless of whether decompression was applied.
|
||||||
|
pb := &peekableBody{Reader: bufio.NewReader(body), closer: body}
|
||||||
|
magic, peekErr := pb.Peek(4)
|
||||||
|
if peekErr == nil || (peekErr == io.EOF && len(magic) >= 2) {
|
||||||
|
switch {
|
||||||
|
case len(magic) >= 2 && magic[0] == 0x1f && magic[1] == 0x8b:
|
||||||
|
gzipReader, gzErr := gzip.NewReader(pb)
|
||||||
|
if gzErr != nil {
|
||||||
|
_ = pb.Close()
|
||||||
|
return nil, fmt.Errorf("magic-byte gzip: failed to create reader: %w", gzErr)
|
||||||
|
}
|
||||||
|
return &compositeReadCloser{
|
||||||
|
Reader: gzipReader,
|
||||||
|
closers: []func() error{
|
||||||
|
gzipReader.Close,
|
||||||
|
pb.Close,
|
||||||
|
},
|
||||||
|
}, nil
|
||||||
|
case len(magic) >= 4 && magic[0] == 0x28 && magic[1] == 0xb5 && magic[2] == 0x2f && magic[3] == 0xfd:
|
||||||
|
decoder, zdErr := zstd.NewReader(pb)
|
||||||
|
if zdErr != nil {
|
||||||
|
_ = pb.Close()
|
||||||
|
return nil, fmt.Errorf("magic-byte zstd: failed to create reader: %w", zdErr)
|
||||||
|
}
|
||||||
|
return &compositeReadCloser{
|
||||||
|
Reader: decoder,
|
||||||
|
closers: []func() error{
|
||||||
|
func() error { decoder.Close(); return nil },
|
||||||
|
pb.Close,
|
||||||
|
},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return pb, nil
|
||||||
}
|
}
|
||||||
encodings := strings.Split(contentEncoding, ",")
|
encodings := strings.Split(contentEncoding, ",")
|
||||||
for _, raw := range encodings {
|
for _, raw := range encodings {
|
||||||
@@ -844,11 +887,15 @@ func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string,
|
|||||||
r.Header.Set("User-Agent", hdrDefault(hd.UserAgent, "claude-cli/2.1.63 (external, cli)"))
|
r.Header.Set("User-Agent", hdrDefault(hd.UserAgent, "claude-cli/2.1.63 (external, cli)"))
|
||||||
}
|
}
|
||||||
r.Header.Set("Connection", "keep-alive")
|
r.Header.Set("Connection", "keep-alive")
|
||||||
r.Header.Set("Accept-Encoding", "gzip, deflate, br, zstd")
|
|
||||||
if stream {
|
if stream {
|
||||||
r.Header.Set("Accept", "text/event-stream")
|
r.Header.Set("Accept", "text/event-stream")
|
||||||
|
// SSE streams must not be compressed: the downstream scanner reads
|
||||||
|
// line-delimited text and cannot parse compressed bytes. Using
|
||||||
|
// "identity" tells the upstream to send an uncompressed stream.
|
||||||
|
r.Header.Set("Accept-Encoding", "identity")
|
||||||
} else {
|
} else {
|
||||||
r.Header.Set("Accept", "application/json")
|
r.Header.Set("Accept", "application/json")
|
||||||
|
r.Header.Set("Accept-Encoding", "gzip, deflate, br, zstd")
|
||||||
}
|
}
|
||||||
// Keep OS/Arch mapping dynamic (not configurable).
|
// Keep OS/Arch mapping dynamic (not configurable).
|
||||||
// They intentionally continue to derive from runtime.GOOS/runtime.GOARCH.
|
// They intentionally continue to derive from runtime.GOOS/runtime.GOARCH.
|
||||||
@@ -857,6 +904,12 @@ func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string,
|
|||||||
attrs = auth.Attributes
|
attrs = auth.Attributes
|
||||||
}
|
}
|
||||||
util.ApplyCustomHeadersFromAttrs(r, attrs)
|
util.ApplyCustomHeadersFromAttrs(r, attrs)
|
||||||
|
// Re-enforce Accept-Encoding: identity after ApplyCustomHeadersFromAttrs, which
|
||||||
|
// may override it with a user-configured value. Compressed SSE breaks the line
|
||||||
|
// scanner regardless of user preference, so this is non-negotiable for streams.
|
||||||
|
if stream {
|
||||||
|
r.Header.Set("Accept-Encoding", "identity")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func claudeCreds(a *cliproxyauth.Auth) (apiKey, baseURL string) {
|
func claudeCreds(a *cliproxyauth.Auth) (apiKey, baseURL string) {
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ package executor
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"bytes"
|
"bytes"
|
||||||
|
"compress/gzip"
|
||||||
"context"
|
"context"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
@@ -9,6 +10,7 @@ import (
|
|||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
|
"github.com/klauspost/compress/zstd"
|
||||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||||
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||||
@@ -583,3 +585,385 @@ func testClaudeExecutorInvalidCompressedErrorBody(
|
|||||||
t.Fatalf("expected status code 400, got: %v", err)
|
t.Fatalf("expected status code 400, got: %v", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestClaudeExecutor_ExecuteStream_SetsIdentityAcceptEncoding verifies that streaming
|
||||||
|
// requests use Accept-Encoding: identity so the upstream cannot respond with a
|
||||||
|
// compressed SSE body that would silently break the line scanner.
|
||||||
|
func TestClaudeExecutor_ExecuteStream_SetsIdentityAcceptEncoding(t *testing.T) {
|
||||||
|
var gotEncoding, gotAccept string
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
gotEncoding = r.Header.Get("Accept-Encoding")
|
||||||
|
gotAccept = r.Header.Get("Accept")
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
_, _ = w.Write([]byte("data: {\"type\":\"message_stop\"}\n\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
executor := NewClaudeExecutor(&config.Config{})
|
||||||
|
auth := &cliproxyauth.Auth{Attributes: map[string]string{
|
||||||
|
"api_key": "key-123",
|
||||||
|
"base_url": server.URL,
|
||||||
|
}}
|
||||||
|
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
|
||||||
|
|
||||||
|
result, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "claude-3-5-sonnet-20241022",
|
||||||
|
Payload: payload,
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FromString("claude"),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ExecuteStream error: %v", err)
|
||||||
|
}
|
||||||
|
for chunk := range result.Chunks {
|
||||||
|
if chunk.Err != nil {
|
||||||
|
t.Fatalf("unexpected chunk error: %v", chunk.Err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if gotEncoding != "identity" {
|
||||||
|
t.Errorf("Accept-Encoding = %q, want %q", gotEncoding, "identity")
|
||||||
|
}
|
||||||
|
if gotAccept != "text/event-stream" {
|
||||||
|
t.Errorf("Accept = %q, want %q", gotAccept, "text/event-stream")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestClaudeExecutor_Execute_SetsCompressedAcceptEncoding verifies that non-streaming
|
||||||
|
// requests keep the full accept-encoding to allow response compression (which
|
||||||
|
// decodeResponseBody handles correctly).
|
||||||
|
func TestClaudeExecutor_Execute_SetsCompressedAcceptEncoding(t *testing.T) {
|
||||||
|
var gotEncoding, gotAccept string
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
gotEncoding = r.Header.Get("Accept-Encoding")
|
||||||
|
gotAccept = r.Header.Get("Accept")
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = w.Write([]byte(`{"id":"msg_1","type":"message","model":"claude-3-5-sonnet-20241022","role":"assistant","content":[{"type":"text","text":"hi"}],"usage":{"input_tokens":1,"output_tokens":1}}`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
executor := NewClaudeExecutor(&config.Config{})
|
||||||
|
auth := &cliproxyauth.Auth{Attributes: map[string]string{
|
||||||
|
"api_key": "key-123",
|
||||||
|
"base_url": server.URL,
|
||||||
|
}}
|
||||||
|
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
|
||||||
|
|
||||||
|
_, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "claude-3-5-sonnet-20241022",
|
||||||
|
Payload: payload,
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FromString("claude"),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if gotEncoding != "gzip, deflate, br, zstd" {
|
||||||
|
t.Errorf("Accept-Encoding = %q, want %q", gotEncoding, "gzip, deflate, br, zstd")
|
||||||
|
}
|
||||||
|
if gotAccept != "application/json" {
|
||||||
|
t.Errorf("Accept = %q, want %q", gotAccept, "application/json")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestClaudeExecutor_ExecuteStream_GzipSuccessBodyDecoded verifies that a streaming
|
||||||
|
// HTTP 200 response with Content-Encoding: gzip is correctly decompressed before
|
||||||
|
// the line scanner runs, so SSE chunks are not silently dropped.
|
||||||
|
func TestClaudeExecutor_ExecuteStream_GzipSuccessBodyDecoded(t *testing.T) {
|
||||||
|
var buf bytes.Buffer
|
||||||
|
gz := gzip.NewWriter(&buf)
|
||||||
|
_, _ = gz.Write([]byte("data: {\"type\":\"message_stop\"}\n"))
|
||||||
|
_ = gz.Close()
|
||||||
|
compressedBody := buf.Bytes()
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Header().Set("Content-Encoding", "gzip")
|
||||||
|
_, _ = w.Write(compressedBody)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
executor := NewClaudeExecutor(&config.Config{})
|
||||||
|
auth := &cliproxyauth.Auth{Attributes: map[string]string{
|
||||||
|
"api_key": "key-123",
|
||||||
|
"base_url": server.URL,
|
||||||
|
}}
|
||||||
|
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
|
||||||
|
|
||||||
|
result, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "claude-3-5-sonnet-20241022",
|
||||||
|
Payload: payload,
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FromString("claude"),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ExecuteStream error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var combined strings.Builder
|
||||||
|
for chunk := range result.Chunks {
|
||||||
|
if chunk.Err != nil {
|
||||||
|
t.Fatalf("chunk error: %v", chunk.Err)
|
||||||
|
}
|
||||||
|
combined.Write(chunk.Payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
if combined.Len() == 0 {
|
||||||
|
t.Fatal("expected at least one chunk from gzip-encoded SSE body, got none (body was not decompressed)")
|
||||||
|
}
|
||||||
|
if !strings.Contains(combined.String(), "message_stop") {
|
||||||
|
t.Errorf("expected SSE content in chunks, got: %q", combined.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestDecodeResponseBody_MagicByteGzipNoHeader verifies that decodeResponseBody
|
||||||
|
// detects gzip-compressed content via magic bytes even when Content-Encoding is absent.
|
||||||
|
func TestDecodeResponseBody_MagicByteGzipNoHeader(t *testing.T) {
|
||||||
|
const plaintext = "data: {\"type\":\"message_stop\"}\n"
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
gz := gzip.NewWriter(&buf)
|
||||||
|
_, _ = gz.Write([]byte(plaintext))
|
||||||
|
_ = gz.Close()
|
||||||
|
|
||||||
|
rc := io.NopCloser(&buf)
|
||||||
|
decoded, err := decodeResponseBody(rc, "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("decodeResponseBody error: %v", err)
|
||||||
|
}
|
||||||
|
defer decoded.Close()
|
||||||
|
|
||||||
|
got, err := io.ReadAll(decoded)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ReadAll error: %v", err)
|
||||||
|
}
|
||||||
|
if string(got) != plaintext {
|
||||||
|
t.Errorf("decoded = %q, want %q", got, plaintext)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestDecodeResponseBody_PlainTextNoHeader verifies that decodeResponseBody returns
|
||||||
|
// plain text untouched when Content-Encoding is absent and no magic bytes match.
|
||||||
|
func TestDecodeResponseBody_PlainTextNoHeader(t *testing.T) {
|
||||||
|
const plaintext = "data: {\"type\":\"message_stop\"}\n"
|
||||||
|
rc := io.NopCloser(strings.NewReader(plaintext))
|
||||||
|
decoded, err := decodeResponseBody(rc, "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("decodeResponseBody error: %v", err)
|
||||||
|
}
|
||||||
|
defer decoded.Close()
|
||||||
|
|
||||||
|
got, err := io.ReadAll(decoded)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ReadAll error: %v", err)
|
||||||
|
}
|
||||||
|
if string(got) != plaintext {
|
||||||
|
t.Errorf("decoded = %q, want %q", got, plaintext)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestClaudeExecutor_ExecuteStream_GzipNoContentEncodingHeader verifies the full
|
||||||
|
// pipeline: when the upstream returns a gzip-compressed SSE body WITHOUT setting
|
||||||
|
// Content-Encoding (a misbehaving upstream), the magic-byte sniff in
|
||||||
|
// decodeResponseBody still decompresses it, so chunks reach the caller.
|
||||||
|
func TestClaudeExecutor_ExecuteStream_GzipNoContentEncodingHeader(t *testing.T) {
|
||||||
|
var buf bytes.Buffer
|
||||||
|
gz := gzip.NewWriter(&buf)
|
||||||
|
_, _ = gz.Write([]byte("data: {\"type\":\"message_stop\"}\n"))
|
||||||
|
_ = gz.Close()
|
||||||
|
compressedBody := buf.Bytes()
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
// Intentionally omit Content-Encoding to simulate misbehaving upstream.
|
||||||
|
_, _ = w.Write(compressedBody)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
executor := NewClaudeExecutor(&config.Config{})
|
||||||
|
auth := &cliproxyauth.Auth{Attributes: map[string]string{
|
||||||
|
"api_key": "key-123",
|
||||||
|
"base_url": server.URL,
|
||||||
|
}}
|
||||||
|
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
|
||||||
|
|
||||||
|
result, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "claude-3-5-sonnet-20241022",
|
||||||
|
Payload: payload,
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FromString("claude"),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ExecuteStream error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var combined strings.Builder
|
||||||
|
for chunk := range result.Chunks {
|
||||||
|
if chunk.Err != nil {
|
||||||
|
t.Fatalf("chunk error: %v", chunk.Err)
|
||||||
|
}
|
||||||
|
combined.Write(chunk.Payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
if combined.Len() == 0 {
|
||||||
|
t.Fatal("expected chunks from gzip body without Content-Encoding header, got none (magic-byte sniff failed)")
|
||||||
|
}
|
||||||
|
if !strings.Contains(combined.String(), "message_stop") {
|
||||||
|
t.Errorf("unexpected chunk content: %q", combined.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestClaudeExecutor_ExecuteStream_AcceptEncodingOverrideCannotBypassIdentity verifies
|
||||||
|
// that injecting Accept-Encoding via auth.Attributes cannot override the stream
|
||||||
|
// path's enforced identity encoding.
|
||||||
|
func TestClaudeExecutor_ExecuteStream_AcceptEncodingOverrideCannotBypassIdentity(t *testing.T) {
|
||||||
|
var gotEncoding string
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
gotEncoding = r.Header.Get("Accept-Encoding")
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
_, _ = w.Write([]byte("data: {\"type\":\"message_stop\"}\n\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
executor := NewClaudeExecutor(&config.Config{})
|
||||||
|
// Inject Accept-Encoding via the custom header attribute mechanism.
|
||||||
|
auth := &cliproxyauth.Auth{Attributes: map[string]string{
|
||||||
|
"api_key": "key-123",
|
||||||
|
"base_url": server.URL,
|
||||||
|
"header:Accept-Encoding": "gzip, deflate, br, zstd",
|
||||||
|
}}
|
||||||
|
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
|
||||||
|
|
||||||
|
result, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "claude-3-5-sonnet-20241022",
|
||||||
|
Payload: payload,
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FromString("claude"),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ExecuteStream error: %v", err)
|
||||||
|
}
|
||||||
|
for chunk := range result.Chunks {
|
||||||
|
if chunk.Err != nil {
|
||||||
|
t.Fatalf("unexpected chunk error: %v", chunk.Err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if gotEncoding != "identity" {
|
||||||
|
t.Errorf("Accept-Encoding = %q; stream path must enforce identity regardless of auth.Attributes override", gotEncoding)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestDecodeResponseBody_MagicByteZstdNoHeader verifies that decodeResponseBody
|
||||||
|
// detects zstd-compressed content via magic bytes (28 b5 2f fd) even when
|
||||||
|
// Content-Encoding is absent.
|
||||||
|
func TestDecodeResponseBody_MagicByteZstdNoHeader(t *testing.T) {
|
||||||
|
const plaintext = "data: {\"type\":\"message_stop\"}\n"
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
enc, err := zstd.NewWriter(&buf)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("zstd.NewWriter: %v", err)
|
||||||
|
}
|
||||||
|
_, _ = enc.Write([]byte(plaintext))
|
||||||
|
_ = enc.Close()
|
||||||
|
|
||||||
|
rc := io.NopCloser(&buf)
|
||||||
|
decoded, err := decodeResponseBody(rc, "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("decodeResponseBody error: %v", err)
|
||||||
|
}
|
||||||
|
defer decoded.Close()
|
||||||
|
|
||||||
|
got, err := io.ReadAll(decoded)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ReadAll error: %v", err)
|
||||||
|
}
|
||||||
|
if string(got) != plaintext {
|
||||||
|
t.Errorf("decoded = %q, want %q", got, plaintext)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestClaudeExecutor_Execute_GzipErrorBodyNoContentEncodingHeader verifies that the
|
||||||
|
// error path (4xx) correctly decompresses a gzip body even when the upstream omits
|
||||||
|
// the Content-Encoding header. This closes the gap left by PR #1771, which only
|
||||||
|
// fixed header-declared compression on the error path.
|
||||||
|
func TestClaudeExecutor_Execute_GzipErrorBodyNoContentEncodingHeader(t *testing.T) {
|
||||||
|
const errJSON = `{"type":"error","error":{"type":"invalid_request_error","message":"test error"}}`
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
gz := gzip.NewWriter(&buf)
|
||||||
|
_, _ = gz.Write([]byte(errJSON))
|
||||||
|
_ = gz.Close()
|
||||||
|
compressedBody := buf.Bytes()
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
// Intentionally omit Content-Encoding to simulate misbehaving upstream.
|
||||||
|
w.WriteHeader(http.StatusBadRequest)
|
||||||
|
_, _ = w.Write(compressedBody)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
executor := NewClaudeExecutor(&config.Config{})
|
||||||
|
auth := &cliproxyauth.Auth{Attributes: map[string]string{
|
||||||
|
"api_key": "key-123",
|
||||||
|
"base_url": server.URL,
|
||||||
|
}}
|
||||||
|
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
|
||||||
|
|
||||||
|
_, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "claude-3-5-sonnet-20241022",
|
||||||
|
Payload: payload,
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FromString("claude"),
|
||||||
|
})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected an error for 400 response, got nil")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "test error") {
|
||||||
|
t.Errorf("error message should contain decompressed JSON, got: %q", err.Error())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestClaudeExecutor_ExecuteStream_GzipErrorBodyNoContentEncodingHeader verifies
|
||||||
|
// the same for the streaming executor: 4xx gzip body without Content-Encoding is
|
||||||
|
// decoded and the error message is readable.
|
||||||
|
func TestClaudeExecutor_ExecuteStream_GzipErrorBodyNoContentEncodingHeader(t *testing.T) {
|
||||||
|
const errJSON = `{"type":"error","error":{"type":"invalid_request_error","message":"stream test error"}}`
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
gz := gzip.NewWriter(&buf)
|
||||||
|
_, _ = gz.Write([]byte(errJSON))
|
||||||
|
_ = gz.Close()
|
||||||
|
compressedBody := buf.Bytes()
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
// Intentionally omit Content-Encoding to simulate misbehaving upstream.
|
||||||
|
w.WriteHeader(http.StatusBadRequest)
|
||||||
|
_, _ = w.Write(compressedBody)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
executor := NewClaudeExecutor(&config.Config{})
|
||||||
|
auth := &cliproxyauth.Auth{Attributes: map[string]string{
|
||||||
|
"api_key": "key-123",
|
||||||
|
"base_url": server.URL,
|
||||||
|
}}
|
||||||
|
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
|
||||||
|
|
||||||
|
_, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "claude-3-5-sonnet-20241022",
|
||||||
|
Payload: payload,
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FromString("claude"),
|
||||||
|
})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected an error for 400 response, got nil")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "stream test error") {
|
||||||
|
t.Errorf("error message should contain decompressed JSON, got: %q", err.Error())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user