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

package client

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestIOStreamCapabilityClientUsesTypedPATAuthenticatedWireContract(t *testing.T) {
	rawCapability := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("c", 32)))
	requestNumber := 0
	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		requestNumber++
		require.Equal(t, "Bearer private-pat", request.Header.Get("Authorization"))
		writer.Header().Set("Content-Type", "application/json")
		switch requestNumber {
		case 1:
			require.Equal(t, "/agentcompat/io-stream-capability/register", request.URL.Path)
			var body IOStreamCapabilityRegisterRequest
			require.NoError(t, json.NewDecoder(request.Body).Decode(&body))
			require.Equal(t, IOStreamCapabilityPurposeTerminal, body.Purpose)
			require.Equal(t, uint64(7), body.ServerID)
			require.Zero(t, body.ResourceID)
			_, err := writer.Write([]byte(`{"success":true,"data":{"capability":"` + rawCapability + `"}}`))
			require.NoError(t, err)
		case 2:
			require.Equal(t, "/agentcompat/io-stream-capability/wait", request.URL.Path)
			var body IOStreamCapabilityWaitRequest
			require.NoError(t, json.NewDecoder(request.Body).Decode(&body))
			require.Equal(t, rawCapability, body.Capability.Value())
			_, err := writer.Write([]byte(`{"success":true,"data":{"stream_id":"private-stream"}}`))
			require.NoError(t, err)
		case 3, 4:
			expectedPath := "/agentcompat/io-stream-capability/cancel"
			if requestNumber == 4 {
				expectedPath = "/agentcompat/io-stream-capability/unregister"
			}
			require.Equal(t, expectedPath, request.URL.Path)
			_, err := writer.Write([]byte(`{"success":true,"data":{}}`))
			require.NoError(t, err)
		default:
			t.Fatalf("unexpected request %d", requestNumber)
		}
	}))
	t.Cleanup(server.Close)
	transport := newTestClient(t, Config{BaseURL: server.URL, BearerToken: "private-pat"})
	capabilities := transport.IOStreamCapabilities()

	registered, err := capabilities.Register(context.Background(), IOStreamCapabilityRegisterRequest{Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7})
	require.NoError(t, err)
	require.Equal(t, rawCapability, registered.Capability.Value())
	waited, err := capabilities.Wait(context.Background(), IOStreamCapabilityWaitRequest{Capability: registered.Capability, Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7})
	require.NoError(t, err)
	require.Equal(t, "private-stream", waited.StreamID.Value())
	access := IOStreamCapabilityAccessRequest{Capability: registered.Capability, Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7}
	require.NoError(t, capabilities.Cancel(context.Background(), access))
	require.NoError(t, capabilities.Unregister(context.Background(), access))
	require.Equal(t, 4, requestNumber)
}

func TestIOStreamCapabilityClientValidatesBeforeDispatch(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		t.Fatal("invalid request must not dispatch")
	}))
	t.Cleanup(server.Close)
	transport := newTestClient(t, Config{BaseURL: server.URL})
	capabilities := transport.IOStreamCapabilities()

	_, err := capabilities.Register(context.Background(), IOStreamCapabilityRegisterRequest{Purpose: "unknown", ServerID: 7})
	require.ErrorIs(t, err, ErrInvalidIOStreamCapabilityRequest)
	_, err = capabilities.Register(context.Background(), IOStreamCapabilityRegisterRequest{Purpose: IOStreamCapabilityPurposeNAT, ServerID: 7})
	require.ErrorIs(t, err, ErrInvalidIOStreamCapabilityRequest)
	_, err = capabilities.Wait(context.Background(), IOStreamCapabilityWaitRequest{Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7})
	require.ErrorIs(t, err, ErrInvalidIOStreamCapabilityRequest)
	require.Zero(t, transport.RequestCount())
}

func TestIOStreamCapabilityClientErrorsNeverEchoSensitiveValues(t *testing.T) {
	rawCapability := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("s", 32)))
	capability, err := ParseIOStreamCapability(rawCapability)
	require.NoError(t, err)
	privateValues := []string{rawCapability, "private-stream", "private-pat", "Authorization", "private-creator", "private-server"}
	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		_, writeErr := writer.Write([]byte(`{"success":false,"error":"capability ` + rawCapability + ` stream private-stream Authorization: Bearer private-pat creator private-creator server private-server"}`))
		require.NoError(t, writeErr)
	}))
	t.Cleanup(server.Close)
	transport := newTestClient(t, Config{BaseURL: server.URL, BearerToken: "private-pat"})
	access := IOStreamCapabilityAccessRequest{Capability: capability, Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7}

	for _, invoke := range []func() error{
		func() error {
			_, callErr := transport.IOStreamCapabilities().Wait(context.Background(), IOStreamCapabilityWaitRequest(access))
			return callErr
		},
		func() error { return transport.IOStreamCapabilities().Cancel(context.Background(), access) },
		func() error { return transport.IOStreamCapabilities().Unregister(context.Background(), access) },
	} {
		callErr := invoke()
		require.ErrorIs(t, callErr, ErrIOStreamCapabilityUnavailable)
		for _, privateValue := range privateValues {
			require.NotContains(t, callErr.Error(), privateValue)
		}
	}
}

func TestIOStreamCapabilityClientPreservesCancellationIdentity(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		select {
		case <-request.Context().Done():
		case <-time.After(time.Second):
			t.Error("request context was not canceled")
		}
	}))
	t.Cleanup(server.Close)
	transport := newTestClient(t, Config{BaseURL: server.URL})
	rawCapability := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("x", 32)))
	capability, err := ParseIOStreamCapability(rawCapability)
	require.NoError(t, err)
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	_, err = transport.IOStreamCapabilities().Wait(ctx, IOStreamCapabilityWaitRequest{Capability: capability, Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7})
	require.True(t, errors.Is(err, context.Canceled))
}

func TestIOStreamCapabilityClientMapsTypedNonsecretErrors(t *testing.T) {
	messages := []struct {
		message string
		target  error
	}{
		{message: ioStreamCapabilityConflictMessage, target: ErrIOStreamCapabilityConflict},
		{message: ioStreamCapabilityCleanupMessage, target: ErrIOStreamCapabilityCleanup},
		{message: ioStreamCapabilityUnavailableMessage, target: ErrIOStreamCapabilityUnavailable},
	}
	for _, testCase := range messages {
		t.Run(testCase.message, func(t *testing.T) {
			server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
				writer.Header().Set("Content-Type", "application/json")
				_, err := writer.Write([]byte(`{"success":false,"error":"` + testCase.message + `"}`))
				require.NoError(t, err)
			}))
			t.Cleanup(server.Close)
			transport := newTestClient(t, Config{BaseURL: server.URL})
			rawCapability := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("m", 32)))
			capability, err := ParseIOStreamCapability(rawCapability)
			require.NoError(t, err)
			callErr := transport.IOStreamCapabilities().Unregister(context.Background(), IOStreamCapabilityAccessRequest{Capability: capability, Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7})
			require.ErrorIs(t, callErr, testCase.target)
			require.Equal(t, testCase.target.Error(), callErr.Error())
		})
	}
}