cpa-plugin/internal/plugin/access.go

package plugin

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

	managedaccess "cpa-ext/internal/access"
)

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 req.RequestID != "" {
		a.pending.Store(req.RequestID, key.ID)
	}
	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 || account.Disabled || account.Unavailable {
		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 true
	}
	for _, allowed := range key.Models {
		if strings.EqualFold(strings.TrimSpace(allowed), model) {
			return true
		}
	}
	return false
}

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 ""
}