vps-control/control-plane/backend/agent.go

package main

import (
	"bytes"
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"
)

const maxAgentResponse = 2 << 20

type AgentClient struct {
	http           *http.Client
	credentialsDir string
}

func NewAgentClient(credentialsDir string) *AgentClient {
	transport := &http.Transport{
		Proxy:                 nil,
		DialContext:           (&net.Dialer{Timeout: 3 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
		TLSClientConfig:       &tls.Config{MinVersion: tls.VersionTLS12},
		TLSHandshakeTimeout:   3 * time.Second,
		ResponseHeaderTimeout: 5 * time.Second,
		MaxConnsPerHost:       8,
		MaxIdleConnsPerHost:   4,
		IdleConnTimeout:       30 * time.Second,
	}
	return &AgentClient{
		http: &http.Client{
			Transport: transport,
			Timeout:   7 * time.Second,
			CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
				return errors.New("agent redirects are disabled")
			},
		},
		credentialsDir: credentialsDir,
	}
}

func validateAgentEndpoint(endpoint string) error {
	parsed, err := url.Parse(endpoint)
	if err != nil {
		return errors.New("invalid URL")
	}
	if parsed.Scheme != "http" && parsed.Scheme != "https" {
		return errors.New("scheme must be http or https")
	}
	if parsed.User != nil || (parsed.Path != "" && parsed.Path != "/") || parsed.RawQuery != "" || parsed.Fragment != "" {
		return errors.New("endpoint must contain only scheme, private IP, and port")
	}
	if parsed.Port() == "" {
		return errors.New("endpoint must include a port")
	}
	ip := net.ParseIP(parsed.Hostname())
	if ip == nil || !(ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast()) {
		return errors.New("endpoint must use a private literal IP address")
	}
	return nil
}

func validCredentialRef(reference string) bool {
	if len(reference) < 1 || len(reference) > 128 {
		return false
	}
	for _, char := range reference {
		if !(char >= 'a' && char <= 'z') && !(char >= 'A' && char <= 'Z') &&
			!(char >= '0' && char <= '9') && char != '_' && char != '-' && char != '.' {
			return false
		}
	}
	return reference != "." && reference != ".."
}

func (client *AgentClient) credential(reference string) (string, error) {
	if !validCredentialRef(reference) {
		return "", errors.New("invalid credential reference")
	}
	path := filepath.Join(client.credentialsDir, reference)
	info, err := os.Lstat(path)
	if err != nil {
		return "", err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return "", errors.New("credential must be a regular file")
	}
	// Docker Compose file-backed secrets are mounted read-only as 0444 on some
	// engines. Accept read-only mounts, but never a group/world-writable secret.
	if info.Mode().Perm()&0o022 != 0 {
		return "", errors.New("credential file must not be group/world writable")
	}
	if info.Size() > 4096 {
		return "", errors.New("credential file is too large")
	}
	value, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	token := strings.TrimSpace(string(value))
	if len(token) < 32 {
		return "", errors.New("credential is too short")
	}
	return token, nil
}

func (client *AgentClient) Do(ctx context.Context, server Server, method, path string, body []byte) (int, []byte, error) {
	if err := validateAgentEndpoint(server.AgentEndpoint); err != nil {
		return 0, nil, err
	}
	if !strings.HasPrefix(path, "/v1/") {
		return 0, nil, errors.New("invalid agent API path")
	}
	token, err := client.credential(server.CredentialRef)
	if err != nil {
		return 0, nil, fmt.Errorf("load agent credential: %w", err)
	}
	request, err := http.NewRequestWithContext(ctx, method,
		strings.TrimRight(server.AgentEndpoint, "/")+path, bytes.NewReader(body))
	if err != nil {
		return 0, nil, err
	}
	request.Header.Set("Authorization", "Bearer "+token)
	request.Header.Set("Accept", "application/json")
	if len(body) > 0 {
		request.Header.Set("Content-Type", "application/json")
	}
	response, err := client.http.Do(request)
	if err != nil {
		return 0, nil, err
	}
	defer response.Body.Close()
	limited := io.LimitReader(response.Body, maxAgentResponse+1)
	payload, err := io.ReadAll(limited)
	if err != nil {
		return 0, nil, err
	}
	if len(payload) > maxAgentResponse {
		return 0, nil, errors.New("agent response is too large")
	}
	return response.StatusCode, payload, nil
}