mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-21 02:30:14 +00:00
test(agentcompat): add protocol clients
Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
This commit is contained in:
naiba
naiba/CloudCode
@@ -0,0 +1,249 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestClient_TransferURLConsumptionOmitsAuthorization(t *testing.T) {
|
||||
var server *httptest.Server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
switch request.URL.Path {
|
||||
case "/mcp":
|
||||
require.Equal(t, "Bearer mcp-token", request.Header.Get("Authorization"))
|
||||
var rpcRequest testJSONRPCRequest
|
||||
require.NoError(t, json.NewDecoder(request.Body).Decode(&rpcRequest))
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":` + jsonNumber(rpcRequest.ID) + `,"result":{"structuredContent":{"url":"` + server.URL + `/mcp/download/one-time-token","method":"GET","expires_at":"2030-01-02T03:04:05Z"}}}`))
|
||||
case "/mcp/download/one-time-token":
|
||||
require.Empty(t, request.Header.Get("Authorization"))
|
||||
require.Empty(t, request.Header.Get("Origin"))
|
||||
_, _ = writer.Write([]byte("transfer payload"))
|
||||
default:
|
||||
http.NotFound(writer, request)
|
||||
}
|
||||
}))
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
client := newTestClient(t, Config{
|
||||
BaseURL: server.URL,
|
||||
BearerToken: "mcp-token",
|
||||
Origin: server.URL,
|
||||
RequestTimeout: time.Second,
|
||||
MaxResponseBytes: 1024,
|
||||
MaxTransferBytes: 1024,
|
||||
})
|
||||
transfer, err := RequestDownloadURL(context.Background(), client, DownloadURLRequest{ServerID: 7, Path: "/tmp/report", TTLSeconds: 30})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.MethodGet, transfer.Method)
|
||||
|
||||
var destination bytes.Buffer
|
||||
written, err := client.DownloadTransfer(context.Background(), transfer, &destination)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(len("transfer payload")), written)
|
||||
require.Equal(t, "transfer payload", destination.String())
|
||||
}
|
||||
|
||||
func TestClient_TransferURLRejectsCrossOrigin(t *testing.T) {
|
||||
foreignServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
||||
t.Fatal("cross-origin transfer request was dispatched")
|
||||
}))
|
||||
t.Cleanup(foreignServer.Close)
|
||||
|
||||
client := newTestClient(t, Config{
|
||||
BaseURL: "http://dashboard.example",
|
||||
RequestTimeout: time.Second,
|
||||
MaxResponseBytes: 1024,
|
||||
MaxTransferBytes: 1024,
|
||||
})
|
||||
|
||||
var destination bytes.Buffer
|
||||
_, err := client.DownloadTransfer(context.Background(), TransferURL{URL: foreignServer.URL + "/mcp/download/token", Method: http.MethodGet}, &destination)
|
||||
require.Error(t, err)
|
||||
require.True(t, errors.Is(err, ErrInvalidConfig))
|
||||
}
|
||||
|
||||
func TestClient_TransferURLRejectsCrossOriginRedirect(t *testing.T) {
|
||||
var foreignRequests atomic.Int32
|
||||
foreignServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
||||
foreignRequests.Add(1)
|
||||
}))
|
||||
t.Cleanup(foreignServer.Close)
|
||||
|
||||
dashboardServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
http.Redirect(writer, &http.Request{}, foreignServer.URL+"/stolen-token", http.StatusTemporaryRedirect)
|
||||
}))
|
||||
t.Cleanup(dashboardServer.Close)
|
||||
|
||||
client := newTestClient(t, Config{
|
||||
BaseURL: dashboardServer.URL,
|
||||
RequestTimeout: time.Second,
|
||||
MaxResponseBytes: 1024,
|
||||
MaxTransferBytes: 1024,
|
||||
})
|
||||
|
||||
var destination bytes.Buffer
|
||||
_, err := client.DownloadTransfer(context.Background(), TransferURL{
|
||||
URL: dashboardServer.URL + "/mcp/download/token",
|
||||
Method: http.MethodGet,
|
||||
ExpiresAt: time.Now().Add(time.Minute),
|
||||
}, &destination)
|
||||
require.ErrorIs(t, err, ErrRedirect)
|
||||
require.Zero(t, foreignRequests.Load())
|
||||
}
|
||||
|
||||
func TestClient_TransferURLRejectsSameOriginRedirect(t *testing.T) {
|
||||
var redirected atomic.Int32
|
||||
var dashboardServer *httptest.Server
|
||||
dashboardServer = httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
if request.URL.Path == "/capture" {
|
||||
redirected.Add(1)
|
||||
return
|
||||
}
|
||||
http.Redirect(writer, request, dashboardServer.URL+"/capture", http.StatusTemporaryRedirect)
|
||||
}))
|
||||
t.Cleanup(dashboardServer.Close)
|
||||
|
||||
client := newTestClient(t, Config{BaseURL: dashboardServer.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024})
|
||||
var destination bytes.Buffer
|
||||
_, err := client.DownloadTransfer(context.Background(), TransferURL{
|
||||
URL: dashboardServer.URL + "/mcp/download/token",
|
||||
Method: http.MethodGet,
|
||||
ExpiresAt: time.Now().Add(time.Minute),
|
||||
}, &destination)
|
||||
require.ErrorIs(t, err, ErrRedirect)
|
||||
require.Zero(t, redirected.Load())
|
||||
}
|
||||
|
||||
func TestClient_TransferURLRejectsExpiredCapability(t *testing.T) {
|
||||
client := newTestClient(t, Config{BaseURL: "http://dashboard.example", RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024})
|
||||
var destination bytes.Buffer
|
||||
_, err := client.DownloadTransfer(context.Background(), TransferURL{
|
||||
URL: "http://dashboard.example/mcp/download/token",
|
||||
Method: http.MethodGet,
|
||||
ExpiresAt: time.Now().Add(-time.Second),
|
||||
}, &destination)
|
||||
require.ErrorIs(t, err, ErrTransferExpired)
|
||||
}
|
||||
|
||||
func TestClient_TransferURLRejectsMissingExpiryBeforeDispatch(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
requests.Add(1)
|
||||
writer.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024})
|
||||
var destination bytes.Buffer
|
||||
_, err := client.DownloadTransfer(context.Background(), TransferURL{
|
||||
URL: server.URL + "/mcp/download/token",
|
||||
Method: http.MethodGet,
|
||||
}, &destination)
|
||||
require.ErrorIs(t, err, ErrTransferExpired)
|
||||
require.Zero(t, requests.Load())
|
||||
}
|
||||
|
||||
func TestClient_RequestUploadURLRejectsMissingExpiryBeforeDispatch(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
var server *httptest.Server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
requests.Add(1)
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"structuredContent":{"url":"` + server.URL + `/mcp/upload/token","method":"POST"}}}`))
|
||||
}))
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024})
|
||||
_, err := RequestUploadURL(context.Background(), client, UploadURLRequest{ServerID: 7, Path: "/tmp/report"})
|
||||
require.ErrorIs(t, err, ErrTransferExpired)
|
||||
require.Equal(t, int32(1), requests.Load())
|
||||
}
|
||||
|
||||
func TestClient_TransferRejectsOversizeWithoutWritingPastLimit(t *testing.T) {
|
||||
dashboardServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = writer.Write([]byte(strings.Repeat("x", 65)))
|
||||
}))
|
||||
t.Cleanup(dashboardServer.Close)
|
||||
|
||||
client := newTestClient(t, Config{
|
||||
BaseURL: dashboardServer.URL,
|
||||
RequestTimeout: time.Second,
|
||||
MaxResponseBytes: 1024,
|
||||
MaxTransferBytes: 64,
|
||||
})
|
||||
|
||||
var destination bytes.Buffer
|
||||
written, err := client.DownloadTransfer(context.Background(), TransferURL{
|
||||
URL: dashboardServer.URL + "/mcp/download/token",
|
||||
Method: http.MethodGet,
|
||||
ExpiresAt: time.Now().Add(time.Minute),
|
||||
}, &destination)
|
||||
require.ErrorIs(t, err, ErrTransferTooLarge)
|
||||
require.Equal(t, int64(64), written)
|
||||
require.Len(t, destination.Bytes(), 64)
|
||||
}
|
||||
|
||||
func TestClient_UploadRejectsChunkedBodyWithoutDispatch(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
dashboardServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
||||
requests.Add(1)
|
||||
}))
|
||||
t.Cleanup(dashboardServer.Close)
|
||||
|
||||
client := newTestClient(t, Config{
|
||||
BaseURL: dashboardServer.URL,
|
||||
RequestTimeout: time.Second,
|
||||
MaxResponseBytes: 1024,
|
||||
MaxTransferBytes: 64,
|
||||
})
|
||||
_, err := client.UploadTransfer(context.Background(), TransferURL{
|
||||
URL: dashboardServer.URL + "/mcp/upload/token",
|
||||
Method: http.MethodPost,
|
||||
ExpiresAt: time.Now().Add(time.Minute),
|
||||
}, UploadTransfer{Body: strings.NewReader("hidden body"), ContentLength: 0})
|
||||
require.ErrorIs(t, err, ErrInvalidConfig)
|
||||
require.Zero(t, requests.Load())
|
||||
}
|
||||
|
||||
func TestClient_UploadRejectsMisleadingSuccessEnvelope(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"success":true,"data":{"size":11,"sha256":"abc"}}`))
|
||||
}))
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024})
|
||||
_, err := client.UploadTransfer(context.Background(), TransferURL{
|
||||
URL: server.URL + "/mcp/upload/token",
|
||||
Method: http.MethodPost,
|
||||
ExpiresAt: time.Now().Add(time.Minute),
|
||||
}, UploadTransfer{Body: strings.NewReader("payload"), ContentLength: int64(len("payload"))})
|
||||
require.ErrorIs(t, err, ErrSemanticFailure)
|
||||
}
|
||||
|
||||
func TestClient_UploadRejectsMissingRequiredResultFields(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"size":7}`))
|
||||
}))
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024})
|
||||
_, err := client.UploadTransfer(context.Background(), TransferURL{
|
||||
URL: server.URL + "/mcp/upload/token",
|
||||
Method: http.MethodPost,
|
||||
ExpiresAt: time.Now().Add(time.Minute),
|
||||
}, UploadTransfer{Body: strings.NewReader("payload"), ContentLength: int64(len("payload"))})
|
||||
require.ErrorIs(t, err, ErrSemanticFailure)
|
||||
}
|
||||
Reference in New Issue
Block a user