jfraeysd/fetch_ml
0
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

feat(crypto,auth): harden KMS and improve permission handling

Browse source 
KMS improvements:
- cache.go: add LRU eviction with memory-bounded caches
- provider.go: refactor provider initialization and key rotation
- tenant_keys.go: per-tenant key isolation with envelope encryption

Auth layer updates:
- hybrid.go: refine hybrid auth flow for API key + JWT
- permissions_loader.go: faster permission caching with hot-reload
- validator.go: stricter validation with detailed error messages

Security middleware:
- security.go: add rate limiting headers and CORS refinement

Testing and benchmarks:
- Add KMS cache and protocol unit tests
- Add KMS benchmark tests for encryption throughput
- Update KMS integration tests for tenant isolation
This commit is contained in:
Jeremie Fraeys 2026-03-12 12:04:32 -04:00
parent de83300962
commit 37c4d4e9c7
No known key found for this signature in database
11 changed files with 83 additions and 83 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
internal/auth/hybrid.go
View file

@ -273,13 +273,13 @@ func (h *HybridAuthStore) Close() error {
}
// GetDatabaseStats returns database statistics
func (h *HybridAuthStore) GetDatabaseStats(ctx context.Context) (map[string]interface{}, error) {
func (h *HybridAuthStore) GetDatabaseStats(ctx context.Context) (map[string]any, error) {
h.mu.RLock()
useDB := h.useDB
h.mu.RUnlock()
if !useDB {
return map[string]interface{}{
return map[string]any{
"store_type": "file",
"users": len(h.fileStore.APIKeys),
}, nil
@ -290,7 +290,7 @@ func (h *HybridAuthStore) GetDatabaseStats(ctx context.Context) (map[string]inte
return nil, err
}
return map[string]interface{}{
return map[string]any{
"store_type": "database",
"users": len(users),
"path": "db/fetch_ml.db",

10
internal/auth/permissions_loader.go
View file

@ -31,8 +31,8 @@ type GroupConfig struct {
// HierarchyConfig defines resource hierarchy
type HierarchyConfig struct {
Children map[string]interface{} `yaml:"children"`
Special map[string]string `yaml:"special"`
Children map[string]any `yaml:"children"`
Special map[string]string `yaml:"special"`
}
// DefaultsConfig defines default settings
@ -257,19 +257,19 @@ func (pm *PermissionManager) ValidatePermission(permission string) bool {
}
// GetPermissionHierarchy returns the hierarchy for a resource
func (pm *PermissionManager) GetPermissionHierarchy(resource string) map[string]interface{} {
func (pm *PermissionManager) GetPermissionHierarchy(resource string) map[string]any {
pm.mu.RLock()
defer pm.mu.RUnlock()
if !pm.loaded {
return make(map[string]interface{})
return make(map[string]any)
}
if hierarchy, exists := pm.config.Hierarchy[resource]; exists {
return hierarchy.Children
}
return make(map[string]interface{})
return make(map[string]any)
}
// Global permission manager instance

8
internal/auth/validator.go
View file

@ -83,14 +83,14 @@ func CheckConfigFilePermissions(configPath string) error {
}
// SanitizeConfig removes sensitive information for logging
func (c *Config) SanitizeConfig() map[string]interface{} {
sanitized := map[string]interface{}{
func (c *Config) SanitizeConfig() map[string]any {
sanitized := map[string]any{
"enabled": c.Enabled,
"users": make(map[string]interface{}),
"users": make(map[string]any),
}
for username := range c.APIKeys {
sanitized["users"].(map[string]interface{})[string(username)] = map[string]interface{}{
sanitized["users"].(map[string]any)[string(username)] = map[string]any{
"admin": c.APIKeys[username].Admin,
"hash": strings.Repeat("*", 8) + "...", // Show only prefix
}

4
internal/crypto/kms/cache.go
View file

@ -6,6 +6,8 @@ import (
"fmt"
"sync"
"time"
"github.com/jfraeys/fetch_ml/internal/crypto/kms/config"
)
// DEKCache implements an in-process cache for unwrapped DEKs per ADR-012.
@ -33,7 +35,7 @@ type cacheEntry struct {
}
// NewDEKCache creates a new DEK cache with the specified configuration.
func NewDEKCache(cfg CacheConfig) *DEKCache {
func NewDEKCache(cfg config.CacheConfig) *DEKCache {
return &DEKCache{
entries: make(map[string]*cacheEntry),
ttl: cfg.TTL,

21
internal/crypto/kms/provider.go
View file

@ -58,33 +58,22 @@ const (
// ProviderFactory creates KMS providers from configuration.
type ProviderFactory struct {
config Config
}
// Config aliases from config package.
type Config = config.Config
type VaultConfig = config.VaultConfig
type AWSConfig = config.AWSConfig
type CacheConfig = config.CacheConfig
// DefaultCacheConfig re-exports from config package.
func DefaultCacheConfig() CacheConfig {
return config.DefaultCacheConfig()
config config.Config
}
// NewProviderFactory creates a new provider factory with the given config.
func NewProviderFactory(cfg Config) *ProviderFactory {
func NewProviderFactory(cfg config.Config) *ProviderFactory {
return &ProviderFactory{config: cfg}
}
// CreateProvider instantiates a KMS provider based on the configuration.
func (f *ProviderFactory) CreateProvider() (KMSProvider, error) {
switch f.config.Provider {
case ProviderTypeVault:
case config.ProviderTypeVault:
return providers.NewVaultProvider(f.config.Vault)
case ProviderTypeAWS:
case config.ProviderTypeAWS:
return providers.NewAWSProvider(f.config.AWS)
case ProviderTypeMemory:
case config.ProviderTypeMemory:
return NewMemoryProvider(), nil
default:
return nil, fmt.Errorf("unsupported KMS provider: %s", f.config.Provider)

19
internal/crypto/tenant_keys.go
View file

@ -16,7 +16,8 @@ import (
"time"
"github.com/jfraeys/fetch_ml/internal/audit"
"github.com/jfraeys/fetch_ml/internal/crypto/kms"
kms "github.com/jfraeys/fetch_ml/internal/crypto/kms"
kmsconfig "github.com/jfraeys/fetch_ml/internal/crypto/kms/config"
)
// KeyHierarchy defines the tenant key structure
@ -32,15 +33,15 @@ type KeyHierarchy struct {
// TenantKeyManager manages per-tenant encryption keys using external KMS per ADR-012 through ADR-015.
// Root keys are stored in the KMS; DEKs are generated locally and cached.
type TenantKeyManager struct {
kms kms.KMSProvider // External KMS for root key operations
cache *kms.DEKCache // In-process DEK cache per ADR-012
config kms.Config // KMS configuration
ctx context.Context // Background context for operations
audit *audit.Logger // Audit logger for key operations per ADR-012
kms kms.KMSProvider // External KMS for root key operations
cache *kms.DEKCache // In-process DEK cache per ADR-012
config kmsconfig.Config // KMS configuration
ctx context.Context // Background context for operations
audit *audit.Logger // Audit logger for key operations per ADR-012
}
// NewTenantKeyManager creates a new tenant key manager with KMS integration.
func NewTenantKeyManager(provider kms.KMSProvider, cache *kms.DEKCache, config kms.Config, auditLogger *audit.Logger) *TenantKeyManager {
func NewTenantKeyManager(provider kms.KMSProvider, cache *kms.DEKCache, config kmsconfig.Config, auditLogger *audit.Logger) *TenantKeyManager {
return &TenantKeyManager{
kms: provider,
cache: cache,
@ -234,8 +235,8 @@ type WrappedDEK struct {
// This provides backward compatibility for existing tests.
func NewTestTenantKeyManager(auditLogger *audit.Logger) *TenantKeyManager {
provider := kms.NewMemoryProvider()
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
config := kms.Config{Provider: kms.ProviderTypeMemory}
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
config := kmsconfig.Config{Provider: kmsconfig.ProviderTypeMemory}
return NewTenantKeyManager(provider, cache, config, auditLogger)
}

4
internal/middleware/security.go
View file

@ -269,7 +269,7 @@ func AuditLogger(next http.Handler) http.Handler {
// Log security-relevant events
if statusCode >= 400 || method == "DELETE" || strings.Contains(path, "/admin") {
// Log to security audit system
logSecurityEvent(map[string]interface{}{
logSecurityEvent(map[string]any{
"timestamp": start.Unix(),
"client_ip": clientIP,
"method": method,
@ -333,7 +333,7 @@ func (rw *responseWriter) WriteHeader(code int) {
rw.ResponseWriter.WriteHeader(code)
}
func logSecurityEvent(event map[string]interface{}) {
func logSecurityEvent(event map[string]any) {
// Implementation would send to security monitoring system
// For now, just log (in production, use proper logging)
log.Printf(

9
tests/benchmarks/kms_benchmark_test.go
View file

@ -6,6 +6,7 @@ import (
"github.com/jfraeys/fetch_ml/internal/crypto"
"github.com/jfraeys/fetch_ml/internal/crypto/kms"
kmsconfig "github.com/jfraeys/fetch_ml/internal/crypto/kms/config"
)
// BenchmarkEncryptArtifact measures the full encryption pipeline performance.
@ -73,12 +74,12 @@ func BenchmarkMemoryProvider_Encrypt(b *testing.B) {
provider := kms.NewMemoryProvider()
defer provider.Close()
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
config := kms.Config{
Provider: kms.ProviderTypeMemory,
Cache: kms.DefaultCacheConfig(),
config := kmsconfig.Config{
Provider: kmsconfig.ProviderTypeMemory,
Cache: kmsconfig.DefaultCacheConfig(),
}
tkm := crypto.NewTenantKeyManager(provider, cache, config, nil)

13
tests/integration/kms_integration_test.go
View file

@ -9,6 +9,7 @@ import (
"github.com/jfraeys/fetch_ml/internal/crypto"
"github.com/jfraeys/fetch_ml/internal/crypto/kms"
kmsconfig "github.com/jfraeys/fetch_ml/internal/crypto/kms/config"
"github.com/jfraeys/fetch_ml/internal/crypto/kms/providers"
"github.com/testcontainers/testcontainers-go"
"github.com/testcontainers/testcontainers-go/wait"
@ -76,7 +77,7 @@ func TestVaultProvider_Integration(t *testing.T) {
}
// Create provider config
config := kms.VaultConfig{
config := kmsconfig.VaultConfig{
Address: endpoint,
AuthMethod: "token",
Token: "test-token",
@ -169,7 +170,7 @@ func TestAWSKMSProvider_Integration(t *testing.T) {
}
// Create provider config
config := kms.AWSConfig{
config := kmsconfig.AWSConfig{
Region: "us-east-1",
KeyAliasPrefix: "alias/test-fetchml",
Endpoint: endpoint,
@ -208,13 +209,13 @@ func TestTenantKeyManager_WithMemoryProvider(t *testing.T) {
defer provider.Close()
// Create DEK cache
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
// Create config
config := kms.Config{
Provider: kms.ProviderTypeMemory,
Cache: kms.DefaultCacheConfig(),
config := kmsconfig.Config{
Provider: kmsconfig.ProviderTypeMemory,
Cache: kmsconfig.DefaultCacheConfig(),
}
// Create TenantKeyManager

19
tests/unit/crypto/kms/cache_test.go
View file

@ -6,11 +6,12 @@ import (
"time"
"github.com/jfraeys/fetch_ml/internal/crypto/kms"
kmsconfig "github.com/jfraeys/fetch_ml/internal/crypto/kms/config"
)
// TestDEKCache_PutAndGet tests basic cache put and get operations.
func TestDEKCache_PutAndGet(t *testing.T) {
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
tenantID := "tenant-1"
@ -35,7 +36,7 @@ func TestDEKCache_PutAndGet(t *testing.T) {
// TestDEKCache_GetNonexistent tests getting a non-existent entry.
func TestDEKCache_GetNonexistent(t *testing.T) {
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
_, ok := cache.Get("nonexistent", "nonexistent", "kms-key-1", false)
@ -47,7 +48,7 @@ func TestDEKCache_GetNonexistent(t *testing.T) {
// TestDEKCache_TTLExpiry tests that entries expire after TTL.
func TestDEKCache_TTLExpiry(t *testing.T) {
// Use very short TTL for testing
config := kms.CacheConfig{
config := kmsconfig.CacheConfig{
TTL: 50 * time.Millisecond,
MaxEntries: 100,
GraceWindow: 100 * time.Millisecond,
@ -95,7 +96,7 @@ func TestDEKCache_TTLExpiry(t *testing.T) {
// TestDEKCache_LRUeviction tests LRU eviction when cache is full.
func TestDEKCache_LRUeviction(t *testing.T) {
config := kms.CacheConfig{
config := kmsconfig.CacheConfig{
TTL: 1 * time.Hour, // Long TTL so eviction is due to size
MaxEntries: 3,
GraceWindow: 1 * time.Hour,
@ -141,7 +142,7 @@ func TestDEKCache_LRUeviction(t *testing.T) {
// TestDEKCache_Flush tests flushing entries for a specific tenant.
func TestDEKCache_Flush(t *testing.T) {
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
// Add entries for two tenants
@ -171,7 +172,7 @@ func TestDEKCache_Flush(t *testing.T) {
// TestDEKCache_Clear tests clearing all entries.
func TestDEKCache_Clear(t *testing.T) {
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
// Add entries
cache.Put("tenant-1", "artifact-1", "kms-key-1", []byte("dek-1"))
@ -193,7 +194,7 @@ func TestDEKCache_Clear(t *testing.T) {
// TestDEKCache_Stats tests cache statistics.
func TestDEKCache_Stats(t *testing.T) {
config := kms.DefaultCacheConfig()
config := kmsconfig.DefaultCacheConfig()
cache := kms.NewDEKCache(config)
defer cache.Clear()
@ -223,7 +224,7 @@ func TestDEKCache_Stats(t *testing.T) {
// TestDEKCache_EmptyDEK tests that empty DEK is rejected.
func TestDEKCache_EmptyDEK(t *testing.T) {
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
err := cache.Put("tenant-1", "artifact-1", "kms-key-1", []byte{})
@ -234,7 +235,7 @@ func TestDEKCache_EmptyDEK(t *testing.T) {
// TestDEKCache_Isolation tests that DEKs are isolated between tenants.
func TestDEKCache_Isolation(t *testing.T) {
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
// Same artifact ID, different tenants

53
tests/unit/crypto/kms/protocol_test.go
View file

@ -10,6 +10,7 @@ import (
"github.com/jfraeys/fetch_ml/internal/api"
"github.com/jfraeys/fetch_ml/internal/crypto"
"github.com/jfraeys/fetch_ml/internal/crypto/kms"
kmsconfig "github.com/jfraeys/fetch_ml/internal/crypto/kms/config"
)
func TestProtocolSerialization(t *testing.T) {
@ -30,8 +31,8 @@ func TestProtocolSerialization(t *testing.T) {
t.Errorf("Expected at least 9 bytes, got %d", len(data))
}
// Test error packet
errorPacket := api.NewErrorPacket(api.ErrorCodeAuthenticationFailed, "Auth failed", "Invalid API key")
// Test error packet - uses string error code from errors package
errorPacket := api.NewErrorPacket("AUTHENTICATION_FAILED", "Auth failed", "Invalid API key")
data, err = errorPacket.Serialize()
if err != nil {
t.Fatalf("Failed to serialize error packet: %v", err)
@ -64,18 +65,22 @@ func TestProtocolSerialization(t *testing.T) {
}
}
func TestErrorMessageMapping(t *testing.T) {
tests := map[byte]string{
api.ErrorCodeUnknownError: "Unknown error occurred",
api.ErrorCodeAuthenticationFailed: "Authentication failed",
api.ErrorCodeJobNotFound: "Job not found",
api.ErrorCodeServerOverloaded: "Server is overloaded",
func TestByteCodeFromErrorCode(t *testing.T) {
tests := map[string]byte{
"UNKNOWN_ERROR": api.ErrorCodeUnknownError,
"AUTHENTICATION_FAILED": api.ErrorCodeAuthenticationFailed,
"JOB_NOT_FOUND": api.ErrorCodeJobNotFound,
"SERVER_OVERLOADED": api.ErrorCodeServerOverloaded,
"INVALID_REQUEST": api.ErrorCodeInvalidRequest,
"BAD_REQUEST": api.ErrorCodeInvalidRequest,
"PERMISSION_DENIED": api.ErrorCodePermissionDenied,
"FORBIDDEN": api.ErrorCodePermissionDenied,
}
for code, expected := range tests {
actual := api.GetErrorMessage(code)
if actual != expected {
t.Errorf("Expected error message '%s' for code %d, got '%s'", expected, code, actual)
for code, expectedByte := range tests {
actual := api.ByteCodeFromErrorCode(code)
if actual != expectedByte {
t.Errorf("Expected byte %d for code '%s', got %d", expectedByte, code, actual)
}
}
}
@ -125,12 +130,12 @@ func TestKMSProtocol_EncryptDecrypt(t *testing.T) {
provider := kms.NewMemoryProvider()
defer provider.Close()
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
config := kms.Config{
config := kmsconfig.Config{
Provider: kms.ProviderTypeMemory,
Cache: kms.DefaultCacheConfig(),
Cache: kmsconfig.DefaultCacheConfig(),
}
tkm := crypto.NewTenantKeyManager(provider, cache, config, nil)
@ -181,12 +186,12 @@ func TestKMSProtocol_MultiTenantIsolation(t *testing.T) {
provider := kms.NewMemoryProvider()
defer provider.Close()
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
config := kms.Config{
config := kmsconfig.Config{
Provider: kms.ProviderTypeMemory,
Cache: kms.DefaultCacheConfig(),
Cache: kmsconfig.DefaultCacheConfig(),
}
tkm := crypto.NewTenantKeyManager(provider, cache, config, nil)
@ -231,12 +236,12 @@ func TestKMSProtocol_CacheHit(t *testing.T) {
provider := kms.NewMemoryProvider()
defer provider.Close()
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
config := kms.Config{
config := kmsconfig.Config{
Provider: kms.ProviderTypeMemory,
Cache: kms.DefaultCacheConfig(),
Cache: kmsconfig.DefaultCacheConfig(),
}
tkm := crypto.NewTenantKeyManager(provider, cache, config, nil)
@ -271,12 +276,12 @@ func TestKMSProtocol_KeyRotation(t *testing.T) {
provider := kms.NewMemoryProvider()
defer provider.Close()
cache := kms.NewDEKCache(kms.DefaultCacheConfig())
cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
defer cache.Clear()
config := kms.Config{
config := kmsconfig.Config{
Provider: kms.ProviderTypeMemory,
Cache: kms.DefaultCacheConfig(),
Cache: kmsconfig.DefaultCacheConfig(),
}
tkm := crypto.NewTenantKeyManager(provider, cache, config, nil)