nezha_domains/integration/agentcompat/internal/client/websocket_test.go

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")
	}
}