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

package client

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type TransferURL struct {
	URL       string    `json:"url"`
	Method    string    `json:"method"`
	ExpiresAt time.Time `json:"expires_at"`
}

type DownloadURLRequest struct {
	ServerID   uint64 `json:"server_id"`
	Path       string `json:"path"`
	TTLSeconds int    `json:"ttl_seconds,omitempty"`
}

type UploadURLRequest struct {
	ServerID      uint64 `json:"server_id"`
	Path          string `json:"path"`
	TTLSeconds    int    `json:"ttl_seconds,omitempty"`
	Mode          string `json:"mode,omitempty"`
	CreateDirs    bool   `json:"create_dirs,omitempty"`
	IfMatchSHA256 string `json:"if_match_sha256,omitempty"`
}

type UploadTransfer struct {
	Body          io.Reader
	ContentLength int64
	SHA256        string
}

type UploadResult struct {
	Size   int64  `json:"size"`
	SHA256 string `json:"sha256"`
}

type uploadResultPayload struct {
	Size   *int64  `json:"size"`
	SHA256 *string `json:"sha256"`
}

func RequestDownloadURL(ctx context.Context, client *Client, request DownloadURLRequest) (TransferURL, error) {
	result, err := CallTool[DownloadURLRequest, TransferURL](ctx, client, ToolCall[DownloadURLRequest]{Name: "fs.download_url", Arguments: request})
	if err != nil {
		return TransferURL{}, err
	}
	return client.validateTransferURL(result.StructuredContent, http.MethodGet)
}

func RequestUploadURL(ctx context.Context, client *Client, request UploadURLRequest) (TransferURL, error) {
	result, err := CallTool[UploadURLRequest, TransferURL](ctx, client, ToolCall[UploadURLRequest]{Name: "fs.upload_url", Arguments: request})
	if err != nil {
		return TransferURL{}, err
	}
	return client.validateTransferURL(result.StructuredContent, http.MethodPost)
}

func (client *Client) DownloadTransfer(ctx context.Context, transfer TransferURL, destination io.Writer) (int64, error) {
	validated, err := client.validateTransferURL(transfer, http.MethodGet)
	if err != nil {
		return 0, err
	}
	requestContext, cancel := client.transferContext(ctx)
	defer cancel()
	request, err := http.NewRequestWithContext(requestContext, http.MethodGet, validated.URL, nil)
	if err != nil {
		return 0, errorsNewRedacted("create transfer request", err)
	}
	response, err := client.transferHTTPClient().Do(request)
	if err != nil {
		if requestContext.Err() != nil {
			return 0, fmt.Errorf("download transfer: %w", requestContext.Err())
		}
		return 0, errorsNewRedacted("download transfer", err)
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return 0, &HTTPError{StatusCode: response.StatusCode}
	}
	written, err := io.Copy(destination, io.LimitReader(response.Body, client.maxTransferBytes))
	if err != nil {
		return written, fmt.Errorf("copy download transfer: %w", err)
	}
	var overflow [1]byte
	read, err := response.Body.Read(overflow[:])
	if err != nil && err != io.EOF {
		return written, fmt.Errorf("probe download transfer size: %w", err)
	}
	if read > 0 {
		return written, ErrTransferTooLarge
	}
	return written, nil
}

func (client *Client) UploadTransfer(ctx context.Context, transfer TransferURL, upload UploadTransfer) (UploadResult, error) {
	validated, err := client.validateTransferURL(transfer, http.MethodPost)
	if err != nil {
		return UploadResult{}, err
	}
	if upload.Body == nil || upload.ContentLength <= 0 {
		return UploadResult{}, fmt.Errorf("upload content length: %w", ErrInvalidConfig)
	}
	if upload.ContentLength > client.maxTransferBytes {
		return UploadResult{}, ErrTransferTooLarge
	}
	transferURL, err := url.Parse(validated.URL)
	if err != nil {
		return UploadResult{}, errorsNewRedacted("parse transfer URL", err)
	}
	if upload.SHA256 != "" {
		query := transferURL.Query()
		query.Set("sha256", upload.SHA256)
		transferURL.RawQuery = query.Encode()
	}
	requestContext, cancel := client.transferContext(ctx)
	defer cancel()
	request, err := http.NewRequestWithContext(requestContext, http.MethodPost, transferURL.String(), io.LimitReader(upload.Body, upload.ContentLength))
	if err != nil {
		return UploadResult{}, errorsNewRedacted("create transfer request", err)
	}
	request.ContentLength = upload.ContentLength
	response, err := client.transferHTTPClient().Do(request)
	if err != nil {
		if requestContext.Err() != nil {
			return UploadResult{}, fmt.Errorf("upload transfer: %w", requestContext.Err())
		}
		return UploadResult{}, errorsNewRedacted("upload transfer", err)
	}
	defer response.Body.Close()
	body, err := readBounded(response.Body, client.maxResponseBytes)
	if err != nil {
		return UploadResult{}, err
	}
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return UploadResult{}, &HTTPError{StatusCode: response.StatusCode, Message: Redact(string(body))}
	}
	var payload uploadResultPayload
	if err := json.NewDecoder(bytes.NewReader(body)).Decode(&payload); err != nil {
		return UploadResult{}, fmt.Errorf("decode upload result: %w", err)
	}
	if payload.Size == nil || payload.SHA256 == nil || *payload.Size <= 0 || *payload.SHA256 == "" {
		return UploadResult{}, fmt.Errorf("decode upload result: %w", ErrSemanticFailure)
	}
	return UploadResult{Size: *payload.Size, SHA256: *payload.SHA256}, nil
}

func (client *Client) transferHTTPClient() *http.Client {
	clone := *client.httpClient
	clone.Jar = nil
	clone.CheckRedirect = rejectRedirect
	return &clone
}

func (client *Client) validateTransferURL(transfer TransferURL, expectedMethod string) (TransferURL, error) {
	parsed, err := url.Parse(transfer.URL)
	if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") {
		return TransferURL{}, fmt.Errorf("transfer URL: %w", ErrInvalidConfig)
	}
	if parsed.User != nil || !client.sameOrigin(parsed) {
		return TransferURL{}, fmt.Errorf("transfer origin: %w", ErrInvalidConfig)
	}
	if transfer.ExpiresAt.IsZero() || !time.Now().Before(transfer.ExpiresAt) {
		return TransferURL{}, ErrTransferExpired
	}
	if !strings.EqualFold(transfer.Method, expectedMethod) {
		return TransferURL{}, fmt.Errorf("transfer method: %w", ErrInvalidConfig)
	}
	transfer.Method = expectedMethod
	return transfer, nil
}

func (client *Client) sameOrigin(candidate *url.URL) bool {
	return strings.EqualFold(candidate.Scheme, client.baseURL.Scheme) && strings.EqualFold(candidate.Host, client.baseURL.Host)
}