internal/mcpgateway/oauth_store.go
Rama main · réplica de lectura del repositorio original.
package mcpgateway
import (
"context"
"errors"
"fmt"
"maps"
"slices"
"time"
"github.com/goccy/go-json"
"github.com/enterpilot/gomodel/internal/encryption"
)
// oauthSecretKind and oauthSecretField name sealed OAuth records in their
// AAD, so a record cannot be moved to another server or table.
const (
oauthSecretKind = "mcp_server_oauth" //nolint:gosec // An AAD label, not a credential.
oauthSecretField = "data"
)
// oauthRecord is the OAuth client registration and token set the gateway
// holds for one upstream server after an operator signed in. It is runtime
// state, never config, and never leaves the gateway: admin views expose
// OAuthView only. The whole record is sealed as one value, so client secrets
// and tokens get the same encryption at rest as MCP headers.
type oauthRecord struct {
// ServerURL is the upstream URL the tokens were issued for; a server
// whose URL changed afterwards no longer uses them.
ServerURL string `json:"server_url"`
// Resource is the RFC 8707 resource indicator sent on token requests.
Resource string `json:"resource"`
Issuer string `json:"issuer"`
TokenURL string `json:"token_url"`
// RedirectURI is the callback the client was registered and authorized
// with; a stored client is only reused for the same issuer and callback.
RedirectURI string `json:"redirect_uri"`
ClientID string `json:"client_id"`
ClientSecret string `json:"client_secret,omitempty"`
// AuthStyle is the oauth2.AuthStyle for token requests.
AuthStyle int `json:"auth_style"`
// Registered marks a client obtained by dynamic client registration.
Registered bool `json:"registered,omitempty"`
Scopes []string `json:"scopes,omitempty"`
AccessToken string `json:"access_token,omitempty"`
RefreshToken string `json:"refresh_token,omitempty"`
TokenType string `json:"token_type,omitempty"`
Expiry time.Time `json:"expiry"`
SignedInAt time.Time `json:"signed_in_at"`
// UpdatedAt changes on every token refresh, so an instance can tell a
// newer stored token from its own.
UpdatedAt time.Time `json:"updated_at"`
// RefreshOwner and RefreshUntil lease the next token refresh to one
// gateway instance, so replicas sharing the store never spend the same
// refresh token twice (see oauthSource.renew).
RefreshOwner string `json:"refresh_owner,omitempty"`
RefreshUntil time.Time `json:"refresh_until,omitzero"`
}
// sameSignIn reports whether other comes from the same sign-in as r.
func (r oauthRecord) sameSignIn(other oauthRecord) bool {
return r.SignedInAt.Equal(other.SignedInAt) && r.ClientID == other.ClientID
}
// leasedElsewhere reports whether another instance holds the refresh lease.
func (r oauthRecord) leasedElsewhere(owner string, now time.Time) bool {
return r.RefreshOwner != "" && r.RefreshOwner != owner && now.Before(r.RefreshUntil)
}
// oauthRowStore persists sealed OAuth records by server name. Both store
// backends implement it.
type oauthRowStore interface {
listOAuth(ctx context.Context) (map[string]string, error)
getOAuth(ctx context.Context, name string) (string, error)
putOAuth(ctx context.Context, name, data string) error
deleteOAuth(ctx context.Context, name string) error
swapOAuth(ctx context.Context, name, current, next string) (bool, error)
}
// oauthStore seals records on the way in and opens them on the way out. A
// nil box stores them in plaintext, like MCP headers without
// GOMODEL_ENCRYPTION_KEY.
type oauthStore struct {
rows oauthRowStore
box *encryption.Box
}
func newOAuthStore(store Store, box *encryption.Box) *oauthStore {
rows, ok := store.(oauthRowStore)
if !ok {
return nil
}
return &oauthStore{rows: rows, box: box}
}
func (s *oauthStore) seal(name string, record oauthRecord) (string, error) {
data, err := json.Marshal(record)
if err != nil {
return "", fmt.Errorf("encode mcp oauth record: %w", err)
}
value := string(data)
if err := s.box.SealFields(oauthSecretKind, name, encryption.Field{Name: oauthSecretField, Value: &value}); err != nil {
return "", err
}
return value, nil
}
func (s *oauthStore) open(name, value string) (oauthRecord, error) {
if err := s.box.OpenFields(oauthSecretKind, name, encryption.Field{Name: oauthSecretField, Value: &value}); err != nil {
return oauthRecord{}, err
}
var record oauthRecord
if err := json.Unmarshal([]byte(value), &record); err != nil {
// The decoder error can quote the input, which holds tokens.
return oauthRecord{}, fmt.Errorf("decode mcp oauth record for %q: invalid JSON", name)
}
return record, nil
}
// list returns every record by server name. A record that cannot be opened
// is skipped and reported in the joined error, so one bad row does not sign
// every other server out.
func (s *oauthStore) list(ctx context.Context) (map[string]oauthRecord, error) {
rows, err := s.rows.listOAuth(ctx)
if err != nil {
return nil, err
}
records := make(map[string]oauthRecord, len(rows))
var errs []error
for _, name := range slices.Sorted(maps.Keys(rows)) {
record, err := s.open(name, rows[name])
if err != nil {
errs = append(errs, err)
continue
}
records[name] = record
}
return records, errors.Join(errs...)
}
func (s *oauthStore) get(ctx context.Context, name string) (*oauthRecord, error) {
value, err := s.rows.getOAuth(ctx, name)
if err != nil {
return nil, err
}
record, err := s.open(name, value)
if err != nil {
return nil, err
}
return &record, nil
}
// put stores a new sign-in, replacing any earlier one. An error wrapping
// encryption.ErrSealUnconfirmed means the record was written.
func (s *oauthStore) put(ctx context.Context, name string, record oauthRecord) error {
value, err := s.seal(name, record)
if err != nil {
return err
}
if err := s.rows.putOAuth(ctx, name, value); err != nil {
return err
}
return s.confirmKey(ctx, name, record, value)
}
// oauthUpdateAttempts bounds how often update re-reads a row that changed
// under it, such as one a data key rotation re-sealed meanwhile.
const oauthUpdateAttempts = 3
// update replaces the stored record with record only while the stored one
// still satisfies same, so a refresh from an old sign-in never undoes a
// sign-out or a newer sign-in. A row that changed between the read and the
// write is read again, so a re-seal by a key rotation does not lose a
// rotated refresh token. It reports whether it wrote.
func (s *oauthStore) update(ctx context.Context, name string, record oauthRecord, same func(stored oauthRecord) bool) (bool, error) {
for range oauthUpdateAttempts {
current, err := s.rows.getOAuth(ctx, name)
if errors.Is(err, ErrNotFound) {
return false, nil
}
if err != nil {
return false, err
}
stored, err := s.open(name, current)
if err != nil {
return false, err
}
if !same(stored) {
return false, nil
}
next, err := s.seal(name, record)
if err != nil {
return false, err
}
swapped, err := s.rows.swapOAuth(ctx, name, current, next)
if err != nil {
return false, err
}
if swapped {
return true, s.confirmKey(ctx, name, record, next)
}
}
return false, fmt.Errorf("save mcp oauth record for %q: the row kept changing", name)
}
// confirmKey re-seals a just-written record when a data key rotation
// activated a new key while the save was in flight (see Box.ConfirmSeal).
// The swap is conditional, so a newer save is never overwritten.
func (s *oauthStore) confirmKey(ctx context.Context, name string, record oauthRecord, written string) error {
return s.box.ConfirmSeal(func() error {
resealed, err := s.seal(name, record)
if err != nil {
return err
}
_, err = s.rows.swapOAuth(ctx, name, written, resealed)
return err
}, encryption.Field{Name: oauthSecretField, Value: &written})
}
func (s *oauthStore) delete(ctx context.Context, name string) error {
return s.rows.deleteOAuth(ctx, name)
}
// reencrypt re-seals every record stored in plaintext or under a data key
// other than the active one. Writes are conditional, so a refresh that
// landed meanwhile is retried rather than overwritten.
func (s *oauthStore) reencrypt(ctx context.Context) (encryption.Report, error) {
const entity = "mcp_server_oauth"
rows, err := s.rows.listOAuth(ctx)
if err != nil {
return encryption.Report{Entity: entity}, err
}
return encryption.ReencryptRows(entity, slices.Sorted(maps.Keys(rows)), func(name string) (encryption.RowOutcome, error) {
current, err := s.rows.getOAuth(ctx, name)
if errors.Is(err, ErrNotFound) {
return encryption.RowUnchanged, nil
}
if err != nil {
return 0, err
}
if !s.box.NeedsReseal(encryption.Field{Name: oauthSecretField, Value: ¤t}) {
return encryption.RowUnchanged, nil
}
record, err := s.open(name, current)
if err != nil {
return 0, err
}
next, err := s.seal(name, record)
if err != nil {
return 0, err
}
swapped, err := s.rows.swapOAuth(ctx, name, current, next)
if err != nil {
return 0, err
}
if !swapped {
return encryption.RowConflict, nil
}
return encryption.RowRewritten, nil
})
}