AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,262 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user