test(agentcompat): add protocol clients

Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
2026-09-19 17:50:12 +00:00 · 2026-07-20 04:48:15 +00:00
parent 8a0e84b214
commit 1c3e926819
23 changed files with 3329 additions and 0 deletions
@@ -0,0 +1,230 @@
+package client
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/require"
+)
+
+func TestClient_WebSocketFrameReassembly(t *testing.T) {
+	upgrader := websocket.Upgrader{WriteBufferSize: 4}
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		require.Equal(t, "Bearer ws-token", request.Header.Get("Authorization"))
+		require.NotEmpty(t, request.Header.Get("Origin"))
+		connection, err := upgrader.Upgrade(writer, request, nil)
+		require.NoError(t, err)
+		defer connection.Close()
+
+		textWriter, err := connection.NextWriter(websocket.TextMessage)
+		require.NoError(t, err)
+		_, err = textWriter.Write([]byte("hello "))
+		require.NoError(t, err)
+		_, err = textWriter.Write([]byte("world"))
+		require.NoError(t, err)
+		require.NoError(t, textWriter.Close())
+
+		binaryWriter, err := connection.NextWriter(websocket.BinaryMessage)
+		require.NoError(t, err)
+		_, err = binaryWriter.Write([]byte{1, 2})
+		require.NoError(t, err)
+		_, err = binaryWriter.Write([]byte{3, 4})
+		require.NoError(t, err)
+		require.NoError(t, binaryWriter.Close())
+	}))
+	t.Cleanup(server.Close)
+
+	client := newTestClient(t, Config{
+		BaseURL:          server.URL,
+		BearerToken:      "ws-token",
+		Origin:           server.URL,
+		RequestTimeout:   time.Second,
+		MaxResponseBytes: 1024,
+	})
+	connection, err := client.DialWebSocket(context.Background(), "/stream")
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, connection.Close()) })
+
+	textFrame, err := connection.ReadFrame(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, FrameText, textFrame.Type)
+	require.Equal(t, []byte("hello world"), textFrame.Payload)
+
+	binaryFrame, err := connection.ReadFrame(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, FrameBinary, binaryFrame.Type)
+	require.Equal(t, []byte{1, 2, 3, 4}, binaryFrame.Payload)
+}
+
+func TestClient_WebSocketRejectsOversize(t *testing.T) {
+	upgrader := websocket.Upgrader{}
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		connection, err := upgrader.Upgrade(writer, request, nil)
+		require.NoError(t, err)
+		defer connection.Close()
+		require.NoError(t, connection.WriteMessage(websocket.BinaryMessage, []byte(strings.Repeat("x", 65))))
+	}))
+	t.Cleanup(server.Close)
+
+	client := newTestClient(t, Config{
+		BaseURL:          server.URL,
+		RequestTimeout:   time.Second,
+		MaxResponseBytes: 64,
+	})
+	connection, err := client.DialWebSocket(context.Background(), "/oversize")
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, connection.Close()) })
+
+	_, err = connection.ReadFrame(context.Background())
+	require.ErrorIs(t, err, ErrResponseTooLarge)
+}
+
+func TestClient_DialWebSocket_ReturnsTypedHandshakeFailure(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		http.Error(writer, "permission denied", http.StatusForbidden)
+	}))
+	t.Cleanup(server.Close)
+	webSocketClient := newTestClient(t, Config{BaseURL: server.URL})
+
+	_, err := webSocketClient.DialWebSocket(context.Background(), "/terminal")
+
+	var handshakeError *WebSocketHandshakeError
+	require.ErrorAs(t, err, &handshakeError)
+	require.Equal(t, http.StatusForbidden, handshakeError.StatusCode)
+	require.Contains(t, handshakeError.Message, "permission denied")
+}
+
+func TestClient_WebSocketDeadline(t *testing.T) {
+	upgrader := websocket.Upgrader{}
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		connection, err := upgrader.Upgrade(writer, request, nil)
+		require.NoError(t, err)
+		defer connection.Close()
+		_, _, _ = connection.ReadMessage()
+	}))
+	t.Cleanup(server.Close)
+
+	client := newTestClient(t, Config{
+		BaseURL:          server.URL,
+		RequestTimeout:   time.Second,
+		MaxResponseBytes: 1024,
+	})
+	connection, err := client.DialWebSocket(context.Background(), "/deadline")
+	require.NoError(t, err)
+	connection.timeout = 20 * time.Millisecond
+
+	_, err = connection.ReadFrame(context.Background())
+	require.True(t, errors.Is(err, context.DeadlineExceeded), err)
+	require.NoError(t, connection.Close())
+}
+
+func TestClient_RedactsAuthorization(t *testing.T) {
+	secret := "nzp_super-secret"
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
+		writer.Header().Set("Content-Type", "application/json")
+		_, _ = writer.Write([]byte(`{"success":false,"error":"Authorization: Bearer ` + secret + `"}`))
+	}))
+	t.Cleanup(server.Close)
+
+	client := newTestClient(t, Config{
+		BaseURL:          server.URL,
+		BearerToken:      secret,
+		Origin:           server.URL,
+		RequestTimeout:   time.Second,
+		MaxResponseBytes: 1024,
+	})
+	_, err := REST[struct{}, semanticResult](context.Background(), client, RESTRequest[struct{}]{Method: http.MethodGet, Path: "/redaction"})
+	require.Error(t, err)
+	require.NotContains(t, err.Error(), secret)
+	require.Contains(t, err.Error(), "[REDACTED]")
+
+	redacted := Redact("request failed with Authorization: Bearer " + secret)
+	require.False(t, strings.Contains(redacted, secret))
+	require.Contains(t, redacted, "Authorization: Bearer [REDACTED]")
+
+	quoted := Redact(`{"Authorization":"Bearer ` + secret + `"}`)
+	require.NotContains(t, quoted, secret)
+	require.Contains(t, quoted, "[REDACTED]")
+
+	query := Redact("https://dashboard.example/mcp/download/path-token?access_token=" + secret + "&X-Amz-Signature=signature-secret")
+	require.NotContains(t, query, secret)
+	require.NotContains(t, query, "signature-secret")
+
+	httpError := &HTTPError{StatusCode: http.StatusBadRequest, Message: Redact("Authorization: Bearer " + secret)}
+	require.NotContains(t, httpError.Message, secret)
+	rpcError := &RPCError{Code: -32603, Message: Redact("token=" + secret)}
+	require.NotContains(t, rpcError.Message, secret)
+}
+
+func TestClient_RedactsCredentialClasses(t *testing.T) {
+	secret := "sensitive-value"
+	inputs := []string{
+		"X-CSRF-Token: " + secret,
+		"password=" + secret,
+		"https://dashboard.example/path?access_token=" + secret,
+		"jwt_token: " + secret,
+	}
+	for _, input := range inputs {
+		redacted := Redact(input)
+		require.NotContains(t, redacted, secret)
+		require.Contains(t, redacted, "[REDACTED]")
+	}
+}
+
+func TestClient_WebSocketReadDeadline(t *testing.T) {
+	upgrader := websocket.Upgrader{}
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		connection, err := upgrader.Upgrade(writer, request, nil)
+		require.NoError(t, err)
+		defer connection.Close()
+		<-request.Context().Done()
+	}))
+	t.Cleanup(server.Close)
+
+	client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: 20 * time.Millisecond, MaxResponseBytes: 1024})
+	connection, err := client.DialWebSocket(context.Background(), "/stream")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = connection.Close() })
+
+	_, err = connection.ReadFrame(context.Background())
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+}
+
+func TestClient_WebSocketWriteHonorsParentCancellation(t *testing.T) {
+	upgrader := websocket.Upgrader{}
+	serverReady := make(chan struct{})
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		connection, err := upgrader.Upgrade(writer, request, nil)
+		require.NoError(t, err)
+		defer connection.Close()
+		close(serverReady)
+		<-request.Context().Done()
+	}))
+	t.Cleanup(server.Close)
+
+	client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: 5 * time.Second, MaxResponseBytes: 1024})
+	connection, err := client.DialWebSocket(context.Background(), "/blocked-write")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = connection.Close() })
+
+	writeContext, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	result := make(chan error, 1)
+	<-serverReady
+	go func() {
+		result <- connection.WriteFrame(writeContext, Frame{Type: FrameBinary, Payload: make([]byte, 128<<20)})
+	}()
+	cancel()
+	select {
+	case err = <-result:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("WebSocket write did not stop after parent cancellation")
+	}
+}