122 lines
3.5 KiB
Go
122 lines
3.5 KiB
Go
package agent
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// TokenProvider supplies the current bearer token. A static token is served
|
|
// forever; a worker key is exchanged at the userservice for short-lived JWTs
|
|
// and refreshed before they expire (mirroring the former Python worker).
|
|
type TokenProvider interface {
|
|
Token() (string, error)
|
|
Refresh() error
|
|
}
|
|
|
|
// StaticToken serves a fixed token forever; empty means no Authorization.
|
|
type StaticToken struct{ token string }
|
|
|
|
func (s *StaticToken) Token() (string, error) { return s.token, nil }
|
|
func (s *StaticToken) Refresh() error { return nil }
|
|
|
|
// WorkerKeyToken exchanges a long-lived worker key for short-lived JWTs.
|
|
type WorkerKeyToken struct {
|
|
userserviceURL string
|
|
workerKey string
|
|
timeout time.Duration
|
|
leeway float64
|
|
mu sync.Mutex
|
|
token string
|
|
refreshAt time.Time
|
|
}
|
|
|
|
func NewWorkerKeyToken(userserviceURL, workerKey string, timeout time.Duration) *WorkerKeyToken {
|
|
return &WorkerKeyToken{
|
|
userserviceURL: strings.TrimRight(userserviceURL, "/"),
|
|
workerKey: workerKey,
|
|
timeout: timeout,
|
|
leeway: 0.2,
|
|
}
|
|
}
|
|
|
|
// Token returns the current token, exchanging first when missing or stale.
|
|
func (p *WorkerKeyToken) Token() (string, error) {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
if p.token == "" || time.Now().After(p.refreshAt) {
|
|
if err := p.exchangeLocked(); err != nil {
|
|
return "", err
|
|
}
|
|
}
|
|
return p.token, nil
|
|
}
|
|
|
|
// Refresh forces an immediate exchange.
|
|
func (p *WorkerKeyToken) Refresh() error {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
return p.exchangeLocked()
|
|
}
|
|
|
|
func (p *WorkerKeyToken) exchangeLocked() error {
|
|
payload, err := json.Marshal(map[string]string{"key": p.workerKey})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
request, err := http.NewRequestWithContext(context.Background(), http.MethodPost, p.userserviceURL+"/worker-tokens/exchange", bytes.NewReader(payload))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
request.Header.Set("Content-Type", "application/json")
|
|
client := &http.Client{Timeout: p.timeout}
|
|
response, err := client.Do(request)
|
|
if err != nil {
|
|
return fmt.Errorf("worker key exchange request failed")
|
|
}
|
|
defer func() { _ = response.Body.Close() }()
|
|
if response.StatusCode != http.StatusOK {
|
|
return fmt.Errorf("worker key exchange rejected with status %d", response.StatusCode)
|
|
}
|
|
raw, err := io.ReadAll(io.LimitReader(response.Body, 1<<20))
|
|
if err != nil {
|
|
return fmt.Errorf("worker key exchange request failed")
|
|
}
|
|
var data map[string]any
|
|
if err := json.Unmarshal(raw, &data); err != nil {
|
|
return fmt.Errorf("worker key exchange response is invalid")
|
|
}
|
|
token, _ := data["token"].(string)
|
|
if token == "" {
|
|
return fmt.Errorf("worker key exchange response is missing a token")
|
|
}
|
|
var ttl time.Duration
|
|
switch value := data["expires_in"].(type) {
|
|
case float64:
|
|
ttl = time.Duration(value * float64(time.Second))
|
|
case int:
|
|
ttl = time.Duration(value) * time.Second
|
|
}
|
|
p.token = token
|
|
p.refreshAt = time.Time{}
|
|
if ttl > 0 {
|
|
p.refreshAt = time.Now().Add(time.Duration(float64(ttl) * (1.0 - p.leeway)))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// NewTokenProvider picks the strategy: a worker key (with userservice) wins
|
|
// over a static bearer token.
|
|
func NewTokenProvider(workerKey, userserviceURL, bearerToken string, timeout time.Duration) TokenProvider {
|
|
if workerKey != "" && userserviceURL != "" {
|
|
return NewWorkerKeyToken(userserviceURL, workerKey, timeout)
|
|
}
|
|
return &StaticToken{token: bearerToken}
|
|
}
|