270 lines
11 KiB
Go
270 lines
11 KiB
Go
package tables
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/maximhq/bifrost/core/schemas"
|
|
"github.com/maximhq/bifrost/framework/encrypt"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
// TableVirtualKeyProviderConfigKey is the join table for the many2many relationship
|
|
// between TableVirtualKeyProviderConfig and TableKey
|
|
type TableVirtualKeyProviderConfigKey struct {
|
|
TableVirtualKeyProviderConfigID uint `gorm:"primaryKey;uniqueIndex:idx_vk_provider_config_key"`
|
|
TableKeyID uint `gorm:"primaryKey;uniqueIndex:idx_vk_provider_config_key"`
|
|
}
|
|
|
|
// TableName sets the table name for the join table
|
|
func (TableVirtualKeyProviderConfigKey) TableName() string {
|
|
return "governance_virtual_key_provider_config_keys"
|
|
}
|
|
|
|
// TableVirtualKeyProviderConfig represents a provider configuration for a virtual key
|
|
type TableVirtualKeyProviderConfig struct {
|
|
ID uint `gorm:"primaryKey;autoIncrement" json:"id"`
|
|
VirtualKeyID string `gorm:"type:varchar(255);not null" json:"virtual_key_id"`
|
|
Provider string `gorm:"type:varchar(50);not null" json:"provider"`
|
|
Weight *float64 `json:"weight"`
|
|
AllowedModels schemas.WhiteList `gorm:"type:text;serializer:json" json:"allowed_models"` // ["*"] allows all models; empty denies all (deny-by-default)
|
|
AllowAllKeys bool `gorm:"default:false" json:"allow_all_keys"` // True means all keys allowed; false with empty Keys means no keys allowed (deny-by-default)
|
|
RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"`
|
|
|
|
// Relationships
|
|
RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID;onDelete:CASCADE" json:"rate_limit,omitempty"`
|
|
Budgets []TableBudget `gorm:"foreignKey:ProviderConfigID;constraint:OnDelete:CASCADE" json:"budgets,omitempty"` // Multiple budgets with different reset intervals
|
|
Keys []TableKey `gorm:"many2many:governance_virtual_key_provider_config_keys;constraint:OnDelete:CASCADE" json:"keys"` // Empty means all keys allowed for this provider
|
|
}
|
|
|
|
// TableName sets the table name for each model
|
|
func (TableVirtualKeyProviderConfig) TableName() string {
|
|
return "governance_virtual_key_provider_configs"
|
|
}
|
|
|
|
// UnmarshalJSON custom unmarshaller to handle "key_ids" ([]string) config-file format
|
|
func (pc *TableVirtualKeyProviderConfig) UnmarshalJSON(data []byte) error {
|
|
type Alias TableVirtualKeyProviderConfig
|
|
type TempProviderConfig struct {
|
|
Alias
|
|
KeyIDs []string `json:"key_ids"` // Config file format: key identifiers (TableKey.KeyID); use ["*"] to allow all keys, empty denies all
|
|
}
|
|
|
|
var temp TempProviderConfig
|
|
if err := json.Unmarshal(data, &temp); err != nil {
|
|
return err
|
|
}
|
|
|
|
// Copy all standard fields
|
|
*pc = TableVirtualKeyProviderConfig(temp.Alias)
|
|
|
|
// If key_ids is provided, convert to Keys or set AllowAllKeys
|
|
if len(temp.KeyIDs) > 0 && len(pc.Keys) == 0 {
|
|
// ["*"] means allow all keys
|
|
if len(temp.KeyIDs) == 1 && temp.KeyIDs[0] == "*" {
|
|
pc.AllowAllKeys = true
|
|
pc.Keys = nil
|
|
} else {
|
|
pc.AllowAllKeys = false
|
|
pc.Keys = make([]TableKey, len(temp.KeyIDs))
|
|
for i, keyID := range temp.KeyIDs {
|
|
pc.Keys[i] = TableKey{KeyID: keyID}
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// BeforeSave validates WhiteList fields before GORM persists the record.
|
|
func (pc *TableVirtualKeyProviderConfig) BeforeSave(tx *gorm.DB) error {
|
|
if err := pc.AllowedModels.Validate(); err != nil {
|
|
return fmt.Errorf("invalid allowed_models: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// MarshalJSON custom marshaller to ensure AllowedModels is always an array (never null)
|
|
func (pc TableVirtualKeyProviderConfig) MarshalJSON() ([]byte, error) {
|
|
type Alias TableVirtualKeyProviderConfig
|
|
|
|
// Ensure AllowedModels is an empty slice instead of nil
|
|
allowedModels := pc.AllowedModels
|
|
if allowedModels == nil {
|
|
allowedModels = []string{}
|
|
}
|
|
|
|
return json.Marshal(&struct {
|
|
Alias
|
|
AllowedModels []string `json:"allowed_models"`
|
|
}{
|
|
Alias: Alias(pc),
|
|
AllowedModels: allowedModels,
|
|
})
|
|
}
|
|
|
|
// AfterFind hook for TableVirtualKeyProviderConfig to clear sensitive data from associated keys
|
|
func (pc *TableVirtualKeyProviderConfig) AfterFind(tx *gorm.DB) error {
|
|
if pc.Keys != nil {
|
|
// Clear sensitive data from associated keys, keeping only key IDs and non-sensitive metadata
|
|
for i := range pc.Keys {
|
|
key := &pc.Keys[i]
|
|
|
|
// Clear the actual API key value
|
|
key.Value = *schemas.NewEnvVar("")
|
|
|
|
// Clear all Azure-related sensitive fields
|
|
key.AzureEndpoint = nil
|
|
key.AzureAPIVersion = nil
|
|
key.AzureClientID = nil
|
|
key.AzureClientSecret = nil
|
|
key.AzureTenantID = nil
|
|
key.AzureScopesJSON = nil
|
|
key.AzureKeyConfig = nil
|
|
|
|
// Clear all Vertex-related sensitive fields
|
|
key.VertexProjectID = nil
|
|
key.VertexProjectNumber = nil
|
|
key.VertexRegion = nil
|
|
key.VertexAuthCredentials = nil
|
|
key.VertexKeyConfig = nil
|
|
|
|
// Clear all Bedrock-related sensitive fields
|
|
key.BedrockAccessKey = nil
|
|
key.BedrockSecretKey = nil
|
|
key.BedrockSessionToken = nil
|
|
key.BedrockRegion = nil
|
|
key.BedrockARN = nil
|
|
key.BedrockRoleARN = nil
|
|
key.BedrockExternalID = nil
|
|
key.BedrockRoleSessionName = nil
|
|
key.BedrockKeyConfig = nil
|
|
|
|
pc.Keys[i] = *key
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type TableVirtualKeyMCPConfig struct {
|
|
ID uint `gorm:"primaryKey;autoIncrement" json:"id"`
|
|
VirtualKeyID string `gorm:"type:varchar(255);not null;uniqueIndex:idx_vk_mcpclient" json:"virtual_key_id"`
|
|
MCPClientID uint `gorm:"not null;uniqueIndex:idx_vk_mcpclient" json:"mcp_client_id"`
|
|
MCPClient TableMCPClient `gorm:"foreignKey:MCPClientID" json:"mcp_client"`
|
|
ToolsToExecute schemas.WhiteList `gorm:"type:text;serializer:json" json:"tools_to_execute"`
|
|
|
|
// MCPClientName is used during config file parsing to resolve the MCP client by name.
|
|
// This field is not persisted to the database - it's only used to capture
|
|
// "mcp_client_name" from config.json and then resolve it to MCPClientID.
|
|
MCPClientName string `gorm:"-" json:"-"`
|
|
}
|
|
|
|
// TableName sets the table name for each model
|
|
func (TableVirtualKeyMCPConfig) TableName() string {
|
|
return "governance_virtual_key_mcp_configs"
|
|
}
|
|
|
|
// BeforeSave validates WhiteList fields before GORM persists the record.
|
|
func (mc *TableVirtualKeyMCPConfig) BeforeSave(tx *gorm.DB) error {
|
|
if err := mc.ToolsToExecute.Validate(); err != nil {
|
|
return fmt.Errorf("invalid tools_to_execute: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// UnmarshalJSON custom unmarshaller to handle both "mcp_client_id" (database format)
|
|
// and "mcp_client_name" (config file format) for MCP client references.
|
|
func (mc *TableVirtualKeyMCPConfig) UnmarshalJSON(data []byte) error {
|
|
// Temporary struct to capture all fields including mcp_client_name
|
|
type Alias TableVirtualKeyMCPConfig
|
|
type TempMCPConfig struct {
|
|
Alias
|
|
MCPClientName string `json:"mcp_client_name"` // Config file format: MCP client name
|
|
}
|
|
var temp TempMCPConfig
|
|
if err := json.Unmarshal(data, &temp); err != nil {
|
|
return err
|
|
}
|
|
// Copy all standard fields
|
|
*mc = TableVirtualKeyMCPConfig(temp.Alias)
|
|
// Capture mcp_client_name for later resolution to MCPClientID
|
|
if temp.MCPClientName != "" {
|
|
mc.MCPClientName = temp.MCPClientName
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// TableVirtualKey represents a virtual key with budget, rate limits, and team/customer association
|
|
type TableVirtualKey struct {
|
|
ID string `gorm:"primaryKey;type:varchar(255)" json:"id"`
|
|
Name string `gorm:"uniqueIndex:idx_virtual_key_name;type:varchar(255);not null" json:"name"`
|
|
Description string `gorm:"type:text" json:"description,omitempty"`
|
|
Value string `gorm:"uniqueIndex:idx_virtual_key_value;type:text;not null" json:"value"` // The virtual key value
|
|
IsActive bool `gorm:"default:true" json:"is_active"`
|
|
ProviderConfigs []TableVirtualKeyProviderConfig `gorm:"foreignKey:VirtualKeyID;constraint:OnDelete:CASCADE" json:"provider_configs"` // Empty means no providers allowed (deny-by-default)
|
|
MCPConfigs []TableVirtualKeyMCPConfig `gorm:"foreignKey:VirtualKeyID;constraint:OnDelete:CASCADE" json:"mcp_configs"`
|
|
|
|
// Foreign key relationships (mutually exclusive: either TeamID or CustomerID, not both)
|
|
TeamID *string `gorm:"type:varchar(255);index" json:"team_id,omitempty"`
|
|
CustomerID *string `gorm:"type:varchar(255);index" json:"customer_id,omitempty"`
|
|
RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"`
|
|
|
|
// Deprecated
|
|
// Calendar aligned is not the property of virtual key but its property of the budget and ratelimit
|
|
// So in the migration we will move this to the budget/ratelimit table
|
|
// And this won't be referred
|
|
CalendarAligned bool `gorm:"default:false" json:"calendar_aligned"` // When true, all budgets under this VK reset at clean calendar boundaries
|
|
|
|
// Relationships
|
|
Team *TableTeam `gorm:"foreignKey:TeamID" json:"team,omitempty"`
|
|
Customer *TableCustomer `gorm:"foreignKey:CustomerID" json:"customer,omitempty"`
|
|
RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID;onDelete:CASCADE" json:"rate_limit,omitempty"`
|
|
Budgets []TableBudget `gorm:"foreignKey:VirtualKeyID;constraint:OnDelete:CASCADE" json:"budgets,omitempty"` // Multiple budgets with different reset intervals
|
|
|
|
// Config hash is used to detect the changes synced from config.json file
|
|
// Every time we sync the config.json file, we will update the config hash
|
|
ConfigHash string `gorm:"type:varchar(255);null" json:"config_hash"`
|
|
|
|
EncryptionStatus string `gorm:"type:varchar(20);default:'plain_text'" json:"-"`
|
|
ValueHash string `gorm:"type:varchar(64);index:idx_virtual_key_value_hash,unique" json:"-"`
|
|
|
|
CreatedAt time.Time `gorm:"index;not null" json:"created_at"`
|
|
UpdatedAt time.Time `gorm:"index;not null" json:"updated_at"`
|
|
}
|
|
|
|
// TableName sets the table name for each model
|
|
func (TableVirtualKey) TableName() string { return "governance_virtual_keys" }
|
|
|
|
// BeforeSave is a GORM hook that enforces mutual exclusion (team vs customer), computes
|
|
// a SHA-256 hash of the plaintext value for indexed lookups, and encrypts the virtual key
|
|
// value before writing to the database.
|
|
func (vk *TableVirtualKey) BeforeSave(tx *gorm.DB) error {
|
|
// Enforce mutual exclusion: VK can belong to either Team OR Customer, not both
|
|
if vk.TeamID != nil && vk.CustomerID != nil {
|
|
return fmt.Errorf("virtual key cannot belong to both team and customer")
|
|
}
|
|
|
|
// Hash must be computed before encryption (from plaintext value)
|
|
if vk.Value != "" {
|
|
vk.ValueHash = encrypt.HashSHA256(vk.Value)
|
|
}
|
|
if encrypt.IsEnabled() && vk.Value != "" {
|
|
if err := encryptString(&vk.Value); err != nil {
|
|
return fmt.Errorf("failed to encrypt virtual key value: %w", err)
|
|
}
|
|
vk.EncryptionStatus = EncryptionStatusEncrypted
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// AfterFind is a GORM hook that decrypts the virtual key value after reading from the database.
|
|
func (vk *TableVirtualKey) AfterFind(tx *gorm.DB) error {
|
|
if vk.EncryptionStatus == EncryptionStatusEncrypted {
|
|
if err := decryptString(&vk.Value); err != nil {
|
|
return fmt.Errorf("failed to decrypt virtual key value: %w", err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|