366 lines
12 KiB
Go
366 lines
12 KiB
Go
package plugin
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
managedaccess "billing/internal/access"
|
|
"billing/internal/pricing"
|
|
"billing/internal/repository"
|
|
)
|
|
|
|
const (
|
|
hostAuthListMethod = "host.auth.list"
|
|
callerScopeMetadata = "caller_scope"
|
|
selectedAuthMetadata = "selected_auth_id"
|
|
)
|
|
|
|
type HostCaller interface {
|
|
Call(method string, request []byte) ([]byte, error)
|
|
}
|
|
|
|
func (a *App) SetHostCaller(caller HostCaller) {
|
|
a.mu.Lock()
|
|
a.host = caller
|
|
a.mu.Unlock()
|
|
}
|
|
|
|
func (a *App) frontendIdentifier() ([]byte, error) {
|
|
return OKEnvelope(FrontendIdentifierResponse{Identifier: PluginName})
|
|
}
|
|
|
|
func (a *App) authenticate(raw []byte) ([]byte, error) {
|
|
var req FrontendAuthRequest
|
|
if err := json.Unmarshal(raw, &req); err != nil {
|
|
return nil, fmt.Errorf("解析下游认证请求: %w", err)
|
|
}
|
|
credential := requestCredential(req.Headers)
|
|
if credential == "" {
|
|
return OKEnvelope(FrontendAuthResponse{Authenticated: false})
|
|
}
|
|
a.mu.RLock()
|
|
store := a.store
|
|
closed := a.closed
|
|
a.mu.RUnlock()
|
|
if closed || store == nil {
|
|
return OKEnvelope(FrontendAuthResponse{Authenticated: false})
|
|
}
|
|
key, err := store.ManagedKeyByCredential(context.Background(), credential)
|
|
if err != nil || key.Status != managedaccess.StatusActive {
|
|
return OKEnvelope(FrontendAuthResponse{Authenticated: false})
|
|
}
|
|
a.scopes.Store(managedaccess.CallerScope(key.ID), key.ID)
|
|
return OKEnvelope(FrontendAuthResponse{
|
|
Authenticated: true,
|
|
Principal: key.ID,
|
|
Metadata: map[string]string{"key_name": key.Name},
|
|
})
|
|
}
|
|
|
|
func requestCredential(headers http.Header) string {
|
|
authorization := strings.TrimSpace(headers.Get("Authorization"))
|
|
if fields := strings.Fields(authorization); len(fields) == 2 && strings.EqualFold(fields[0], "Bearer") {
|
|
return fields[1]
|
|
}
|
|
return strings.TrimSpace(headers.Get("X-Api-Key"))
|
|
}
|
|
|
|
func (a *App) interceptRequest(raw []byte, afterAuth bool) ([]byte, error) {
|
|
var req RequestInterceptRequest
|
|
if err := json.Unmarshal(raw, &req); err != nil {
|
|
return nil, fmt.Errorf("解析请求拦截事件: %w", err)
|
|
}
|
|
keyID := metadataString(req.Metadata, callerScopeMetadata)
|
|
if keyID == "" && req.RequestID != "" {
|
|
if pending, ok := a.pending.Load(req.RequestID); ok {
|
|
keyID, _ = pending.(string)
|
|
}
|
|
}
|
|
if keyID == "" {
|
|
return OKEnvelope(RequestInterceptResponse{})
|
|
}
|
|
key, err := a.managedKeyByReference(keyID)
|
|
if err != nil || key.Status != managedaccess.StatusActive {
|
|
return OKEnvelope(denyRequest(http.StatusForbidden, "access_credential_inactive", "Key 已停用"))
|
|
}
|
|
requestedModel := strings.TrimSpace(req.RequestedModel)
|
|
if requestedModel == "" {
|
|
requestedModel = strings.TrimSpace(req.Model)
|
|
}
|
|
if !managedKeyAllowsModel(key, requestedModel) {
|
|
return OKEnvelope(denyRequest(http.StatusForbidden, "model_not_allowed", "该 Key 不允许使用此模型"))
|
|
}
|
|
if !afterAuth {
|
|
store, ok := a.currentStore()
|
|
if !ok {
|
|
return OKEnvelope(denyRequest(http.StatusServiceUnavailable, "billing_unavailable", "额度数据库不可用"))
|
|
}
|
|
_, authorizeErr := store.AuthorizeBilling(context.Background(), key.ID, req.RequestID, time.Now().UTC())
|
|
switch {
|
|
case errors.Is(authorizeErr, repository.ErrBillingQuotaExhausted):
|
|
return OKEnvelope(denyRequest(http.StatusTooManyRequests, "billing_quota_exhausted", "该 Key 的可用余额已用尽"))
|
|
case errors.Is(authorizeErr, repository.ErrBillingConcurrency):
|
|
return OKEnvelope(denyRequest(http.StatusTooManyRequests, "billing_concurrency_exceeded", "该 Key 的并发请求已达到上限"))
|
|
case authorizeErr != nil:
|
|
return OKEnvelope(denyRequest(http.StatusServiceUnavailable, "billing_unavailable", "额度检查暂时不可用"))
|
|
}
|
|
if req.RequestID != "" {
|
|
a.pending.Store(req.RequestID, key.ID)
|
|
}
|
|
}
|
|
if afterAuth {
|
|
billingModel := strings.TrimSpace(req.Model)
|
|
if billingModel == "" {
|
|
billingModel = requestedModel
|
|
}
|
|
if _, found := a.priceForModel(billingModel); !found {
|
|
return OKEnvelope(denyRequest(http.StatusServiceUnavailable, "billing_price_unavailable", "该模型尚未配置价格"))
|
|
}
|
|
}
|
|
if afterAuth && key.RouteMode == managedaccess.RouteStrict {
|
|
account, accountErr := a.upstreamAccount(key.UpstreamAccountID)
|
|
selectedAuthID := metadataString(req.Metadata, selectedAuthMetadata)
|
|
if accountErr != nil || selectedAuthID == "" || selectedAuthID != account.CPAAuthID {
|
|
return OKEnvelope(denyRequest(http.StatusServiceUnavailable, "bound_upstream_unavailable", "指定的上游账号不可用"))
|
|
}
|
|
}
|
|
return OKEnvelope(RequestInterceptResponse{})
|
|
}
|
|
|
|
func (a *App) pickScheduler(raw []byte) ([]byte, error) {
|
|
var req SchedulerPickRequest
|
|
if err := json.Unmarshal(raw, &req); err != nil {
|
|
return nil, fmt.Errorf("解析调度请求: %w", err)
|
|
}
|
|
a.observeSchedulerCandidates(req.Candidates)
|
|
keyID := metadataString(req.Options.Metadata, callerScopeMetadata)
|
|
if keyID == "" {
|
|
return OKEnvelope(SchedulerPickResponse{Handled: false})
|
|
}
|
|
key, err := a.managedKeyByReference(keyID)
|
|
if err != nil || key.Status != managedaccess.StatusActive {
|
|
return ErrorEnvelope("access_credential_inactive", "Key 已停用", http.StatusForbidden), nil
|
|
}
|
|
if key.RouteMode != managedaccess.RouteStrict {
|
|
return OKEnvelope(SchedulerPickResponse{Handled: false})
|
|
}
|
|
account, err := a.upstreamAccount(key.UpstreamAccountID)
|
|
if err != nil {
|
|
return ErrorEnvelope("bound_upstream_unavailable", "指定的上游账号不可用", http.StatusServiceUnavailable), nil
|
|
}
|
|
for _, candidate := range req.Candidates {
|
|
if candidate.ID == account.CPAAuthID {
|
|
return OKEnvelope(SchedulerPickResponse{AuthID: candidate.ID, Handled: true})
|
|
}
|
|
}
|
|
return ErrorEnvelope("bound_upstream_unavailable", "指定的上游账号不在当前可用候选中", http.StatusServiceUnavailable), nil
|
|
}
|
|
|
|
func (a *App) observeSchedulerCandidates(candidates []SchedulerAuthCandidate) {
|
|
now := time.Now().UTC()
|
|
for _, candidate := range candidates {
|
|
if strings.TrimSpace(candidate.ID) == "" {
|
|
continue
|
|
}
|
|
display := firstNonEmpty(
|
|
candidate.Attributes["label"], candidate.Attributes["email"], candidate.Attributes["account"],
|
|
candidate.Attributes["name"], candidate.ID,
|
|
)
|
|
a.observedUpstreams.Store(candidate.ID, managedaccess.UpstreamAccount{
|
|
CPAAuthID: candidate.ID, Provider: candidate.Provider, DisplayName: display,
|
|
Status: candidate.Status, Priority: candidate.Priority, LastSeenAt: now,
|
|
})
|
|
}
|
|
}
|
|
|
|
func (a *App) flushObservedUpstreams(ctx context.Context, store interface {
|
|
SyncUpstreamAccounts(context.Context, []managedaccess.UpstreamAccount) error
|
|
}) error {
|
|
var accounts []managedaccess.UpstreamAccount
|
|
a.observedUpstreams.Range(func(_, value any) bool {
|
|
if account, ok := value.(managedaccess.UpstreamAccount); ok {
|
|
accounts = append(accounts, account)
|
|
}
|
|
return true
|
|
})
|
|
if len(accounts) == 0 {
|
|
return nil
|
|
}
|
|
return store.SyncUpstreamAccounts(ctx, accounts)
|
|
}
|
|
|
|
func managedKeyAllowsModel(key managedaccess.ManagedKey, model string) bool {
|
|
if key.AllModels {
|
|
return true
|
|
}
|
|
model = strings.ToLower(strings.TrimSpace(model))
|
|
if model == "" {
|
|
return false
|
|
}
|
|
for _, allowed := range key.Models {
|
|
if modelPatternMatches(allowed, model) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func modelPatternMatches(pattern, model string) bool {
|
|
pattern = normalizeModelName(pattern)
|
|
model = normalizeModelName(model)
|
|
if pattern == "" || model == "" {
|
|
return false
|
|
}
|
|
if pattern == "*" {
|
|
return true
|
|
}
|
|
parts := strings.Split(pattern, "*")
|
|
position := 0
|
|
if !strings.HasPrefix(pattern, "*") {
|
|
if !strings.HasPrefix(model, parts[0]) {
|
|
return false
|
|
}
|
|
position = len(parts[0])
|
|
parts = parts[1:]
|
|
}
|
|
for index, part := range parts {
|
|
if part == "" {
|
|
continue
|
|
}
|
|
last := index == len(parts)-1 && !strings.HasSuffix(pattern, "*")
|
|
if last {
|
|
return strings.HasSuffix(model[position:], part)
|
|
}
|
|
offset := strings.Index(model[position:], part)
|
|
if offset < 0 {
|
|
return false
|
|
}
|
|
position += offset + len(part)
|
|
}
|
|
return strings.HasSuffix(pattern, "*") || position == len(model)
|
|
}
|
|
|
|
func normalizeModelName(model string) string {
|
|
return strings.ToLower(strings.TrimSpace(model))
|
|
}
|
|
|
|
func (a *App) priceForModel(model string) (pricing.Policy, bool) {
|
|
a.mu.RLock()
|
|
defer a.mu.RUnlock()
|
|
policy, found := a.prices[normalizeModelName(model)]
|
|
return policy, found
|
|
}
|
|
|
|
func denyRequest(status int, code, message string) RequestInterceptResponse {
|
|
body, _ := json.Marshal(map[string]any{"error": map[string]string{"code": code, "message": message}})
|
|
return RequestInterceptResponse{
|
|
Terminate: true, StatusCode: status,
|
|
ResponseHeaders: http.Header{"Content-Type": []string{"application/json; charset=utf-8"}},
|
|
ResponseBody: body,
|
|
}
|
|
}
|
|
|
|
func metadataString(metadata map[string]any, key string) string {
|
|
if metadata == nil {
|
|
return ""
|
|
}
|
|
value, _ := metadata[key].(string)
|
|
return strings.TrimSpace(value)
|
|
}
|
|
|
|
func (a *App) managedKeyByID(id string) (managedaccess.ManagedKey, error) {
|
|
a.mu.RLock()
|
|
store := a.store
|
|
a.mu.RUnlock()
|
|
if store == nil {
|
|
return managedaccess.ManagedKey{}, errors.New("Key 数据库尚未初始化")
|
|
}
|
|
return store.ManagedKeyByID(context.Background(), id)
|
|
}
|
|
|
|
func (a *App) managedKeyByReference(reference string) (managedaccess.ManagedKey, error) {
|
|
a.mu.RLock()
|
|
store := a.store
|
|
a.mu.RUnlock()
|
|
if store == nil {
|
|
return managedaccess.ManagedKey{}, errors.New("Key 数据库尚未初始化")
|
|
}
|
|
if id, ok := a.scopes.Load(reference); ok {
|
|
if keyID, valid := id.(string); valid {
|
|
return store.ManagedKeyByID(context.Background(), keyID)
|
|
}
|
|
}
|
|
return store.ManagedKeyByReference(context.Background(), reference)
|
|
}
|
|
|
|
func (a *App) upstreamAccount(id string) (managedaccess.UpstreamAccount, error) {
|
|
a.mu.RLock()
|
|
store := a.store
|
|
a.mu.RUnlock()
|
|
if store == nil {
|
|
return managedaccess.UpstreamAccount{}, errors.New("上游账号数据库尚未初始化")
|
|
}
|
|
return store.UpstreamAccountByID(context.Background(), id)
|
|
}
|
|
|
|
func (a *App) syncUpstreamAccounts(ctx context.Context) error {
|
|
a.mu.RLock()
|
|
caller := a.host
|
|
store := a.store
|
|
a.mu.RUnlock()
|
|
if caller == nil {
|
|
return errors.New("CPA host callback 尚未初始化")
|
|
}
|
|
if store == nil {
|
|
return errors.New("上游账号数据库尚未初始化")
|
|
}
|
|
raw, err := caller.Call(hostAuthListMethod, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("调用 host.auth.list: %w", err)
|
|
}
|
|
var envelope Envelope
|
|
if err := json.Unmarshal(raw, &envelope); err != nil {
|
|
return fmt.Errorf("解析 host.auth.list envelope: %w", err)
|
|
}
|
|
if !envelope.OK {
|
|
if envelope.Error != nil {
|
|
return errors.New(envelope.Error.Message)
|
|
}
|
|
return errors.New("host.auth.list 失败")
|
|
}
|
|
var response struct {
|
|
Files []HostAuthFileEntry `json:"files"`
|
|
}
|
|
if err := json.Unmarshal(envelope.Result, &response); err != nil {
|
|
return fmt.Errorf("解析 host.auth.list 结果: %w", err)
|
|
}
|
|
now := time.Now().UTC()
|
|
accounts := make([]managedaccess.UpstreamAccount, 0, len(response.Files))
|
|
for _, entry := range response.Files {
|
|
if strings.TrimSpace(entry.ID) == "" {
|
|
continue
|
|
}
|
|
display := firstNonEmpty(entry.Label, entry.Email, entry.Account, entry.ProjectID, entry.Name, entry.ID)
|
|
accounts = append(accounts, managedaccess.UpstreamAccount{
|
|
CPAAuthID: entry.ID, CPAAuthIndex: entry.AuthIndex, Provider: firstNonEmpty(entry.Provider, entry.Type),
|
|
DisplayName: display, Status: entry.Status, StatusMessage: entry.StatusMessage,
|
|
Disabled: entry.Disabled, Unavailable: entry.Unavailable, Priority: entry.Priority, LastSeenAt: now,
|
|
})
|
|
}
|
|
return store.SyncUpstreamAccounts(ctx, accounts)
|
|
}
|
|
|
|
func firstNonEmpty(values ...string) string {
|
|
for _, value := range values {
|
|
if value = strings.TrimSpace(value); value != "" {
|
|
return value
|
|
}
|
|
}
|
|
return ""
|
|
}
|