Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
1802 changed files with 503346 additions and 2 deletions
387
backend/internal/home/certificate.go
Normal file
387
backend/internal/home/certificate.go
Normal file
|
|
@ -0,0 +1,387 @@
|
|||
package home
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/rsa"
|
||||
"crypto/sha256"
|
||||
"crypto/x509"
|
||||
"crypto/x509/pkix"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"encoding/pem"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
)
|
||||
|
||||
const homeCertificateRequestTimeout = 30 * time.Second
|
||||
|
||||
type homeJWTClaims struct {
|
||||
CertificateID string `json:"certificate_id"`
|
||||
ClusterID string `json:"cluster_id"`
|
||||
CAFingerprint string `json:"ca_fingerprint"`
|
||||
EnrollmentSecret string `json:"enrollment_secret"`
|
||||
IP string `json:"ip"`
|
||||
Port int `json:"port"`
|
||||
IssuedAt int64 `json:"iat"`
|
||||
}
|
||||
|
||||
type certificateRequestResponse struct {
|
||||
OK bool `json:"ok"`
|
||||
Certificate string `json:"certificate"`
|
||||
CA string `json:"ca"`
|
||||
}
|
||||
|
||||
type certificatePaths struct {
|
||||
Dir string
|
||||
ClientCert string
|
||||
ClientKey string
|
||||
CACert string
|
||||
}
|
||||
|
||||
// ConfigFromJWT prepares a Home config from the JWT and ensures local mTLS files exist.
|
||||
func ConfigFromJWT(ctx context.Context, rawJWT string) (config.HomeConfig, error) {
|
||||
claims, errClaims := parseHomeJWTClaims(rawJWT)
|
||||
if errClaims != nil {
|
||||
return config.HomeConfig{}, errClaims
|
||||
}
|
||||
paths, errPaths := defaultCertificatePaths()
|
||||
if errPaths != nil {
|
||||
return config.HomeConfig{}, errPaths
|
||||
}
|
||||
if errEnsure := ensureHomeCertificateFiles(ctx, claims, paths); errEnsure != nil {
|
||||
return config.HomeConfig{}, errEnsure
|
||||
}
|
||||
return config.HomeConfig{
|
||||
Enabled: true,
|
||||
NodeID: strings.TrimSpace(claims.CertificateID),
|
||||
Host: strings.TrimSpace(claims.IP),
|
||||
Port: claims.Port,
|
||||
TLS: config.HomeTLSConfig{
|
||||
Enable: true,
|
||||
CACert: paths.CACert,
|
||||
ClientCert: paths.ClientCert,
|
||||
ClientKey: paths.ClientKey,
|
||||
UseTargetServerName: true,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func parseHomeJWTClaims(rawJWT string) (homeJWTClaims, error) {
|
||||
var claims homeJWTClaims
|
||||
parts := strings.Split(strings.TrimSpace(rawJWT), ".")
|
||||
if len(parts) != 3 {
|
||||
return claims, fmt.Errorf("home jwt is invalid")
|
||||
}
|
||||
payload, errDecode := decodeJWTPart(parts[1])
|
||||
if errDecode != nil {
|
||||
return claims, errDecode
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(payload, &claims); errUnmarshal != nil {
|
||||
return claims, errUnmarshal
|
||||
}
|
||||
if strings.TrimSpace(claims.CertificateID) == "" {
|
||||
return claims, fmt.Errorf("home jwt certificate_id is required")
|
||||
}
|
||||
if strings.TrimSpace(claims.ClusterID) == "" {
|
||||
return claims, fmt.Errorf("home jwt cluster_id is required")
|
||||
}
|
||||
if normalizeFingerprint(claims.CAFingerprint) == "" {
|
||||
return claims, fmt.Errorf("home jwt ca_fingerprint is required")
|
||||
}
|
||||
if strings.TrimSpace(claims.EnrollmentSecret) == "" {
|
||||
return claims, fmt.Errorf("home jwt enrollment_secret is required")
|
||||
}
|
||||
if strings.TrimSpace(claims.IP) == "" || claims.Port <= 0 {
|
||||
return claims, fmt.Errorf("home jwt target address is invalid")
|
||||
}
|
||||
return claims, nil
|
||||
}
|
||||
|
||||
func decodeJWTPart(part string) ([]byte, error) {
|
||||
if decoded, errDecode := base64.RawURLEncoding.DecodeString(part); errDecode == nil {
|
||||
return decoded, nil
|
||||
}
|
||||
return base64.URLEncoding.DecodeString(part)
|
||||
}
|
||||
|
||||
func defaultCertificatePaths() (certificatePaths, error) {
|
||||
homeDir, errHome := os.UserHomeDir()
|
||||
if errHome != nil {
|
||||
return certificatePaths{}, errHome
|
||||
}
|
||||
dir := filepath.Join(homeDir, ".cli-proxy-api")
|
||||
return certificatePaths{
|
||||
Dir: dir,
|
||||
ClientCert: filepath.Join(dir, "client-crt.pem"),
|
||||
ClientKey: filepath.Join(dir, "client-key.pem"),
|
||||
CACert: filepath.Join(dir, "home-ca-crt.pem"),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func ensureHomeCertificateFiles(ctx context.Context, claims homeJWTClaims, paths certificatePaths) error {
|
||||
if fileExists(paths.ClientCert) && fileExists(paths.ClientKey) {
|
||||
if !fileExists(paths.CACert) {
|
||||
return fmt.Errorf("home ca certificate file is missing")
|
||||
}
|
||||
if errVerify := verifyCACertificateFile(paths.CACert, claims.CAFingerprint); errVerify != nil {
|
||||
return errVerify
|
||||
}
|
||||
if errChmod := chmodCertificateFiles(paths); errChmod != nil {
|
||||
return errChmod
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if errMkdir := os.MkdirAll(paths.Dir, 0o700); errMkdir != nil {
|
||||
return errMkdir
|
||||
}
|
||||
key, errKey := loadOrCreateClientKey(paths.ClientKey)
|
||||
if errKey != nil {
|
||||
return errKey
|
||||
}
|
||||
csrPEM, errCSR := createClientCSR(claims.CertificateID, key)
|
||||
if errCSR != nil {
|
||||
return errCSR
|
||||
}
|
||||
response, errRequest := requestClientCertificate(ctx, claims, csrPEM)
|
||||
if errRequest != nil {
|
||||
return errRequest
|
||||
}
|
||||
if strings.TrimSpace(response.Certificate) == "" || strings.TrimSpace(response.CA) == "" {
|
||||
return fmt.Errorf("home certificate response is incomplete")
|
||||
}
|
||||
if errVerify := verifyCACertificatePEM([]byte(response.CA), claims.CAFingerprint); errVerify != nil {
|
||||
return errVerify
|
||||
}
|
||||
if errWrite := writeFile0600(paths.ClientCert, []byte(response.Certificate)); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
if errWrite := writeFile0600(paths.CACert, []byte(response.CA)); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func verifyCACertificateFile(path string, expectedFingerprint string) error {
|
||||
raw, errRead := os.ReadFile(path)
|
||||
if errRead != nil {
|
||||
return errRead
|
||||
}
|
||||
return verifyCACertificatePEM(raw, expectedFingerprint)
|
||||
}
|
||||
|
||||
func verifyCACertificatePEM(raw []byte, expectedFingerprint string) error {
|
||||
actual, errFingerprint := certificateFingerprintPEM(raw)
|
||||
if errFingerprint != nil {
|
||||
return errFingerprint
|
||||
}
|
||||
expected := normalizeFingerprint(expectedFingerprint)
|
||||
if expected == "" {
|
||||
return fmt.Errorf("home ca fingerprint is required")
|
||||
}
|
||||
if actual != expected {
|
||||
return fmt.Errorf("home ca fingerprint mismatch")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func certificateFingerprintPEM(raw []byte) (string, error) {
|
||||
block, _ := pem.Decode(raw)
|
||||
if block == nil || block.Type != "CERTIFICATE" {
|
||||
return "", fmt.Errorf("home ca certificate pem is invalid")
|
||||
}
|
||||
cert, errParse := x509.ParseCertificate(block.Bytes)
|
||||
if errParse != nil {
|
||||
return "", errParse
|
||||
}
|
||||
sum := sha256.Sum256(cert.Raw)
|
||||
return hex.EncodeToString(sum[:]), nil
|
||||
}
|
||||
|
||||
func normalizeFingerprint(fingerprint string) string {
|
||||
fingerprint = strings.TrimSpace(strings.ToLower(fingerprint))
|
||||
fingerprint = strings.ReplaceAll(fingerprint, ":", "")
|
||||
fingerprint = strings.ReplaceAll(fingerprint, " ", "")
|
||||
return fingerprint
|
||||
}
|
||||
|
||||
func loadOrCreateClientKey(path string) (*rsa.PrivateKey, error) {
|
||||
if fileExists(path) {
|
||||
raw, errRead := os.ReadFile(path)
|
||||
if errRead != nil {
|
||||
return nil, errRead
|
||||
}
|
||||
key, errParse := parseRSAPrivateKeyPEM(raw)
|
||||
if errParse != nil {
|
||||
return nil, errParse
|
||||
}
|
||||
if errChmod := os.Chmod(path, 0o600); errChmod != nil {
|
||||
return nil, errChmod
|
||||
}
|
||||
return key, nil
|
||||
}
|
||||
key, errKey := rsa.GenerateKey(rand.Reader, 2048)
|
||||
if errKey != nil {
|
||||
return nil, errKey
|
||||
}
|
||||
raw := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
|
||||
if errWrite := writeFile0600(path, raw); errWrite != nil {
|
||||
return nil, errWrite
|
||||
}
|
||||
return key, nil
|
||||
}
|
||||
|
||||
func writeFile0600(path string, raw []byte) error {
|
||||
if errWrite := os.WriteFile(path, raw, 0o600); errWrite != nil {
|
||||
return errWrite
|
||||
}
|
||||
return os.Chmod(path, 0o600)
|
||||
}
|
||||
|
||||
func chmodCertificateFiles(paths certificatePaths) error {
|
||||
for _, path := range []string{paths.ClientCert, paths.ClientKey, paths.CACert} {
|
||||
if errChmod := os.Chmod(path, 0o600); errChmod != nil {
|
||||
return errChmod
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func parseRSAPrivateKeyPEM(raw []byte) (*rsa.PrivateKey, error) {
|
||||
block, _ := pem.Decode(raw)
|
||||
if block == nil {
|
||||
return nil, fmt.Errorf("client key pem is invalid")
|
||||
}
|
||||
switch block.Type {
|
||||
case "RSA PRIVATE KEY":
|
||||
return x509.ParsePKCS1PrivateKey(block.Bytes)
|
||||
case "PRIVATE KEY":
|
||||
key, errParse := x509.ParsePKCS8PrivateKey(block.Bytes)
|
||||
if errParse != nil {
|
||||
return nil, errParse
|
||||
}
|
||||
rsaKey, ok := key.(*rsa.PrivateKey)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("client key is not rsa")
|
||||
}
|
||||
return rsaKey, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("client key pem type %q is unsupported", block.Type)
|
||||
}
|
||||
}
|
||||
|
||||
func createClientCSR(certificateID string, key *rsa.PrivateKey) ([]byte, error) {
|
||||
certificateID = strings.TrimSpace(certificateID)
|
||||
if certificateID == "" {
|
||||
return nil, fmt.Errorf("certificate id is required")
|
||||
}
|
||||
template := &x509.CertificateRequest{
|
||||
Subject: pkix.Name{
|
||||
CommonName: certificateID,
|
||||
},
|
||||
}
|
||||
der, errCreate := x509.CreateCertificateRequest(rand.Reader, template, key)
|
||||
if errCreate != nil {
|
||||
return nil, errCreate
|
||||
}
|
||||
return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: der}), nil
|
||||
}
|
||||
|
||||
func requestClientCertificate(ctx context.Context, claims homeJWTClaims, csrPEM []byte) (certificateRequestResponse, error) {
|
||||
var response certificateRequestResponse
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
dialCtx, cancel := context.WithTimeout(ctx, homeCertificateRequestTimeout)
|
||||
defer cancel()
|
||||
addr := net.JoinHostPort(strings.TrimSpace(claims.IP), strconv.Itoa(claims.Port))
|
||||
conn, errDial := (&net.Dialer{}).DialContext(dialCtx, "tcp", addr)
|
||||
if errDial != nil {
|
||||
return response, errDial
|
||||
}
|
||||
defer func() {
|
||||
_ = conn.Close()
|
||||
}()
|
||||
if deadline, ok := dialCtx.Deadline(); ok {
|
||||
_ = conn.SetDeadline(deadline)
|
||||
}
|
||||
if _, errWrite := conn.Write(encodeRESPArray("CERTIFICATE", "REQUEST", claims.CertificateID, claims.EnrollmentSecret, string(csrPEM))); errWrite != nil {
|
||||
return response, errWrite
|
||||
}
|
||||
raw, errRead := readRESPBulk(bufio.NewReader(conn))
|
||||
if errRead != nil {
|
||||
return response, errRead
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(raw, &response); errUnmarshal != nil {
|
||||
return response, errUnmarshal
|
||||
}
|
||||
if !response.OK {
|
||||
return response, fmt.Errorf("home certificate request failed")
|
||||
}
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func encodeRESPArray(args ...string) []byte {
|
||||
var buf bytes.Buffer
|
||||
buf.WriteString("*")
|
||||
buf.WriteString(strconv.Itoa(len(args)))
|
||||
buf.WriteString("\r\n")
|
||||
for _, arg := range args {
|
||||
buf.WriteString("$")
|
||||
buf.WriteString(strconv.Itoa(len(arg)))
|
||||
buf.WriteString("\r\n")
|
||||
buf.WriteString(arg)
|
||||
buf.WriteString("\r\n")
|
||||
}
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func readRESPBulk(reader *bufio.Reader) ([]byte, error) {
|
||||
prefix, errRead := reader.ReadByte()
|
||||
if errRead != nil {
|
||||
return nil, errRead
|
||||
}
|
||||
switch prefix {
|
||||
case '$':
|
||||
line, errLine := reader.ReadString('\n')
|
||||
if errLine != nil {
|
||||
return nil, errLine
|
||||
}
|
||||
size, errSize := strconv.Atoi(strings.TrimSpace(line))
|
||||
if errSize != nil {
|
||||
return nil, errSize
|
||||
}
|
||||
if size < 0 {
|
||||
return nil, fmt.Errorf("home certificate request returned nil")
|
||||
}
|
||||
payload := make([]byte, size+2)
|
||||
if _, errFull := io.ReadFull(reader, payload); errFull != nil {
|
||||
return nil, errFull
|
||||
}
|
||||
return payload[:size], nil
|
||||
case '-':
|
||||
line, errLine := reader.ReadString('\n')
|
||||
if errLine != nil {
|
||||
return nil, errLine
|
||||
}
|
||||
return nil, fmt.Errorf("%s", strings.TrimSpace(line))
|
||||
default:
|
||||
return nil, fmt.Errorf("home certificate request returned unsupported resp prefix %q", prefix)
|
||||
}
|
||||
}
|
||||
|
||||
func fileExists(path string) bool {
|
||||
info, errStat := os.Stat(path)
|
||||
return errStat == nil && !info.IsDir()
|
||||
}
|
||||
2010
backend/internal/home/client.go
Normal file
2010
backend/internal/home/client.go
Normal file
File diff suppressed because it is too large
Load diff
2239
backend/internal/home/client_test.go
Normal file
2239
backend/internal/home/client_test.go
Normal file
File diff suppressed because it is too large
Load diff
287
backend/internal/home/concurrency_release.go
Normal file
287
backend/internal/home/concurrency_release.go
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
package home
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
|
||||
)
|
||||
|
||||
// ConcurrencyReleaseFrame is the cumulative release accepted by Home for one credential and model.
|
||||
type ConcurrencyReleaseFrame struct {
|
||||
CredentialID string `json:"credential_id"`
|
||||
Model string `json:"model"`
|
||||
ReleaseSeq int64 `json:"release_seq"`
|
||||
}
|
||||
|
||||
type releaseState struct {
|
||||
Latest int64
|
||||
Acked int64
|
||||
waiters map[int64][]chan struct{}
|
||||
}
|
||||
|
||||
type releaseFlusher struct {
|
||||
mu sync.Mutex
|
||||
groups map[executionregistry.ReleaseGroup]releaseState
|
||||
flushInterval time.Duration
|
||||
maxBackoff time.Duration
|
||||
configProvider func() internalconfig.CredentialConcurrencyConfig
|
||||
send func(context.Context, ConcurrencyReleaseFrame) error
|
||||
wake chan struct{}
|
||||
force chan context.Context
|
||||
}
|
||||
|
||||
func newReleaseFlusher(flushInterval, maxBackoff time.Duration, send func(context.Context, ConcurrencyReleaseFrame) error) *releaseFlusher {
|
||||
return &releaseFlusher{
|
||||
groups: make(map[executionregistry.ReleaseGroup]releaseState),
|
||||
flushInterval: flushInterval,
|
||||
maxBackoff: maxBackoff,
|
||||
send: send,
|
||||
wake: make(chan struct{}, 1),
|
||||
force: make(chan context.Context, 1),
|
||||
}
|
||||
}
|
||||
|
||||
// NewReleaseFlusher creates a flusher that reads timing updates from the current limiter configuration.
|
||||
func NewReleaseFlusher(configProvider func() internalconfig.CredentialConcurrencyConfig, send func(context.Context, ConcurrencyReleaseFrame) error) *releaseFlusher {
|
||||
flusher := newReleaseFlusher(0, 0, send)
|
||||
flusher.SetConfigProvider(configProvider)
|
||||
return flusher
|
||||
}
|
||||
|
||||
func (f *releaseFlusher) SetConfigProvider(provider func() internalconfig.CredentialConcurrencyConfig) {
|
||||
if f == nil {
|
||||
return
|
||||
}
|
||||
f.mu.Lock()
|
||||
f.configProvider = provider
|
||||
f.mu.Unlock()
|
||||
f.signal()
|
||||
}
|
||||
|
||||
// SetSender replaces the Home lifetime used for subsequent release attempts.
|
||||
func (f *releaseFlusher) SetSender(send func(context.Context, ConcurrencyReleaseFrame) error) {
|
||||
if f == nil {
|
||||
return
|
||||
}
|
||||
f.mu.Lock()
|
||||
f.send = send
|
||||
f.mu.Unlock()
|
||||
f.signal()
|
||||
}
|
||||
|
||||
// MarkDirty records the latest cumulative sequence for one release group and
|
||||
// returns a ticket completed when Home acknowledges that sequence.
|
||||
func (f *releaseFlusher) MarkDirty(group executionregistry.ReleaseGroup, sequence int64) *executionregistry.ReleaseTicket {
|
||||
if f == nil || sequence <= 0 || group.CredentialID == "" || group.Model == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
f.mu.Lock()
|
||||
state := f.groups[group]
|
||||
if sequence <= state.Acked {
|
||||
close(done)
|
||||
} else {
|
||||
if state.waiters == nil {
|
||||
state.waiters = make(map[int64][]chan struct{})
|
||||
}
|
||||
state.waiters[sequence] = append(state.waiters[sequence], done)
|
||||
if sequence > state.Latest {
|
||||
state.Latest = sequence
|
||||
}
|
||||
f.groups[group] = state
|
||||
}
|
||||
f.mu.Unlock()
|
||||
f.signal()
|
||||
return executionregistry.NewReleaseTicket(group, sequence, done)
|
||||
}
|
||||
|
||||
// Run sends dirty groups until its lifetime is cancelled.
|
||||
func (f *releaseFlusher) Run(ctx context.Context) {
|
||||
if f == nil {
|
||||
return
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
|
||||
timer := time.NewTimer(0)
|
||||
defer timer.Stop()
|
||||
delay := f.timings().flushInterval
|
||||
backingOff := false
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-f.wake:
|
||||
if !backingOff {
|
||||
resetReleaseTimer(timer, 0)
|
||||
}
|
||||
case forceCtx := <-f.force:
|
||||
resetReleaseTimer(timer, 0)
|
||||
failed := f.flush(forceCtx)
|
||||
delay, backingOff = f.nextDelay(delay, failed)
|
||||
resetReleaseTimer(timer, delay)
|
||||
case <-timer.C:
|
||||
failed := f.flush(ctx)
|
||||
delay, backingOff = f.nextDelay(delay, failed)
|
||||
timer.Reset(delay)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (f *releaseFlusher) nextDelay(delay time.Duration, failed bool) (time.Duration, bool) {
|
||||
timings := f.timings()
|
||||
if !failed {
|
||||
return timings.flushInterval, false
|
||||
}
|
||||
delay *= 2
|
||||
if delay < timings.flushInterval {
|
||||
delay = timings.flushInterval
|
||||
}
|
||||
if delay > timings.maxBackoff {
|
||||
delay = timings.maxBackoff
|
||||
}
|
||||
return delay, true
|
||||
}
|
||||
|
||||
func resetReleaseTimer(timer *time.Timer, delay time.Duration) {
|
||||
if !timer.Stop() {
|
||||
select {
|
||||
case <-timer.C:
|
||||
default:
|
||||
}
|
||||
}
|
||||
timer.Reset(delay)
|
||||
}
|
||||
|
||||
type releaseFlusherTimings struct {
|
||||
flushInterval time.Duration
|
||||
maxBackoff time.Duration
|
||||
}
|
||||
|
||||
func (f *releaseFlusher) timings() releaseFlusherTimings {
|
||||
defaults := internalconfig.CredentialConcurrencyConfig{}.WithDefaults()
|
||||
timings := releaseFlusherTimings{flushInterval: f.flushInterval, maxBackoff: f.maxBackoff}
|
||||
|
||||
f.mu.Lock()
|
||||
provider := f.configProvider
|
||||
f.mu.Unlock()
|
||||
if provider != nil {
|
||||
cfg := provider().WithDefaults()
|
||||
timings.flushInterval = cfg.ReleaseFlushInterval
|
||||
timings.maxBackoff = cfg.ReleaseMaxBackoff
|
||||
}
|
||||
if timings.flushInterval <= 0 {
|
||||
timings.flushInterval = defaults.ReleaseFlushInterval
|
||||
}
|
||||
if timings.maxBackoff < timings.flushInterval {
|
||||
timings.maxBackoff = timings.flushInterval
|
||||
}
|
||||
return timings
|
||||
}
|
||||
|
||||
func (f *releaseFlusher) flush(ctx context.Context) bool {
|
||||
if f == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
f.mu.Lock()
|
||||
send := f.send
|
||||
pending := make(map[executionregistry.ReleaseGroup]int64, len(f.groups))
|
||||
for group, state := range f.groups {
|
||||
if state.Latest > state.Acked {
|
||||
pending[group] = state.Latest
|
||||
}
|
||||
}
|
||||
f.mu.Unlock()
|
||||
if send == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
failed := false
|
||||
for group, sequence := range pending {
|
||||
errSend := send(ctx, ConcurrencyReleaseFrame{
|
||||
CredentialID: group.CredentialID,
|
||||
Model: group.Model,
|
||||
ReleaseSeq: sequence,
|
||||
})
|
||||
if errSend != nil {
|
||||
failed = true
|
||||
continue
|
||||
}
|
||||
f.mu.Lock()
|
||||
state := f.groups[group]
|
||||
if sequence > state.Acked {
|
||||
state.Acked = sequence
|
||||
for waiterSequence, waiters := range state.waiters {
|
||||
if waiterSequence <= state.Acked {
|
||||
for _, done := range waiters {
|
||||
close(done)
|
||||
}
|
||||
delete(state.waiters, waiterSequence)
|
||||
}
|
||||
}
|
||||
}
|
||||
f.groups[group] = state
|
||||
f.mu.Unlock()
|
||||
}
|
||||
return failed
|
||||
}
|
||||
|
||||
// Flush waits for all currently dirty groups to be acknowledged within ctx.
|
||||
func (f *releaseFlusher) Flush(ctx context.Context) error {
|
||||
if f == nil {
|
||||
return nil
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
f.forceFlush(ctx)
|
||||
ticker := time.NewTicker(time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
if f.idle() {
|
||||
return nil
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (f *releaseFlusher) idle() bool {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
for _, state := range f.groups {
|
||||
if state.Latest > state.Acked {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (f *releaseFlusher) signal() {
|
||||
if f == nil {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case f.wake <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (f *releaseFlusher) forceFlush(ctx context.Context) {
|
||||
if f == nil {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case f.force <- ctx:
|
||||
default:
|
||||
}
|
||||
}
|
||||
505
backend/internal/home/concurrency_release_test.go
Normal file
505
backend/internal/home/concurrency_release_test.go
Normal file
|
|
@ -0,0 +1,505 @@
|
|||
package home
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
|
||||
)
|
||||
|
||||
func concurrencyReleaseFrameFromFixture(t *testing.T) ConcurrencyReleaseFrame {
|
||||
t.Helper()
|
||||
raw, errRead := os.ReadFile(filepath.Join("testdata", "concurrency_release.json"))
|
||||
if errRead != nil {
|
||||
t.Fatal(errRead)
|
||||
}
|
||||
|
||||
var frame ConcurrencyReleaseFrame
|
||||
if errUnmarshal := json.Unmarshal(raw, &frame); errUnmarshal != nil {
|
||||
t.Fatal(errUnmarshal)
|
||||
}
|
||||
return frame
|
||||
}
|
||||
|
||||
func TestConcurrencyReleaseFrameFixture(t *testing.T) {
|
||||
raw, errRead := os.ReadFile(filepath.Join("testdata", "concurrency_release.json"))
|
||||
if errRead != nil {
|
||||
t.Fatal(errRead)
|
||||
}
|
||||
frame := concurrencyReleaseFrameFromFixture(t)
|
||||
if frame != (ConcurrencyReleaseFrame{CredentialID: "cred-1", Model: "gpt", ReleaseSeq: 1}) {
|
||||
t.Fatalf("fixture frame = %#v", frame)
|
||||
}
|
||||
marshaled, errMarshal := json.Marshal(frame)
|
||||
if errMarshal != nil {
|
||||
t.Fatal(errMarshal)
|
||||
}
|
||||
if !bytes.Equal(marshaled, bytes.TrimSpace(raw)) {
|
||||
t.Fatalf("marshaled frame = %q, want fixture %q", marshaled, bytes.TrimSpace(raw))
|
||||
}
|
||||
}
|
||||
|
||||
type recordingReleaseSender struct {
|
||||
mu sync.Mutex
|
||||
failures int
|
||||
frames []ConcurrencyReleaseFrame
|
||||
acked []ConcurrencyReleaseFrame
|
||||
sent chan struct{}
|
||||
}
|
||||
|
||||
func (s *recordingReleaseSender) Send(_ context.Context, frame ConcurrencyReleaseFrame) error {
|
||||
s.mu.Lock()
|
||||
s.frames = append(s.frames, frame)
|
||||
failed := s.failures > 0
|
||||
if failed {
|
||||
s.failures--
|
||||
} else {
|
||||
s.acked = append(s.acked, frame)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
select {
|
||||
case s.sent <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
if failed {
|
||||
return errors.New("temporary Home failure")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *recordingReleaseSender) LastSequence() int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if len(s.acked) == 0 {
|
||||
return 0
|
||||
}
|
||||
return s.acked[len(s.acked)-1].ReleaseSeq
|
||||
}
|
||||
|
||||
func (s *recordingReleaseSender) WaitForSequence(sequence int64, timeout time.Duration) bool {
|
||||
timer := time.NewTimer(timeout)
|
||||
defer timer.Stop()
|
||||
for {
|
||||
if s.LastSequence() == sequence {
|
||||
return true
|
||||
}
|
||||
select {
|
||||
case <-timer.C:
|
||||
return false
|
||||
case <-s.sent:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseFlusherRetriesLatestCumulativeSequence(t *testing.T) {
|
||||
sender := &recordingReleaseSender{failures: 1, sent: make(chan struct{}, 8)}
|
||||
flusher := newReleaseFlusher(10*time.Millisecond, 40*time.Millisecond, sender.Send)
|
||||
group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"}
|
||||
flusher.MarkDirty(group, 1)
|
||||
flusher.MarkDirty(group, 3)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
go flusher.Run(ctx)
|
||||
|
||||
if !sender.WaitForSequence(3, 500*time.Millisecond) {
|
||||
t.Fatalf("last sequence = %d, want 3", sender.LastSequence())
|
||||
}
|
||||
if sender.LastSequence() != 3 {
|
||||
t.Fatalf("last sequence = %d, want 3", sender.LastSequence())
|
||||
}
|
||||
}
|
||||
|
||||
type blockingReleaseSender struct {
|
||||
started chan struct{}
|
||||
release chan struct{}
|
||||
frames chan ConcurrencyReleaseFrame
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (s *blockingReleaseSender) Send(_ context.Context, frame ConcurrencyReleaseFrame) error {
|
||||
s.once.Do(func() { close(s.started) })
|
||||
select {
|
||||
case s.frames <- frame:
|
||||
default:
|
||||
}
|
||||
<-s.release
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestReleaseFlusherDoesNotLoseASequenceMarkedDuringSend(t *testing.T) {
|
||||
sender := &blockingReleaseSender{
|
||||
started: make(chan struct{}),
|
||||
release: make(chan struct{}),
|
||||
frames: make(chan ConcurrencyReleaseFrame, 4),
|
||||
}
|
||||
flusher := newReleaseFlusher(time.Millisecond, 10*time.Millisecond, sender.Send)
|
||||
group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"}
|
||||
flusher.MarkDirty(group, 1)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
go flusher.Run(ctx)
|
||||
|
||||
select {
|
||||
case <-sender.started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("release flusher did not begin sending")
|
||||
}
|
||||
flusher.MarkDirty(group, 2)
|
||||
close(sender.release)
|
||||
|
||||
deadline := time.NewTimer(time.Second)
|
||||
defer deadline.Stop()
|
||||
for {
|
||||
select {
|
||||
case frame := <-sender.frames:
|
||||
if frame.ReleaseSeq == 2 {
|
||||
return
|
||||
}
|
||||
case <-deadline.C:
|
||||
t.Fatal("release flusher did not send the latest sequence")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseFlusherUsesCurrentLimiterConfig(t *testing.T) {
|
||||
flusher := newReleaseFlusher(time.Hour, 2*time.Hour, func(context.Context, ConcurrencyReleaseFrame) error { return nil })
|
||||
flusher.SetConfigProvider(func() internalconfig.CredentialConcurrencyConfig {
|
||||
return internalconfig.CredentialConcurrencyConfig{
|
||||
ReleaseFlushInterval: 5 * time.Millisecond,
|
||||
ReleaseMaxBackoff: 25 * time.Millisecond,
|
||||
}
|
||||
})
|
||||
if got := flusher.timings(); got.flushInterval != 5*time.Millisecond || got.maxBackoff != 25*time.Millisecond {
|
||||
t.Fatalf("timings = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseFlusherStopsWithLifetime(t *testing.T) {
|
||||
sender := &recordingReleaseSender{sent: make(chan struct{}, 1)}
|
||||
flusher := newReleaseFlusher(time.Hour, time.Hour, sender.Send)
|
||||
done := make(chan struct{})
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go func() {
|
||||
defer close(done)
|
||||
flusher.Run(ctx)
|
||||
}()
|
||||
cancel()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("release flusher did not stop with its lifetime")
|
||||
}
|
||||
}
|
||||
|
||||
type timedReleaseAttempt struct {
|
||||
at time.Time
|
||||
frame ConcurrencyReleaseFrame
|
||||
failed bool
|
||||
}
|
||||
|
||||
type outageReleaseSender struct {
|
||||
mu sync.Mutex
|
||||
outage bool
|
||||
attempts []timedReleaseAttempt
|
||||
sent chan struct{}
|
||||
}
|
||||
|
||||
func (s *outageReleaseSender) Send(_ context.Context, frame ConcurrencyReleaseFrame) error {
|
||||
s.mu.Lock()
|
||||
failed := s.outage
|
||||
s.attempts = append(s.attempts, timedReleaseAttempt{at: time.Now(), frame: frame, failed: failed})
|
||||
s.mu.Unlock()
|
||||
select {
|
||||
case s.sent <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
if failed {
|
||||
return errors.New("temporary Home outage")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *outageReleaseSender) SetOutage(outage bool) {
|
||||
s.mu.Lock()
|
||||
s.outage = outage
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *outageReleaseSender) WaitForAttempts(count int, timeout time.Duration) []timedReleaseAttempt {
|
||||
timer := time.NewTimer(timeout)
|
||||
defer timer.Stop()
|
||||
for {
|
||||
s.mu.Lock()
|
||||
attempts := append([]timedReleaseAttempt(nil), s.attempts...)
|
||||
s.mu.Unlock()
|
||||
if len(attempts) >= count {
|
||||
return attempts
|
||||
}
|
||||
select {
|
||||
case <-timer.C:
|
||||
return attempts
|
||||
case <-s.sent:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseFlusherCoalescesDirtyWakesDuringFailureBackoff(t *testing.T) {
|
||||
const (
|
||||
flushInterval = 20 * time.Millisecond
|
||||
maxBackoff = 80 * time.Millisecond
|
||||
tolerance = 10 * time.Millisecond
|
||||
)
|
||||
|
||||
sender := &outageReleaseSender{outage: true, sent: make(chan struct{}, 32)}
|
||||
flusher := newReleaseFlusher(flushInterval, maxBackoff, sender.Send)
|
||||
group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"}
|
||||
flusher.MarkDirty(group, 1)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
flusher.Run(ctx)
|
||||
}()
|
||||
defer func() {
|
||||
cancel()
|
||||
<-done
|
||||
}()
|
||||
|
||||
stopReleases := make(chan struct{})
|
||||
producerDone := make(chan struct{})
|
||||
latest := int64(1)
|
||||
go func() {
|
||||
defer close(producerDone)
|
||||
ticker := time.NewTicker(time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stopReleases:
|
||||
return
|
||||
case <-ticker.C:
|
||||
latest++
|
||||
flusher.MarkDirty(group, latest)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
attempts := sender.WaitForAttempts(3, time.Second)
|
||||
close(stopReleases)
|
||||
<-producerDone
|
||||
if len(attempts) < 3 {
|
||||
t.Fatalf("attempt count = %d, want at least 3", len(attempts))
|
||||
}
|
||||
for _, attempt := range attempts[:3] {
|
||||
if !attempt.failed {
|
||||
t.Fatal("release unexpectedly succeeded during outage")
|
||||
}
|
||||
}
|
||||
if got := attempts[1].at.Sub(attempts[0].at); got < 2*flushInterval-tolerance {
|
||||
t.Fatalf("first retry delay = %s, want at least %s", got, 2*flushInterval-tolerance)
|
||||
}
|
||||
if got := attempts[2].at.Sub(attempts[1].at); got < maxBackoff-tolerance {
|
||||
t.Fatalf("second retry delay = %s, want at least %s", got, maxBackoff-tolerance)
|
||||
}
|
||||
|
||||
latest++
|
||||
recoverySequence := latest
|
||||
recoveryStart := attempts[2].at
|
||||
sender.SetOutage(false)
|
||||
flusher.MarkDirty(group, recoverySequence)
|
||||
|
||||
attempts = sender.WaitForAttempts(4, time.Second)
|
||||
if len(attempts) < 4 {
|
||||
t.Fatalf("attempt count after recovery = %d, want at least 4", len(attempts))
|
||||
}
|
||||
recovered := attempts[3]
|
||||
if recovered.failed || recovered.frame.ReleaseSeq != recoverySequence {
|
||||
t.Fatalf("recovery attempt = %#v, want successful sequence %d", recovered, recoverySequence)
|
||||
}
|
||||
if got := recovered.at.Sub(recoveryStart); got < maxBackoff-tolerance {
|
||||
t.Fatalf("recovery retry delay = %s, want at least %s", got, maxBackoff-tolerance)
|
||||
}
|
||||
}
|
||||
|
||||
type boundedForceReleaseSender struct {
|
||||
attempts chan context.Context
|
||||
calls int
|
||||
}
|
||||
|
||||
func (s *boundedForceReleaseSender) Send(ctx context.Context, _ ConcurrencyReleaseFrame) error {
|
||||
s.calls++
|
||||
select {
|
||||
case s.attempts <- ctx:
|
||||
default:
|
||||
}
|
||||
if s.calls == 1 {
|
||||
return errors.New("temporary Home failure")
|
||||
}
|
||||
<-ctx.Done()
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
func TestReleaseFlusherFlushForceUsesBoundedContext(t *testing.T) {
|
||||
sender := &boundedForceReleaseSender{attempts: make(chan context.Context, 2)}
|
||||
flusher := newReleaseFlusher(time.Second, time.Second, sender.Send)
|
||||
group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"}
|
||||
flusher.MarkDirty(group, 1)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
flusher.Run(ctx)
|
||||
}()
|
||||
defer func() {
|
||||
cancel()
|
||||
<-done
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-sender.attempts:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("release flusher did not make the initial failed attempt")
|
||||
}
|
||||
|
||||
flushCtx, cancelFlush := context.WithTimeout(context.Background(), 40*time.Millisecond)
|
||||
defer cancelFlush()
|
||||
if errFlush := flusher.Flush(flushCtx); !errors.Is(errFlush, context.DeadlineExceeded) {
|
||||
t.Fatalf("Flush() error = %v, want deadline exceeded", errFlush)
|
||||
}
|
||||
|
||||
select {
|
||||
case forceCtx := <-sender.attempts:
|
||||
if _, ok := forceCtx.Deadline(); !ok {
|
||||
t.Fatal("forced release attempt did not receive the bounded Flush context")
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Flush() did not bypass the normal retry interval")
|
||||
}
|
||||
}
|
||||
|
||||
func TestScopeEndBlocksDrainUntilReleaseSinkFlushesFinalSequence(t *testing.T) {
|
||||
sender := &recordingReleaseSender{sent: make(chan struct{}, 2)}
|
||||
flusher := newReleaseFlusher(time.Hour, time.Hour, sender.Send)
|
||||
releaseCtx, cancelRelease := context.WithCancel(context.Background())
|
||||
releaseDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(releaseDone)
|
||||
flusher.Run(releaseCtx)
|
||||
}()
|
||||
defer func() {
|
||||
cancelRelease()
|
||||
<-releaseDone
|
||||
}()
|
||||
|
||||
registry := executionregistry.New()
|
||||
sinkStarted := make(chan struct{})
|
||||
unblockSink := make(chan struct{})
|
||||
registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, sequence int64) {
|
||||
close(sinkStarted)
|
||||
<-unblockSink
|
||||
flusher.MarkDirty(group, sequence)
|
||||
})
|
||||
pending, errBegin := registry.BeginDispatch()
|
||||
if errBegin != nil {
|
||||
t.Fatal(errBegin)
|
||||
}
|
||||
scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{CredentialID: "cred-1", Model: "gpt", Accounted: true})
|
||||
if errInstall != nil {
|
||||
t.Fatal(errInstall)
|
||||
}
|
||||
|
||||
endDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(endDone)
|
||||
scope.End("complete")
|
||||
}()
|
||||
select {
|
||||
case <-sinkStarted:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Scope.End() did not call the release sink")
|
||||
}
|
||||
|
||||
drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancelDrain()
|
||||
drainDone := make(chan error, 1)
|
||||
go func() { drainDone <- registry.Drain(drainCtx) }()
|
||||
|
||||
select {
|
||||
case errDrain := <-drainDone:
|
||||
t.Fatalf("Drain() returned before the release sink completed: %v", errDrain)
|
||||
case <-time.After(20 * time.Millisecond):
|
||||
}
|
||||
|
||||
mutexAvailable := make(chan struct{})
|
||||
go func() {
|
||||
registry.SetReleaseSink(nil)
|
||||
close(mutexAvailable)
|
||||
}()
|
||||
select {
|
||||
case <-mutexAvailable:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("release sink blocked the registry mutex")
|
||||
}
|
||||
if _, errBegin := registry.BeginDispatch(); !errors.Is(errBegin, executionregistry.ErrRegistryNotAccepting) {
|
||||
t.Fatalf("BeginDispatch() error = %v, want ErrRegistryNotAccepting", errBegin)
|
||||
}
|
||||
|
||||
close(unblockSink)
|
||||
select {
|
||||
case <-endDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Scope.End() did not complete after the release sink unblocked")
|
||||
}
|
||||
if errDrain := <-drainDone; errDrain != nil {
|
||||
t.Fatalf("Drain() error = %v", errDrain)
|
||||
}
|
||||
|
||||
flushCtx, cancelFlush := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancelFlush()
|
||||
if errFlush := flusher.Flush(flushCtx); errFlush != nil {
|
||||
t.Fatalf("Flush() error = %v", errFlush)
|
||||
}
|
||||
if got := sender.LastSequence(); got != 1 {
|
||||
t.Fatalf("final flushed sequence = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseFlusherSenderReplacementPreservesTicket(t *testing.T) {
|
||||
flusher := newReleaseFlusher(time.Hour, time.Hour, func(context.Context, ConcurrencyReleaseFrame) error {
|
||||
return errors.New("old Home unavailable")
|
||||
})
|
||||
group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"}
|
||||
ticket := flusher.MarkDirty(group, 1)
|
||||
if ticket == nil {
|
||||
t.Fatal("MarkDirty() ticket = nil")
|
||||
}
|
||||
if failed := flusher.flush(context.Background()); !failed {
|
||||
t.Fatal("old sender release attempt did not fail")
|
||||
}
|
||||
|
||||
flusher.SetSender(func(_ context.Context, frame ConcurrencyReleaseFrame) error {
|
||||
if frame.CredentialID != group.CredentialID || frame.Model != group.Model || frame.ReleaseSeq != 1 {
|
||||
t.Fatalf("replacement sender frame = %#v", frame)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if failed := flusher.flush(context.Background()); failed {
|
||||
t.Fatal("replacement sender release attempt failed")
|
||||
}
|
||||
waitCtx, cancelWait := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancelWait()
|
||||
if errWait := ticket.Wait(waitCtx); errWait != nil {
|
||||
t.Fatalf("ticket did not survive sender replacement: %v", errWait)
|
||||
}
|
||||
}
|
||||
27
backend/internal/home/global.go
Normal file
27
backend/internal/home/global.go
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
package home
|
||||
|
||||
import "sync/atomic"
|
||||
|
||||
var currentClient atomic.Pointer[Client]
|
||||
|
||||
// SetCurrent sets the active home client used by runtime integrations.
|
||||
func SetCurrent(client *Client) {
|
||||
currentClient.Store(client)
|
||||
}
|
||||
|
||||
// Current returns the active home client instance, if any.
|
||||
func Current() *Client {
|
||||
return currentClient.Load()
|
||||
}
|
||||
|
||||
// ClearCurrent removes the active home client.
|
||||
func ClearCurrent() {
|
||||
currentClient.Store(nil)
|
||||
}
|
||||
|
||||
// ClearCurrentIf removes the active client only when it is client.
|
||||
func ClearCurrentIf(client *Client) {
|
||||
if client != nil {
|
||||
currentClient.CompareAndSwap(client, nil)
|
||||
}
|
||||
}
|
||||
182
backend/internal/home/in_flight_contract_test.go
Normal file
182
backend/internal/home/in_flight_contract_test.go
Normal file
|
|
@ -0,0 +1,182 @@
|
|||
package home
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCredentialInFlightWireContractFixture(t *testing.T) {
|
||||
raw, errRead := os.ReadFile(filepath.Join("testdata", "credential_in_flight_contract.json"))
|
||||
if errRead != nil {
|
||||
t.Fatalf("ReadFile() error = %v", errRead)
|
||||
}
|
||||
fixture, errDecode := decodeInFlightContractFixture(raw)
|
||||
if errDecode != nil {
|
||||
t.Fatalf("decodeInFlightContractFixture() error = %v", errDecode)
|
||||
}
|
||||
if fixture.Part.Kind != InFlightFramePart || fixture.Part.PartIndex == nil || *fixture.Part.PartIndex != 0 || fixture.Part.PartCount == nil || *fixture.Part.PartCount != 1 {
|
||||
t.Fatalf("part = %#v", fixture.Part)
|
||||
}
|
||||
if fixture.Part.Aggregates[0].Status != InFlightAccounted || fixture.Part.Aggregates[1].Status != InFlightUnaccounted {
|
||||
t.Fatalf("statuses = %#v", fixture.Part.Aggregates)
|
||||
}
|
||||
if fixture.Overflow.Kind != InFlightFrameOverflow || fixture.Overflow.AggregateGroupCount != 100001 {
|
||||
t.Fatalf("overflow = %#v", fixture.Overflow)
|
||||
}
|
||||
assertInFlightContractFields(t)
|
||||
assertRequiredInFlightJSONKeys(t, raw, []string{"config", "part", "overflow"})
|
||||
assertInFlightFixtureKeys(t, fixture)
|
||||
}
|
||||
|
||||
func TestCredentialInFlightWireContractRejectsInvalidJSON(t *testing.T) {
|
||||
raw, errRead := os.ReadFile(filepath.Join("testdata", "credential_in_flight_contract.json"))
|
||||
if errRead != nil {
|
||||
t.Fatalf("ReadFile() error = %v", errRead)
|
||||
}
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
raw []byte
|
||||
}{
|
||||
{name: "unknown frame owner field", raw: bytes.Replace(raw, []byte(`"kind": "part"`), []byte(`"kind": "part", "node_id": "node-a"`), 1)},
|
||||
{name: "unknown aggregate owner field", raw: bytes.Replace(raw, []byte(`"credential_id": "cred-a"`), []byte(`"credential_id": "cred-a", "fingerprint": "owner"`), 1)},
|
||||
{name: "unknown detail secret field", raw: bytes.Replace(raw, []byte(`"request_id": "req-1"`), []byte(`"request_id": "req-1", "secret": "secret"`), 1)},
|
||||
{name: "unknown overflow secret field", raw: bytes.Replace(raw, []byte(`"aggregate_group_count": 100001`), []byte(`"aggregate_group_count": 100001, "api_key": "secret"`), 1)},
|
||||
{name: "trailing JSON", raw: append(append([]byte{}, raw...), []byte(` {"part": {}}`)...)},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if _, errDecode := decodeInFlightContractFixture(test.raw); errDecode == nil {
|
||||
t.Fatal("decodeInFlightContractFixture() error = nil")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type inFlightContractFixture struct {
|
||||
Part InFlightSnapshotFrame
|
||||
Overflow InFlightSnapshotFrame
|
||||
PartJSON json.RawMessage
|
||||
OverflowJSON json.RawMessage
|
||||
}
|
||||
|
||||
func decodeInFlightContractFixture(raw []byte) (inFlightContractFixture, error) {
|
||||
var fixture inFlightContractFixture
|
||||
var document struct {
|
||||
Config json.RawMessage `json:"config"`
|
||||
Part InFlightSnapshotFrame `json:"part"`
|
||||
Overflow InFlightSnapshotFrame `json:"overflow"`
|
||||
}
|
||||
decoder := json.NewDecoder(bytes.NewReader(raw))
|
||||
decoder.DisallowUnknownFields()
|
||||
if errDecode := decoder.Decode(&document); errDecode != nil {
|
||||
return fixture, errDecode
|
||||
}
|
||||
if errDecode := decoder.Decode(&struct{}{}); errDecode == nil {
|
||||
return fixture, errors.New("unexpected trailing JSON")
|
||||
} else if errDecode != io.EOF {
|
||||
return fixture, errDecode
|
||||
}
|
||||
documentRaw := struct {
|
||||
Part json.RawMessage `json:"part"`
|
||||
Overflow json.RawMessage `json:"overflow"`
|
||||
}{}
|
||||
if errDecode := json.Unmarshal(raw, &documentRaw); errDecode != nil {
|
||||
return fixture, errDecode
|
||||
}
|
||||
fixture.Part = document.Part
|
||||
fixture.Overflow = document.Overflow
|
||||
fixture.PartJSON = documentRaw.Part
|
||||
fixture.OverflowJSON = documentRaw.Overflow
|
||||
return fixture, nil
|
||||
}
|
||||
|
||||
func assertInFlightContractFields(t *testing.T) {
|
||||
t.Helper()
|
||||
assertOrderedInFlightJSONFields(t, reflect.TypeOf(InFlightSnapshotFrame{}), []inFlightJSONField{
|
||||
{name: "Kind", tag: "kind"},
|
||||
{name: "Revision", tag: "revision"},
|
||||
{name: "ObservedAt", tag: "observed_at"},
|
||||
{name: "BarrierRevision", tag: "barrier_revision"},
|
||||
{name: "PartIndex", tag: "part_index,omitempty"},
|
||||
{name: "PartCount", tag: "part_count,omitempty"},
|
||||
{name: "DetailsTruncated", tag: "details_truncated,omitempty"},
|
||||
{name: "Aggregates", tag: "aggregates,omitempty"},
|
||||
{name: "Details", tag: "details,omitempty"},
|
||||
{name: "AggregateGroupCount", tag: "aggregate_group_count,omitempty"},
|
||||
})
|
||||
assertOrderedInFlightJSONFields(t, reflect.TypeOf(InFlightAggregate{}), []inFlightJSONField{
|
||||
{name: "CredentialID", tag: "credential_id"},
|
||||
{name: "Model", tag: "model"},
|
||||
{name: "Status", tag: "status"},
|
||||
{name: "Count", tag: "count"},
|
||||
})
|
||||
assertOrderedInFlightJSONFields(t, reflect.TypeOf(InFlightRequestDetail{}), []inFlightJSONField{
|
||||
{name: "RequestID", tag: "request_id"},
|
||||
{name: "CredentialID", tag: "credential_id"},
|
||||
{name: "Model", tag: "model"},
|
||||
{name: "RequestKind", tag: "request_kind"},
|
||||
{name: "StartedAt", tag: "started_at"},
|
||||
})
|
||||
}
|
||||
|
||||
func assertInFlightFixtureKeys(t *testing.T, fixture inFlightContractFixture) {
|
||||
t.Helper()
|
||||
assertRequiredInFlightJSONKeys(t, fixture.PartJSON, []string{"kind", "revision", "observed_at", "barrier_revision", "part_index", "part_count", "details_truncated", "aggregates", "details"})
|
||||
assertRequiredInFlightJSONKeys(t, fixture.OverflowJSON, []string{"kind", "revision", "observed_at", "barrier_revision", "aggregate_group_count"})
|
||||
|
||||
var part struct {
|
||||
Aggregates []json.RawMessage `json:"aggregates"`
|
||||
Details []json.RawMessage `json:"details"`
|
||||
}
|
||||
if errDecode := json.Unmarshal(fixture.PartJSON, &part); errDecode != nil {
|
||||
t.Fatalf("json.Unmarshal() error = %v", errDecode)
|
||||
}
|
||||
for index, aggregate := range part.Aggregates {
|
||||
assertRequiredInFlightJSONKeys(t, aggregate, []string{"credential_id", "model", "status", "count"})
|
||||
if len(aggregate) == 0 {
|
||||
t.Fatalf("aggregate %d is empty", index)
|
||||
}
|
||||
}
|
||||
for index, detail := range part.Details {
|
||||
assertRequiredInFlightJSONKeys(t, detail, []string{"request_id", "credential_id", "model", "request_kind", "started_at"})
|
||||
if len(detail) == 0 {
|
||||
t.Fatalf("detail %d is empty", index)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type inFlightJSONField struct {
|
||||
name string
|
||||
tag string
|
||||
}
|
||||
|
||||
func assertOrderedInFlightJSONFields(t *testing.T, structType reflect.Type, want []inFlightJSONField) {
|
||||
t.Helper()
|
||||
if structType.NumField() != len(want) {
|
||||
t.Fatalf("%s field count = %d, want %d", structType.Name(), structType.NumField(), len(want))
|
||||
}
|
||||
for index, expected := range want {
|
||||
field := structType.Field(index)
|
||||
if field.Name != expected.name || field.Tag.Get("json") != expected.tag {
|
||||
t.Fatalf("%s field %d = (%q, %q), want (%q, %q)", structType.Name(), index, field.Name, field.Tag.Get("json"), expected.name, expected.tag)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func assertRequiredInFlightJSONKeys(t *testing.T, raw json.RawMessage, required []string) {
|
||||
t.Helper()
|
||||
var fields map[string]json.RawMessage
|
||||
if errDecode := json.Unmarshal(raw, &fields); errDecode != nil {
|
||||
t.Fatalf("json.Unmarshal() error = %v", errDecode)
|
||||
}
|
||||
for _, key := range required {
|
||||
if _, ok := fields[key]; !ok {
|
||||
t.Fatalf("required JSON key %q is missing", key)
|
||||
}
|
||||
}
|
||||
}
|
||||
189
backend/internal/home/kv_helpers.go
Normal file
189
backend/internal/home/kv_helpers.go
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
package home
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
func HashKeyPart(value string) string {
|
||||
sum := sha256.Sum256([]byte(value))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func CurrentKVClient() (*Client, bool, error) {
|
||||
client := Current()
|
||||
if client == nil {
|
||||
return nil, false, nil
|
||||
}
|
||||
if !client.Enabled() {
|
||||
return nil, true, fmt.Errorf("home kv store unavailable: %w", ErrDisabled)
|
||||
}
|
||||
if !client.HeartbeatOK() {
|
||||
return nil, true, fmt.Errorf("home kv store unavailable: %w", ErrNotConnected)
|
||||
}
|
||||
return client, true, nil
|
||||
}
|
||||
|
||||
func KVGetJSONRequired(ctx context.Context, key string, out any) (bool, bool, error) {
|
||||
client, homeMode, errClient := CurrentKVClient()
|
||||
if !homeMode || errClient != nil {
|
||||
return homeMode, false, errClient
|
||||
}
|
||||
raw, found, errGet := client.KVGet(ctx, key)
|
||||
if errGet != nil || !found {
|
||||
return true, false, errGet
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(raw, out); errUnmarshal != nil {
|
||||
return true, false, errUnmarshal
|
||||
}
|
||||
return true, true, nil
|
||||
}
|
||||
|
||||
func KVSetJSONRequired(ctx context.Context, key string, value any, ttl time.Duration) (bool, error) {
|
||||
raw, errMarshal := json.Marshal(value)
|
||||
if errMarshal != nil {
|
||||
return false, errMarshal
|
||||
}
|
||||
return KVSetBytesRequired(ctx, key, raw, ttl)
|
||||
}
|
||||
|
||||
func KVSetBytesRequired(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error) {
|
||||
client, homeMode, errClient := CurrentKVClient()
|
||||
if !homeMode || errClient != nil {
|
||||
return homeMode, errClient
|
||||
}
|
||||
written, errSet := client.KVSet(ctx, key, value, kvSetOptionsForTTL(ttl))
|
||||
if errSet != nil {
|
||||
return true, errSet
|
||||
}
|
||||
if !written {
|
||||
return true, fmt.Errorf("home kv store unavailable")
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func KVSetNXRequired(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, bool, error) {
|
||||
client, homeMode, errClient := CurrentKVClient()
|
||||
if !homeMode || errClient != nil {
|
||||
return homeMode, false, errClient
|
||||
}
|
||||
written, errSet := client.KVSetNX(ctx, key, value, ttl)
|
||||
return true, written, errSet
|
||||
}
|
||||
|
||||
func KVDelRequired(ctx context.Context, keys ...string) (bool, int64, error) {
|
||||
client, homeMode, errClient := CurrentKVClient()
|
||||
if !homeMode || errClient != nil {
|
||||
return homeMode, 0, errClient
|
||||
}
|
||||
deleted, errDel := client.KVDel(ctx, keys...)
|
||||
return true, deleted, errDel
|
||||
}
|
||||
|
||||
func KVExpireRequired(ctx context.Context, key string, ttl time.Duration) (bool, error) {
|
||||
client, homeMode, errClient := CurrentKVClient()
|
||||
if !homeMode || errClient != nil {
|
||||
return homeMode, errClient
|
||||
}
|
||||
_, errExpire := client.KVExpire(ctx, key, ttl)
|
||||
return true, errExpire
|
||||
}
|
||||
|
||||
func KVGetJSONBestEffort(ctx context.Context, key string, out any) (bool, bool) {
|
||||
homeMode, found, errGet := KVGetJSONRequired(ctx, key, out)
|
||||
if errGet != nil {
|
||||
log.Errorf("home kv best-effort get failed prefix=%s: %v", kvLogPrefix(key), errGet)
|
||||
return homeMode, false
|
||||
}
|
||||
return homeMode, found
|
||||
}
|
||||
|
||||
func KVSetJSONBestEffort(ctx context.Context, key string, value any, ttl time.Duration) bool {
|
||||
raw, errMarshal := json.Marshal(value)
|
||||
if errMarshal != nil {
|
||||
log.Errorf("home kv best-effort set failed prefix=%s: %v", kvLogPrefix(key), errMarshal)
|
||||
return false
|
||||
}
|
||||
return KVSetBytesBestEffort(ctx, key, raw, ttl)
|
||||
}
|
||||
|
||||
func KVSetBytesBestEffort(ctx context.Context, key string, value []byte, ttl time.Duration) bool {
|
||||
homeMode, errSet := KVSetBytesRequired(ctx, key, value, ttl)
|
||||
if !homeMode {
|
||||
return false
|
||||
}
|
||||
if errSet != nil {
|
||||
log.Errorf("home kv best-effort set failed prefix=%s: %v", kvLogPrefix(key), errSet)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func KVSetNXBestEffort(ctx context.Context, key string, value []byte, ttl time.Duration) bool {
|
||||
homeMode, written, errSet := KVSetNXRequired(ctx, key, value, ttl)
|
||||
if !homeMode {
|
||||
return false
|
||||
}
|
||||
if errSet != nil {
|
||||
log.Errorf("home kv best-effort setnx failed prefix=%s: %v", kvLogPrefix(key), errSet)
|
||||
return false
|
||||
}
|
||||
return written
|
||||
}
|
||||
|
||||
func KVDelBestEffort(ctx context.Context, keys ...string) bool {
|
||||
homeMode, _, errDel := KVDelRequired(ctx, keys...)
|
||||
if !homeMode {
|
||||
return false
|
||||
}
|
||||
if errDel != nil {
|
||||
log.Errorf("home kv best-effort del failed prefix=%s: %v", kvLogPrefix(firstKVKey(keys)), errDel)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func KVExpireBestEffort(ctx context.Context, key string, ttl time.Duration) bool {
|
||||
homeMode, errExpire := KVExpireRequired(ctx, key, ttl)
|
||||
if !homeMode {
|
||||
return false
|
||||
}
|
||||
if errExpire != nil {
|
||||
log.Errorf("home kv best-effort expire failed prefix=%s: %v", kvLogPrefix(key), errExpire)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func kvSetOptionsForTTL(ttl time.Duration) KVSetOptions {
|
||||
if ttl <= 0 {
|
||||
return KVSetOptions{}
|
||||
}
|
||||
return KVSetOptions{EX: ttl}
|
||||
}
|
||||
|
||||
func kvLogPrefix(key string) string {
|
||||
key = strings.TrimSpace(key)
|
||||
if key == "" {
|
||||
return "unknown"
|
||||
}
|
||||
parts := strings.Split(key, ":")
|
||||
if len(parts) >= 2 {
|
||||
return parts[0] + ":" + parts[1] + ":*"
|
||||
}
|
||||
return parts[0] + ":*"
|
||||
}
|
||||
|
||||
func firstKVKey(keys []string) string {
|
||||
if len(keys) == 0 {
|
||||
return ""
|
||||
}
|
||||
return keys[0]
|
||||
}
|
||||
110
backend/internal/home/kv_helpers_test.go
Normal file
110
backend/internal/home/kv_helpers_test.go
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
package home
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
func TestHashKeyPart(t *testing.T) {
|
||||
first := HashKeyPart("secret-value")
|
||||
again := HashKeyPart("secret-value")
|
||||
other := HashKeyPart("other-value")
|
||||
if first == "" || len(first) != 64 {
|
||||
t.Fatalf("HashKeyPart() = %q, want 64 hex chars", first)
|
||||
}
|
||||
if first != again {
|
||||
t.Fatalf("HashKeyPart() is not stable")
|
||||
}
|
||||
if first == other {
|
||||
t.Fatalf("HashKeyPart() returned same hash for different inputs")
|
||||
}
|
||||
if strings.Contains(first, "secret") || strings.Contains(first, "value") {
|
||||
t.Fatalf("HashKeyPart() leaked input: %q", first)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKVRequiredHelpersReturnNonHomeMode(t *testing.T) {
|
||||
ClearCurrent()
|
||||
t.Cleanup(ClearCurrent)
|
||||
|
||||
var out map[string]string
|
||||
homeMode, found, errGet := KVGetJSONRequired(context.Background(), "key", &out)
|
||||
if errGet != nil {
|
||||
t.Fatalf("KVGetJSONRequired() error = %v", errGet)
|
||||
}
|
||||
if homeMode || found {
|
||||
t.Fatalf("KVGetJSONRequired() = homeMode %v found %v, want false false", homeMode, found)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCurrentKVClientUnavailableErrors(t *testing.T) {
|
||||
t.Cleanup(ClearCurrent)
|
||||
|
||||
disabled := New(config.HomeConfig{Enabled: false})
|
||||
SetCurrent(disabled)
|
||||
if _, homeMode, errClient := CurrentKVClient(); !homeMode || errClient == nil {
|
||||
t.Fatalf("CurrentKVClient(disabled) = homeMode %v err %v, want true error", homeMode, errClient)
|
||||
}
|
||||
|
||||
notReady := New(config.HomeConfig{Enabled: true, Host: "127.0.0.1", Port: 1})
|
||||
SetCurrent(notReady)
|
||||
if _, homeMode, errClient := CurrentKVClient(); !homeMode || errClient == nil {
|
||||
t.Fatalf("CurrentKVClient(no heartbeat) = homeMode %v err %v, want true error", homeMode, errClient)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKVRequiredHelpersPropagateClientErrors(t *testing.T) {
|
||||
client, _ := newRedisCommandTestClient(t, func(args []string) string {
|
||||
return "-ERR home kv unavailable\r\n"
|
||||
})
|
||||
client.heartbeatOK.Store(true)
|
||||
SetCurrent(client)
|
||||
t.Cleanup(ClearCurrent)
|
||||
|
||||
var out map[string]string
|
||||
homeMode, _, errGet := KVGetJSONRequired(context.Background(), "cpa:test:key", &out)
|
||||
if !homeMode || errGet == nil {
|
||||
t.Fatalf("KVGetJSONRequired() = homeMode %v err %v, want true error", homeMode, errGet)
|
||||
}
|
||||
homeMode, errSet := KVSetJSONRequired(context.Background(), "cpa:test:key", map[string]string{"value": "secret"}, 0)
|
||||
if !homeMode || errSet == nil {
|
||||
t.Fatalf("KVSetJSONRequired() = homeMode %v err %v, want true error", homeMode, errSet)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKVBestEffortWriteSwallowsErrorAndRedactsLog(t *testing.T) {
|
||||
client, _ := newRedisCommandTestClient(t, func(args []string) string {
|
||||
return "-ERR home kv unavailable\r\n"
|
||||
})
|
||||
client.heartbeatOK.Store(true)
|
||||
SetCurrent(client)
|
||||
t.Cleanup(ClearCurrent)
|
||||
|
||||
logger := log.StandardLogger()
|
||||
previousOutput := logger.Out
|
||||
previousLevel := log.GetLevel()
|
||||
buffer := &bytes.Buffer{}
|
||||
log.SetOutput(buffer)
|
||||
log.SetLevel(log.ErrorLevel)
|
||||
t.Cleanup(func() {
|
||||
log.SetOutput(previousOutput)
|
||||
log.SetLevel(previousLevel)
|
||||
})
|
||||
|
||||
ok := KVSetJSONBestEffort(context.Background(), "cpa:test:secret-key", map[string]string{"value": "secret-value"}, 0)
|
||||
if ok {
|
||||
t.Fatalf("KVSetJSONBestEffort() = true, want false")
|
||||
}
|
||||
logText := buffer.String()
|
||||
if !strings.Contains(logText, "cpa:test:*") {
|
||||
t.Fatalf("log = %q, want redacted key prefix", logText)
|
||||
}
|
||||
if strings.Contains(logText, "secret-key") || strings.Contains(logText, "secret-value") {
|
||||
t.Fatalf("log leaked key or value: %q", logText)
|
||||
}
|
||||
}
|
||||
42
backend/internal/home/plugin_status.go
Normal file
42
backend/internal/home/plugin_status.go
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
package home
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/homeplugins"
|
||||
)
|
||||
|
||||
const pluginStatusReportTimeout = 10 * time.Second
|
||||
|
||||
// PluginStatusClient defines the interface for pushing plugin status reports.
|
||||
type PluginStatusClient interface {
|
||||
RPushPluginStatus(ctx context.Context, payload []byte) error
|
||||
}
|
||||
|
||||
// ReportPluginStatus marshals the given report, sets NodeID and UpdatedAt,
|
||||
// and pushes it to the provided client with a timeout.
|
||||
func ReportPluginStatus(ctx context.Context, client PluginStatusClient, nodeID string, report homeplugins.SyncReport) error {
|
||||
if client == nil {
|
||||
return fmt.Errorf("home plugin status client is unavailable")
|
||||
}
|
||||
nodeID = strings.TrimSpace(nodeID)
|
||||
if nodeID == "" {
|
||||
return fmt.Errorf("home plugin status node id is empty")
|
||||
}
|
||||
report.NodeID = nodeID
|
||||
report.UpdatedAt = time.Now().UTC()
|
||||
raw, errMarshal := json.Marshal(report)
|
||||
if errMarshal != nil {
|
||||
return errMarshal
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
reportCtx, cancel := context.WithTimeout(ctx, pluginStatusReportTimeout)
|
||||
defer cancel()
|
||||
return client.RPushPluginStatus(reportCtx, raw)
|
||||
}
|
||||
93
backend/internal/home/plugin_status_test.go
Normal file
93
backend/internal/home/plugin_status_test.go
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
package home
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/homeplugins"
|
||||
)
|
||||
|
||||
type recordingPluginStatusClient struct {
|
||||
payload []byte
|
||||
err error
|
||||
}
|
||||
|
||||
func (c *recordingPluginStatusClient) RPushPluginStatus(ctx context.Context, payload []byte) error {
|
||||
c.payload = append([]byte(nil), payload...)
|
||||
return c.err
|
||||
}
|
||||
|
||||
func TestReportPluginStatusPushesNodeReport(t *testing.T) {
|
||||
client := &recordingPluginStatusClient{}
|
||||
report := homeplugins.SyncReport{
|
||||
Task: "plugin-sync",
|
||||
Status: "success",
|
||||
OK: true,
|
||||
Plugins: []homeplugins.PluginInstallStatus{{ID: "sample", InstallStatus: "installed"}},
|
||||
}
|
||||
|
||||
if errReport := ReportPluginStatus(context.Background(), client, " node-1 ", report); errReport != nil {
|
||||
t.Fatalf("ReportPluginStatus() error = %v", errReport)
|
||||
}
|
||||
var payload homeplugins.SyncReport
|
||||
if errUnmarshal := json.Unmarshal(client.payload, &payload); errUnmarshal != nil {
|
||||
t.Fatalf("unmarshal payload: %v", errUnmarshal)
|
||||
}
|
||||
if payload.NodeID != "node-1" || !payload.OK || len(payload.Plugins) != 1 {
|
||||
t.Fatalf("payload = %+v, want node report", payload)
|
||||
}
|
||||
if payload.UpdatedAt.IsZero() {
|
||||
t.Fatal("payload UpdatedAt is zero")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReportPluginStatusPushesEmptyReport(t *testing.T) {
|
||||
client := &recordingPluginStatusClient{}
|
||||
report := homeplugins.SyncReport{
|
||||
Task: "plugin-sync",
|
||||
Status: "success",
|
||||
OK: true,
|
||||
Plugins: []homeplugins.PluginInstallStatus{},
|
||||
}
|
||||
|
||||
if errReport := ReportPluginStatus(context.Background(), client, "node-1", report); errReport != nil {
|
||||
t.Fatalf("ReportPluginStatus() error = %v", errReport)
|
||||
}
|
||||
var payload homeplugins.SyncReport
|
||||
if errUnmarshal := json.Unmarshal(client.payload, &payload); errUnmarshal != nil {
|
||||
t.Fatalf("unmarshal payload: %v", errUnmarshal)
|
||||
}
|
||||
if payload.NodeID != "node-1" || len(payload.Plugins) != 0 {
|
||||
t.Fatalf("payload = %+v, want empty node report", payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReportPluginStatusRequiresNodeID(t *testing.T) {
|
||||
client := &recordingPluginStatusClient{}
|
||||
report := homeplugins.SyncReport{
|
||||
Plugins: []homeplugins.PluginInstallStatus{{ID: "sample", InstallStatus: "failed"}},
|
||||
}
|
||||
|
||||
errReport := ReportPluginStatus(context.Background(), client, " ", report)
|
||||
if errReport == nil || !strings.Contains(errReport.Error(), "node id") {
|
||||
t.Fatalf("ReportPluginStatus() error = %v, want node id error", errReport)
|
||||
}
|
||||
if len(client.payload) != 0 {
|
||||
t.Fatalf("client payload = %s, want none", client.payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReportPluginStatusPropagatesPushError(t *testing.T) {
|
||||
client := &recordingPluginStatusClient{err: errors.New("push failed")}
|
||||
report := homeplugins.SyncReport{
|
||||
Plugins: []homeplugins.PluginInstallStatus{{ID: "sample", InstallStatus: "installed"}},
|
||||
}
|
||||
|
||||
errReport := ReportPluginStatus(context.Background(), client, "node-1", report)
|
||||
if !errors.Is(errReport, client.err) {
|
||||
t.Fatalf("ReportPluginStatus() error = %v, want push failed", errReport)
|
||||
}
|
||||
}
|
||||
66
backend/internal/home/requests.go
Normal file
66
backend/internal/home/requests.go
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
package home
|
||||
|
||||
import "time"
|
||||
|
||||
type authDispatchRequest struct {
|
||||
Type string `json:"type"`
|
||||
Model string `json:"model"`
|
||||
Count int `json:"count"`
|
||||
ConcurrencyProtocol int `json:"concurrency_protocol,omitempty"`
|
||||
SessionID string `json:"session_id,omitempty"`
|
||||
Headers map[string]string `json:"headers,omitempty"`
|
||||
CredentialPolicy string `json:"credential_policy,omitempty"`
|
||||
RetryRound *int `json:"retry_round,omitempty"`
|
||||
ExcludedAuthIDs *[]string `json:"excluded_auth_ids,omitempty"`
|
||||
PinnedAuthID string `json:"pinned_auth_id,omitempty"`
|
||||
}
|
||||
|
||||
type modelsRequest struct {
|
||||
Type string `json:"type"`
|
||||
Headers map[string]string `json:"headers,omitempty"`
|
||||
Query map[string]string `json:"query,omitempty"`
|
||||
}
|
||||
|
||||
type refreshRequest struct {
|
||||
Type string `json:"type"`
|
||||
AuthIndex string `json:"auth_index"`
|
||||
ObservedAccessTokenSHA256 string `json:"access_token_sha256,omitempty"`
|
||||
}
|
||||
|
||||
type InFlightFrameKind string
|
||||
type InFlightAccountedStatus string
|
||||
|
||||
const (
|
||||
InFlightFramePart InFlightFrameKind = "part"
|
||||
InFlightFrameOverflow InFlightFrameKind = "overflow"
|
||||
InFlightAccounted InFlightAccountedStatus = "accounted"
|
||||
InFlightUnaccounted InFlightAccountedStatus = "unaccounted"
|
||||
)
|
||||
|
||||
type InFlightAggregate struct {
|
||||
CredentialID string `json:"credential_id"`
|
||||
Model string `json:"model"`
|
||||
Status InFlightAccountedStatus `json:"status"`
|
||||
Count int64 `json:"count"`
|
||||
}
|
||||
|
||||
type InFlightRequestDetail struct {
|
||||
RequestID string `json:"request_id"`
|
||||
CredentialID string `json:"credential_id"`
|
||||
Model string `json:"model"`
|
||||
RequestKind string `json:"request_kind"`
|
||||
StartedAt time.Time `json:"started_at"`
|
||||
}
|
||||
|
||||
type InFlightSnapshotFrame struct {
|
||||
Kind InFlightFrameKind `json:"kind"`
|
||||
Revision int64 `json:"revision"`
|
||||
ObservedAt time.Time `json:"observed_at"`
|
||||
BarrierRevision int64 `json:"barrier_revision"`
|
||||
PartIndex *int `json:"part_index,omitempty"`
|
||||
PartCount *int `json:"part_count,omitempty"`
|
||||
DetailsTruncated bool `json:"details_truncated,omitempty"`
|
||||
Aggregates []InFlightAggregate `json:"aggregates,omitempty"`
|
||||
Details []InFlightRequestDetail `json:"details,omitempty"`
|
||||
AggregateGroupCount int `json:"aggregate_group_count,omitempty"`
|
||||
}
|
||||
27
backend/internal/home/testdata/concurrency_dispatch_accounted.json
vendored
Normal file
27
backend/internal/home/testdata/concurrency_dispatch_accounted.json
vendored
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
{
|
||||
"model": "gpt",
|
||||
"provider": "codex",
|
||||
"auth_index": "cred-1",
|
||||
"user_api_key": "user-key",
|
||||
"auth": {
|
||||
"id": "cred-1",
|
||||
"provider": "codex",
|
||||
"status": "active",
|
||||
"disabled": false,
|
||||
"unavailable": false,
|
||||
"quota": {
|
||||
"exceeded": false,
|
||||
"next_recover_at": "0001-01-01T00:00:00Z"
|
||||
},
|
||||
"created_at": "0001-01-01T00:00:00Z",
|
||||
"updated_at": "0001-01-01T00:00:00Z",
|
||||
"last_refreshed_at": "0001-01-01T00:00:00Z",
|
||||
"next_refresh_after": "0001-01-01T00:00:00Z",
|
||||
"next_retry_after": "0001-01-01T00:00:00Z"
|
||||
},
|
||||
"concurrency": {
|
||||
"accounted": true,
|
||||
"credential_id": "cred-1",
|
||||
"model": "gpt"
|
||||
}
|
||||
}
|
||||
8
backend/internal/home/testdata/concurrency_dispatch_busy.json
vendored
Normal file
8
backend/internal/home/testdata/concurrency_dispatch_busy.json
vendored
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
{
|
||||
"error": {
|
||||
"type": "credential_concurrency_exceeded",
|
||||
"message": "credential concurrency limit reached",
|
||||
"retryable": true,
|
||||
"retry_after_ms": 750
|
||||
}
|
||||
}
|
||||
1
backend/internal/home/testdata/concurrency_release.json
vendored
Normal file
1
backend/internal/home/testdata/concurrency_release.json
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
{"credential_id":"cred-1","model":"gpt","release_seq":1}
|
||||
52
backend/internal/home/testdata/credential_in_flight_contract.json
vendored
Normal file
52
backend/internal/home/testdata/credential_in_flight_contract.json
vendored
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
{
|
||||
"config": {
|
||||
"snapshot-interval": "2s",
|
||||
"stale-after": "10s",
|
||||
"max-part-bytes": 262144,
|
||||
"max-part-count": 64,
|
||||
"max-revision-bytes": 16777216,
|
||||
"max-aggregate-groups": 100000,
|
||||
"max-details": 10000,
|
||||
"max-string-bytes": 256,
|
||||
"staging-retention": "1m"
|
||||
},
|
||||
"part": {
|
||||
"kind": "part",
|
||||
"revision": 7,
|
||||
"observed_at": "2026-07-21T12:00:00Z",
|
||||
"barrier_revision": 11,
|
||||
"part_index": 0,
|
||||
"part_count": 1,
|
||||
"details_truncated": false,
|
||||
"aggregates": [
|
||||
{
|
||||
"credential_id": "cred-a",
|
||||
"model": "gpt-5",
|
||||
"status": "accounted",
|
||||
"count": 2
|
||||
},
|
||||
{
|
||||
"credential_id": "cred-a",
|
||||
"model": "gpt-5",
|
||||
"status": "unaccounted",
|
||||
"count": 1
|
||||
}
|
||||
],
|
||||
"details": [
|
||||
{
|
||||
"request_id": "req-1",
|
||||
"credential_id": "cred-a",
|
||||
"model": "gpt-5",
|
||||
"request_kind": "sse",
|
||||
"started_at": "2026-07-21T11:59:58Z"
|
||||
}
|
||||
]
|
||||
},
|
||||
"overflow": {
|
||||
"kind": "overflow",
|
||||
"revision": 8,
|
||||
"observed_at": "2026-07-21T12:00:02Z",
|
||||
"barrier_revision": 12,
|
||||
"aggregate_group_count": 100001
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue