Skip to content

chore: remove vault encryption hooks from certain GORM tables in favor of AES-only encryption #4245

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Pratham-Mishra04 merged 1 commit into from
Jun 11, 2026
Merged

chore: remove vault encryption hooks from certain GORM tables in favor of AES-only encryption #4245

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
132 changes: 10 additions & 122 deletions framework/configstore/rdb.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1828,31 +1828,7 @@ func (s *RDBConfigStore) UpdateMCPClientConfig(ctx context.Context, id string, c
}

headersJSONStr := string(headersJSON)
vaultEnabled := tables.VaultIsEnabled()
mcpTableName := tables.TableMCPClient{}.TableName()
vaultPfx := tables.VaultPrefix()
if tables.VaultHooks.Prefix != nil {
vaultPfx = tables.VaultHooks.Prefix()
}
headersPath := fmt.Sprintf("%s/%s/%s/headers_json", vaultPfx, mcpTableName, id)
connPath := fmt.Sprintf("%s/%s/%s/connection_string", vaultPfx, mcpTableName, id)
// StoreString runs before Updates so the vault ref is available to write to DB.
// vaultStoredPaths tracks paths written to vault so we can compensate (remove)
// if the subsequent DB write fails — preventing vault/DB divergence.
// vaultRemovePaths tracks cleared/env-backed paths to remove post-commit so a
// DB failure doesn't leave DB holding a ref to an already-deleted vault entry.
var vaultStoredPaths []string
var vaultRemovePaths []string
if vaultEnabled {
if headersJSONStr != "" && headersJSONStr != "{}" {
if err := tables.VaultHooks.StoreString(ctx, headersPath, &headersJSONStr); err != nil {
return fmt.Errorf("failed to vault mcp headers: %w", err)
}
vaultStoredPaths = append(vaultStoredPaths, headersPath)
} else if tables.VaultHooks.Remove != nil {
vaultRemovePaths = append(vaultRemovePaths, headersPath)
}
} else if encrypt.IsEnabled() && headersJSONStr != "" && headersJSONStr != "{}" {
if encrypt.IsEnabled() && headersJSONStr != "" && headersJSONStr != "{}" {
encrypted, encErr := encrypt.Encrypt(headersJSONStr)
if encErr != nil {
return fmt.Errorf("failed to encrypt mcp headers: %w", encErr)
Expand All @@ -1875,9 +1851,7 @@ func (s *RDBConfigStore) UpdateMCPClientConfig(ctx context.Context, id string, c
"disabled": clientConfigCopy.Disabled,
"updated_at": time.Now(),
}
if vaultEnabled {
updates["encryption_status"] = tables.EncryptionStatusVault
} else if encrypt.IsEnabled() {
if encrypt.IsEnabled() {
updates["encryption_status"] = encryptionStatusEncrypted
}
if clientConfigCopy.OauthConfigID != nil {
Expand All @@ -1901,19 +1875,7 @@ func (s *RDBConfigStore) UpdateMCPClientConfig(ctx context.Context, id string, c
// also sync connection/auth metadata from config.json and persist the hash.
if clientConfigCopy.ConfigHash != "" {
connectionStringToPersist := clientConfigCopy.ConnectionString
if vaultEnabled {
if connectionStringToPersist != nil && !connectionStringToPersist.IsFromEnv() && connectionStringToPersist.GetValue() != "" {
// Mirror TableMCPClient.BeforeSave vault behavior for map-based Updates.
cs := *connectionStringToPersist
if err := tables.VaultHooks.StoreString(ctx, connPath, &cs.Val); err != nil {
return fmt.Errorf("failed to vault mcp connection string: %w", err)
}
connectionStringToPersist = &cs
vaultStoredPaths = append(vaultStoredPaths, connPath)
} else if tables.VaultHooks.Remove != nil {
vaultRemovePaths = append(vaultRemovePaths, connPath)
}
} else if encrypt.IsEnabled() && connectionStringToPersist != nil &&
if encrypt.IsEnabled() && connectionStringToPersist != nil &&
!connectionStringToPersist.IsFromEnv() && connectionStringToPersist.GetValue() != "" {
// Mirror TableMCPClient.BeforeSave behavior for map-based Updates.
cs := *connectionStringToPersist
Expand Down Expand Up @@ -1942,66 +1904,24 @@ func (s *RDBConfigStore) UpdateMCPClientConfig(ctx context.Context, id string, c
}

if err := tx.WithContext(ctx).Model(&existingClient).Updates(updates).Error; err != nil {
// DB write failed — compensate by removing any vault entries we just wrote
// so vault and DB don't diverge (old DB ref would point to nothing).
if tables.VaultHooks.Remove != nil {
for _, p := range vaultStoredPaths {
_ = tables.VaultHooks.Remove(ctx, p)
}
}
return s.parseGormError(err)
}
// Post-commit: best-effort vault removal for cleared/env-backed fields.
// Runs after DB succeeds so a failed Updates doesn't leave DB holding a
// ref to a vault entry we already deleted.
for _, p := range vaultRemovePaths {
_ = tables.VaultHooks.Remove(ctx, p)
}
return nil
})
}

// DeleteMCPClientConfig deletes an MCP client configuration from the database.
func (s *RDBConfigStore) DeleteMCPClientConfig(ctx context.Context, id string) error {
// Find existing client upfront so we can pre-select vault-backed rows.
var existingClient tables.TableMCPClient
if err := s.DB().WithContext(ctx).Where("client_id = ?", id).First(&existingClient).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return fmt.Errorf("MCP client with id '%s' not found", id)
}
return err
}

// Fetch IDs of vault-backed rows before deletion for post-tx cleanup.
var vaultTokenIDs, vaultSessionIDs, vaultCredIDs []string
if tables.VaultHooks.Remove != nil {
var tokens []tables.TableOauthUserToken
if err := s.DB().WithContext(ctx).Select("id").
Where("mcp_client_id = ? AND encryption_status = ?", existingClient.ClientID, tables.EncryptionStatusVault).
Find(&tokens).Error; err == nil {
for _, t := range tokens {
vaultTokenIDs = append(vaultTokenIDs, t.ID)
}
}
var sessions []tables.TableOauthUserSession
if err := s.DB().WithContext(ctx).Select("id").
Where("mcp_client_id = ? AND encryption_status = ?", existingClient.ClientID, tables.EncryptionStatusVault).
Find(&sessions).Error; err == nil {
for _, sess := range sessions {
vaultSessionIDs = append(vaultSessionIDs, sess.ID)
}
}
var creds []tables.TableMCPPerUserHeaderCredential
if err := s.DB().WithContext(ctx).Select("id").
Where("mcp_client_id = ? AND encryption_status = ?", existingClient.ClientID, tables.EncryptionStatusVault).
Find(&creds).Error; err == nil {
for _, c := range creds {
vaultCredIDs = append(vaultCredIDs, c.ID)
return s.DB().Transaction(func(tx *gorm.DB) error {
// Find existing client
var existingClient tables.TableMCPClient
if err := dbForUpdate(tx.WithContext(ctx)).Where("client_id = ?", id).First(&existingClient).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return fmt.Errorf("MCP client with id '%s' not found", id)
}
return err
}
}

err := s.DB().Transaction(func(tx *gorm.DB) error {
// Delete any virtual key MCP configs that reference this client
var configIDs []uint
if err := dbForUpdate(tx.WithContext(ctx)).
Expand Down Expand Up @@ -2038,20 +1958,6 @@ func (s *RDBConfigStore) DeleteMCPClientConfig(ctx context.Context, id string) e
// Delete the client (this will also handle foreign key cascades)
return tx.WithContext(ctx).Delete(&existingClient).Error
})
if err != nil {
return err
}

// Best-effort vault cleanup after successful transaction.
if len(vaultTokenIDs) > 0 || len(vaultSessionIDs) > 0 || len(vaultCredIDs) > 0 {
go func() {
tables.TableOauthUserToken{}.DeleteVaultSecrets(context.Background(), vaultTokenIDs)
tables.TableOauthUserSession{}.DeleteVaultSecrets(context.Background(), vaultSessionIDs)
tables.TableMCPPerUserHeaderCredential{}.DeleteVaultSecrets(context.Background(), vaultCredIDs)
}()
}

return nil
}

// GetVectorStoreConfig retrieves the vector store configuration from the database.
Expand Down Expand Up @@ -5106,40 +5012,22 @@ func (s *RDBConfigStore) DeleteTempTokensByResourceID(ctx context.Context, scope
if len(tx) > 0 {
db = tx[0]
}
var vaultIDs []string
if tables.VaultHooks.Remove != nil {
db.WithContext(ctx).Model(&tables.TempToken{}).
Where("scope = ? AND resource_id = ? AND encryption_status = ?", scope, resourceID, tables.EncryptionStatusVault).
Pluck("id", &vaultIDs)
}
res := db.WithContext(ctx).
Where("scope = ? AND resource_id = ?", scope, resourceID).
Delete(&tables.TempToken{})
if res.Error != nil {
return 0, res.Error
}
if len(vaultIDs) > 0 {
go tables.TempToken{}.DeleteVaultSecrets(context.Background(), vaultIDs)
}
return res.RowsAffected, nil
}

// DeleteExpiredTempTokens hard-deletes rows whose expires_at is at or before
// the given cutoff. Returns the number of rows deleted.
func (s *RDBConfigStore) DeleteExpiredTempTokens(ctx context.Context, before time.Time) (int64, error) {
var vaultIDs []string
if tables.VaultHooks.Remove != nil {
s.DB().WithContext(ctx).Model(&tables.TempToken{}).
Where("expires_at <= ? AND encryption_status = ?", before, tables.EncryptionStatusVault).
Pluck("id", &vaultIDs)
}
res := s.DB().WithContext(ctx).Where("expires_at <= ?", before).Delete(&tables.TempToken{})
if res.Error != nil {
return 0, res.Error
}
if len(vaultIDs) > 0 {
go tables.TempToken{}.DeleteVaultSecrets(context.Background(), vaultIDs)
}
return res.RowsAffected, nil
}

Expand Down
52 changes: 2 additions & 50 deletions framework/configstore/tables/mcp.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,31 +194,7 @@ func (c *TableMCPClient) BeforeSave(tx *gorm.DB) error {
// Encrypt sensitive fields after serialization.
// Always set EncryptionStatus when encryption is enabled so the startup
// batch pass does not re-process this row indefinitely.
if VaultIsEnabled() {
connPath := fmt.Sprintf("%s/%s/%s/%s", VaultPrefix(), c.TableName(), c.ClientID,
tx.Statement.DB.NamingStrategy.ColumnName("", "ConnectionString"))
headersPath := fmt.Sprintf("%s/%s/%s/%s", VaultPrefix(), c.TableName(), c.ClientID,
tx.Statement.DB.NamingStrategy.ColumnName("", "HeadersJSON"))
if c.ConnectionString != nil && !c.ConnectionString.IsFromEnv() && c.ConnectionString.GetValue() != "" {
cs := *c.ConnectionString
if err := vaultEnvVar(tx.Statement.Context, connPath, &cs); err != nil {
return fmt.Errorf("failed to vault mcp connection string: %w", err)
}
c.ConnectionString = &cs
} else {
// Field cleared or switched to env-var — remove any stale vault entry.
removeVaultEnvVar(tx.Statement.Context, connPath, c.ConnectionString)
}
if c.HeadersJSON != "" && c.HeadersJSON != "{}" {
if err := vaultString(tx.Statement.Context, headersPath, &c.HeadersJSON); err != nil {
return fmt.Errorf("failed to vault mcp headers: %w", err)
}
} else {
// Headers cleared — remove any stale vault entry.
removeVaultString(tx.Statement.Context, headersPath, &c.HeadersJSON)
}
c.EncryptionStatus = EncryptionStatusVault
} else if encrypt.IsEnabled() {
if encrypt.IsEnabled() {
if c.ConnectionString != nil && !c.ConnectionString.IsFromEnv() && c.ConnectionString.GetValue() != "" {
// Copy to avoid encrypting the shared ConnectionString through the pointer
cs := *c.ConnectionString
Expand All @@ -245,19 +221,7 @@ func (c *TableMCPClient) BeforeSave(tx *gorm.DB) error {
// AfterFind is a GORM hook that decrypts the connection string and headers (if encrypted)
// and deserializes JSON columns back into runtime structs after reading from the database.
func (c *TableMCPClient) AfterFind(tx *gorm.DB) error {
switch c.EncryptionStatus {
case EncryptionStatusVault:
if c.HeadersJSON != "" && c.HeadersJSON != "{}" {
if err := resolveVaultString(tx.Statement.Context, &c.HeadersJSON); err != nil {
return fmt.Errorf("failed to resolve vault mcp headers: %w", err)
}
}
if c.ConnectionString != nil && !c.ConnectionString.IsFromEnv() && c.ConnectionString.GetValue() != "" {
if err := resolveVaultEnvVar(tx.Statement.Context, c.ConnectionString); err != nil {
return fmt.Errorf("failed to resolve vault mcp connection string: %w", err)
}
}
case EncryptionStatusEncrypted:
if c.EncryptionStatus == EncryptionStatusEncrypted {
if c.HeadersJSON != "" && c.HeadersJSON != "{}" {
decrypted, err := encrypt.Decrypt(c.HeadersJSON)
if err != nil {
Expand Down Expand Up @@ -329,15 +293,3 @@ func (c *TableMCPClient) AfterFind(tx *gorm.DB) error {
}
return nil
}

// AfterDelete hook for best-effort vault cleanup on row deletion.
func (c *TableMCPClient) AfterDelete(tx *gorm.DB) error {
if c.EncryptionStatus != EncryptionStatusVault || VaultHooks.Remove == nil {
return nil
}
connField := tx.Statement.DB.NamingStrategy.ColumnName("", "ConnectionString")
headersField := tx.Statement.DB.NamingStrategy.ColumnName("", "HeadersJSON")
_ = VaultHooks.Remove(tx.Statement.Context, fmt.Sprintf("%s/%s/%s/%s", VaultPrefix(), c.TableName(), c.ClientID, connField))
_ = VaultHooks.Remove(tx.Statement.Context, fmt.Sprintf("%s/%s/%s/%s", VaultPrefix(), c.TableName(), c.ClientID, headersField))
return nil
}
18 changes: 2 additions & 16 deletions framework/configstore/tables/mcp_per_user_headers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,14 +104,7 @@ func (c *TableMCPPerUserHeaderCredential) BeforeSave(tx *gorm.DB) error {
if c.HeadersJSON == "" {
c.HeadersJSON = "{}"
}
if VaultIsEnabled() && c.HeadersJSON != "{}" {
path := fmt.Sprintf("%s/%s/%s/%s", VaultPrefix(), c.TableName(), c.ID,
tx.Statement.DB.NamingStrategy.ColumnName("", "HeadersJSON"))
if err := vaultString(tx.Statement.Context, path, &c.HeadersJSON); err != nil {
return fmt.Errorf("failed to vault mcp per-user header credential headers: %w", err)
}
c.EncryptionStatus = EncryptionStatusVault
} else if encrypt.IsEnabled() {
if encrypt.IsEnabled() {
if err := encryptString(&c.HeadersJSON); err != nil {
return fmt.Errorf("failed to encrypt mcp per-user header credential headers: %w", err)
}
Expand All @@ -122,14 +115,7 @@ func (c *TableMCPPerUserHeaderCredential) BeforeSave(tx *gorm.DB) error {

// AfterFind decrypts HeadersJSON when the row is marked encrypted.
func (c *TableMCPPerUserHeaderCredential) AfterFind(tx *gorm.DB) error {
switch c.EncryptionStatus {
case EncryptionStatusVault:
if c.HeadersJSON != "" && c.HeadersJSON != "{}" {
if err := resolveVaultString(tx.Statement.Context, &c.HeadersJSON); err != nil {
return fmt.Errorf("failed to resolve vault mcp per-user header credential headers: %w", err)
}
}
case EncryptionStatusEncrypted:
if c.EncryptionStatus == EncryptionStatusEncrypted {
if err := decryptString(&c.HeadersJSON); err != nil {
return fmt.Errorf("failed to decrypt mcp per-user header credential headers: %w", err)
}
Expand Down
Loading
Loading