ai-gateway-go/internal/provider/controlplane/service.go

package controlplane

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"aigateway.local/core/internal/provider"
	provideropenai "aigateway.local/core/internal/provider/openai"
)

const maxModelsResponseBytes int64 = 4 << 20

type Service struct {
	repository   *provider.Repository
	cipher       *provider.CredentialCipher
	allowPrivate bool
	client       *http.Client
}

func NewService(repository *provider.Repository, cipher *provider.CredentialCipher, allowPrivate bool) *Service {
	return &Service{
		repository:   repository,
		cipher:       cipher,
		allowPrivate: allowPrivate,
		client: &http.Client{
			Timeout: 10 * time.Second,
			Transport: &http.Transport{
				DialContext:           safeDialContext(allowPrivate),
				ForceAttemptHTTP2:     true,
				MaxIdleConns:          32,
				MaxIdleConnsPerHost:   4,
				IdleConnTimeout:       90 * time.Second,
				TLSHandshakeTimeout:   5 * time.Second,
				ResponseHeaderTimeout: 8 * time.Second,
			},
			CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
				return errors.New("upstream redirects are not allowed")
			},
		},
	}
}

func (s *Service) TestConnection(ctx context.Context, providerID string) (provider.ConnectionTestResult, error) {
	record, adapter, err := s.load(ctx, providerID)
	if err != nil {
		return provider.ConnectionTestResult{}, err
	}
	started := time.Now()
	response, err := s.doModelsRequest(ctx, adapter)
	latency := time.Since(started).Milliseconds()
	if err != nil {
		return provider.ConnectionTestResult{}, fmt.Errorf("%w: request provider %s: %v", provider.ErrProviderUpstream, record.Code, err)
	}
	defer response.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 32<<10))
	result := provider.ConnectionTestResult{
		Connected:  response.StatusCode >= 200 && response.StatusCode < 300,
		StatusCode: response.StatusCode,
		LatencyMS:  latency,
	}
	if result.Connected {
		result.Message = "连接成功"
	} else {
		result.Message = "上游返回 HTTP " + strconv.Itoa(response.StatusCode)
	}
	return result, nil
}

func (s *Service) SyncModels(ctx context.Context, providerID, actorID string) (provider.ModelSyncResult, error) {
	record, adapter, err := s.load(ctx, providerID)
	if err != nil {
		return provider.ModelSyncResult{}, err
	}
	if !hasCapability(record.Capabilities, string(provider.CapabilityModels)) {
		return provider.ModelSyncResult{}, errors.New("供应商未启用 models 能力")
	}
	response, err := s.doModelsRequest(ctx, adapter)
	if err != nil {
		return provider.ModelSyncResult{}, fmt.Errorf("%w: request provider %s: %v", provider.ErrProviderUpstream, record.Code, err)
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 32<<10))
		return provider.ModelSyncResult{}, fmt.Errorf("%w: provider %s returned HTTP %d", provider.ErrProviderUpstream, record.Code, response.StatusCode)
	}
	payload, err := io.ReadAll(io.LimitReader(response.Body, maxModelsResponseBytes+1))
	if err != nil {
		return provider.ModelSyncResult{}, fmt.Errorf("%w: read provider %s response: %v", provider.ErrProviderUpstream, record.Code, err)
	}
	if int64(len(payload)) > maxModelsResponseBytes {
		return provider.ModelSyncResult{}, fmt.Errorf("%w: provider %s model response exceeds 4 MiB", provider.ErrProviderUpstream, record.Code)
	}
	models, err := decodeModels(payload)
	if err != nil {
		return provider.ModelSyncResult{}, fmt.Errorf("%w: provider %s returned invalid model data: %v", provider.ErrProviderUpstream, record.Code, err)
	}
	return s.repository.SyncModels(ctx, providerID, actorID, models)
}

func (s *Service) ListModels(ctx context.Context, providerID string) ([]provider.Model, error) {
	if _, err := s.repository.Get(ctx, providerID); err != nil {
		return nil, err
	}
	return s.repository.ListModels(ctx, providerID)
}

func (s *Service) RotateCredentials(ctx context.Context, actorID string) (provider.CredentialRotationResult, error) {
	result := provider.CredentialRotationResult{
		ActiveVersion: s.cipher.ActiveVersion(), LoadedVersions: s.cipher.Versions(),
	}
	records, err := s.repository.List(ctx)
	if err != nil {
		return provider.CredentialRotationResult{}, err
	}
	rotations := make([]provider.CredentialRotation, 0, len(records))
	for _, record := range records {
		if record.CredentialKEKVersion == result.ActiveVersion {
			result.Skipped++
			continue
		}
		plaintext, err := s.cipher.Decrypt(record.EncryptedCredentials, record.CredentialKEKVersion)
		if err != nil {
			return provider.CredentialRotationResult{}, fmt.Errorf("provider %s credentials cannot be decrypted: %w", record.Code, err)
		}
		ciphertext, version, err := s.cipher.Encrypt(plaintext)
		if err != nil {
			return provider.CredentialRotationResult{}, err
		}
		rotations = append(rotations, provider.CredentialRotation{
			ProviderID: record.ID, FromVersion: record.CredentialKEKVersion,
			ToVersion: version, Ciphertext: ciphertext,
		})
	}
	if err := s.repository.RotateCredentials(ctx, actorID, rotations); err != nil {
		return provider.CredentialRotationResult{}, err
	}
	result.Rotated = len(rotations)
	return result, nil
}

func (s *Service) load(ctx context.Context, providerID string) (provider.Record, provider.Adapter, error) {
	record, err := s.repository.Get(ctx, providerID)
	if err != nil {
		return provider.Record{}, nil, err
	}
	validatedURL, err := provider.ValidateBaseURL(ctx, record.BaseURL, s.allowPrivate)
	if err != nil {
		return provider.Record{}, nil, fmt.Errorf("供应商地址校验失败: %w", err)
	}
	plaintext, err := s.cipher.Decrypt(record.EncryptedCredentials, record.CredentialKEKVersion)
	if err != nil {
		return provider.Record{}, nil, err
	}
	var credentials provider.Credentials
	if err := json.Unmarshal(plaintext, &credentials); err != nil {
		return provider.Record{}, nil, errors.New("供应商凭据格式无效")
	}
	switch record.Adapter {
	case "openai-compatible":
		adapter, err := provideropenai.New(validatedURL, credentials.APIKey)
		return record, adapter, err
	default:
		return provider.Record{}, nil, fmt.Errorf("不支持的供应商适配器 %q", record.Adapter)
	}
}

func (s *Service) doModelsRequest(ctx context.Context, adapter provider.Adapter) (*http.Response, error) {
	target := adapter.Target()
	target.Path = strings.TrimRight(target.Path, "/") + "/v1/models"
	target.RawPath = ""
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, target.String(), nil)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Accept", "application/json")
	adapter.Prepare(request)
	return s.client.Do(request)
}

func decodeModels(payload []byte) ([]provider.DiscoveredModel, error) {
	var response struct {
		Data []json.RawMessage `json:"data"`
	}
	if err := json.Unmarshal(payload, &response); err != nil {
		return nil, err
	}
	seen := make(map[string]struct{}, len(response.Data))
	models := make([]provider.DiscoveredModel, 0, len(response.Data))
	for _, raw := range response.Data {
		var item struct {
			ID      string `json:"id"`
			OwnedBy string `json:"owned_by"`
		}
		if err := json.Unmarshal(raw, &item); err != nil {
			return nil, err
		}
		item.ID = strings.TrimSpace(item.ID)
		if item.ID == "" || len(item.ID) > 512 {
			return nil, errors.New("model id must contain 1 to 512 characters")
		}
		if _, exists := seen[item.ID]; exists {
			continue
		}
		seen[item.ID] = struct{}{}
		models = append(models, provider.DiscoveredModel{
			ProviderModelID: item.ID,
			OwnedBy:         strings.TrimSpace(item.OwnedBy),
			Metadata:        append(json.RawMessage(nil), raw...),
		})
	}
	return models, nil
}

func hasCapability(capabilities []string, expected string) bool {
	for _, capability := range capabilities {
		if capability == expected {
			return true
		}
	}
	return false
}

func safeDialContext(allowPrivate bool) func(context.Context, string, string) (net.Conn, error) {
	dialer := &net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}
	if allowPrivate {
		return dialer.DialContext
	}
	return func(ctx context.Context, network, address string) (net.Conn, error) {
		host, port, err := net.SplitHostPort(address)
		if err != nil {
			return nil, err
		}
		addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
		if err != nil {
			return nil, err
		}
		if len(addresses) == 0 {
			return nil, errors.New("upstream host did not resolve")
		}
		for _, address := range addresses {
			if !isPublicAddress(address.IP) {
				return nil, fmt.Errorf("upstream resolved to blocked address %s", address.IP)
			}
		}
		return dialer.DialContext(ctx, network, net.JoinHostPort(addresses[0].IP.String(), port))
	}
}

func isPublicAddress(ip net.IP) bool {
	return ip != nil && !ip.IsPrivate() && !ip.IsLoopback() && !ip.IsLinkLocalUnicast() &&
		!ip.IsLinkLocalMulticast() && !ip.IsMulticast() && !ip.IsUnspecified()
}

var _ provider.AdminOperations = (*Service)(nil)