alanyeung-s3/aws/internal/kvdb/kvdb.go

package kvdb

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sync"
)

// InMemoryKVDB is a simple in-memory key-value database
type InMemoryKVDB struct {
	users         map[string]*User         // key: AccessKeyID
	bucketConfigs map[string]*BucketConfig // key: accountID:bucketID
	mu            sync.RWMutex
}

// NewInMemoryKVDB creates a new in-memory key-value database
func NewInMemoryKVDB() *InMemoryKVDB {
	return &InMemoryKVDB{
		users:         make(map[string]*User),
		bucketConfigs: make(map[string]*BucketConfig),
	}
}

// GetUser retrieves a user by access key ID
func (db *InMemoryKVDB) GetUser(accessKeyID string) (*User, error) {
	db.mu.RLock()
	defer db.mu.RUnlock()

	user, exists := db.users[accessKeyID]
	if !exists {
		return nil, ErrUserNotFound
	}

	// Return a copy to prevent external modifications
	userCopy := *user
	return &userCopy, nil
}

// CreateUser creates a new user
func (db *InMemoryKVDB) CreateUser(user *User) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	if _, exists := db.users[user.AccessKeyID]; exists {
		return ErrUserAlreadyExists
	}

	// Store a copy to prevent external modifications
	userCopy := *user
	db.users[user.AccessKeyID] = &userCopy
	return nil
}

// UpdateUser updates an existing user
func (db *InMemoryKVDB) UpdateUser(user *User) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	if _, exists := db.users[user.AccessKeyID]; !exists {
		return ErrUserNotFound
	}

	userCopy := *user
	db.users[user.AccessKeyID] = &userCopy
	return nil
}

// DeleteUser deletes a user
func (db *InMemoryKVDB) DeleteUser(accessKeyID string) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	if _, exists := db.users[accessKeyID]; !exists {
		return ErrUserNotFound
	}

	delete(db.users, accessKeyID)
	return nil
}

// ListUsers returns all users
func (db *InMemoryKVDB) ListUsers() ([]*User, error) {
	db.mu.RLock()
	defer db.mu.RUnlock()

	users := make([]*User, 0, len(db.users))
	for _, user := range db.users {
		userCopy := *user
		users = append(users, &userCopy)
	}

	return users, nil
}

// ValidateCredentials validates user credentials
func (db *InMemoryKVDB) ValidateCredentials(accessKeyID, secretAccessKey string) (*User, error) {
	user, err := db.GetUser(accessKeyID)
	if err != nil {
		return nil, err
	}

	if user.SecretAccessKey != secretAccessKey {
		return nil, errors.New("invalid credentials")
	}

	return user, nil
}

// SetBucketConfig sets the configuration for a bucket
func (db *InMemoryKVDB) SetBucketConfig(config *BucketConfig) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	key := config.AccountID + ":" + config.BucketID
	configCopy := *config
	db.bucketConfigs[key] = &configCopy
	return nil
}

// GetBucketConfig gets the configuration for a bucket
func (db *InMemoryKVDB) GetBucketConfig(accountID, bucketID string) (*BucketConfig, error) {
	db.mu.RLock()
	defer db.mu.RUnlock()

	key := accountID + ":" + bucketID
	config, exists := db.bucketConfigs[key]
	if !exists {
		return nil, fmt.Errorf("bucket config not found")
	}

	configCopy := *config
	return &configCopy, nil
}

// DeleteBucketConfig deletes the configuration for a bucket
func (db *InMemoryKVDB) DeleteBucketConfig(accountID, bucketID string) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	key := accountID + ":" + bucketID
	if _, exists := db.bucketConfigs[key]; !exists {
		return fmt.Errorf("bucket config not found")
	}

	delete(db.bucketConfigs, key)
	return nil
}

// ListBucketConfigs lists all bucket configurations for an account
func (db *InMemoryKVDB) ListBucketConfigs(accountID string) ([]*BucketConfig, error) {
	db.mu.RLock()
	defer db.mu.RUnlock()

	var configs []*BucketConfig
	prefix := accountID + ":"

	for key, config := range db.bucketConfigs {
		if len(key) >= len(prefix) && key[:len(prefix)] == prefix {
			configCopy := *config
			configs = append(configs, &configCopy)
		}
	}

	return configs, nil
}

// ResolveBucketName resolves a bucket name to accountID and bucketID
func (db *InMemoryKVDB) ResolveBucketName(bucketName string) (string, string, error) {
	db.mu.RLock()
	defer db.mu.RUnlock()

	for _, config := range db.bucketConfigs {
		if config.BucketName == bucketName {
			return config.AccountID, config.BucketID, nil
		}
	}

	return "", "", fmt.Errorf("bucket not found: %s", bucketName)
}

// Close closes the database (no-op for in-memory)
func (db *InMemoryKVDB) Close() error {
	return nil
}

// FileBasedKVDB is a file-based key-value database with persistence
type FileBasedKVDB struct {
	filePath      string
	users         map[string]*User
	bucketConfigs map[string]*BucketConfig
	mu            sync.RWMutex
}

// NewFileBasedKVDB creates a new file-based key-value database
func NewFileBasedKVDB(filePath string) (*FileBasedKVDB, error) {
	db := &FileBasedKVDB{
		filePath:      filePath,
		users:         make(map[string]*User),
		bucketConfigs: make(map[string]*BucketConfig),
	}

	// Create directory if it doesn't exist
	dir := filepath.Dir(filePath)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, err
	}

	// Load existing data if file exists
	if err := db.load(); err != nil && !os.IsNotExist(err) {
		return nil, err
	}

	return db, nil
}

type fileBasedData struct {
	Users         []*User         `json:"users"`
	BucketConfigs []*BucketConfig `json:"bucket_configs"`
}

// load reads the database from disk
func (db *FileBasedKVDB) load() error {
	data, err := os.ReadFile(db.filePath)
	if err != nil {
		return err
	}

	var fileData fileBasedData
	if err := json.Unmarshal(data, &fileData); err != nil {
		return err
	}

	db.mu.Lock()
	defer db.mu.Unlock()

	for _, user := range fileData.Users {
		db.users[user.AccessKeyID] = user
	}

	for _, config := range fileData.BucketConfigs {
		key := config.AccountID + ":" + config.BucketID
		db.bucketConfigs[key] = config
	}

	return nil
}

// save writes the database to disk
func (db *FileBasedKVDB) save() error {
	users := make([]*User, 0, len(db.users))
	for _, user := range db.users {
		users = append(users, user)
	}

	configs := make([]*BucketConfig, 0, len(db.bucketConfigs))
	for _, config := range db.bucketConfigs {
		configs = append(configs, config)
	}

	fileData := fileBasedData{
		Users:         users,
		BucketConfigs: configs,
	}

	data, err := json.MarshalIndent(fileData, "", "  ")
	if err != nil {
		return err
	}

	return os.WriteFile(db.filePath, data, 0644)
}

// GetUser retrieves a user by access key ID
func (db *FileBasedKVDB) GetUser(accessKeyID string) (*User, error) {
	db.mu.RLock()
	defer db.mu.RUnlock()

	user, exists := db.users[accessKeyID]
	if !exists {
		return nil, ErrUserNotFound
	}

	userCopy := *user
	return &userCopy, nil
}

// CreateUser creates a new user
func (db *FileBasedKVDB) CreateUser(user *User) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	if _, exists := db.users[user.AccessKeyID]; exists {
		return ErrUserAlreadyExists
	}

	userCopy := *user
	db.users[user.AccessKeyID] = &userCopy

	return db.save()
}

// UpdateUser updates an existing user
func (db *FileBasedKVDB) UpdateUser(user *User) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	if _, exists := db.users[user.AccessKeyID]; !exists {
		return ErrUserNotFound
	}

	userCopy := *user
	db.users[user.AccessKeyID] = &userCopy

	return db.save()
}

// DeleteUser deletes a user
func (db *FileBasedKVDB) DeleteUser(accessKeyID string) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	if _, exists := db.users[accessKeyID]; !exists {
		return ErrUserNotFound
	}

	delete(db.users, accessKeyID)

	return db.save()
}

// ListUsers returns all users
func (db *FileBasedKVDB) ListUsers() ([]*User, error) {
	db.mu.RLock()
	defer db.mu.RUnlock()

	users := make([]*User, 0, len(db.users))
	for _, user := range db.users {
		userCopy := *user
		users = append(users, &userCopy)
	}

	return users, nil
}

// ValidateCredentials validates user credentials
func (db *FileBasedKVDB) ValidateCredentials(accessKeyID, secretAccessKey string) (*User, error) {
	user, err := db.GetUser(accessKeyID)
	if err != nil {
		return nil, err
	}

	if user.SecretAccessKey != secretAccessKey {
		return nil, errors.New("invalid credentials")
	}

	return user, nil
}

// SetBucketConfig sets the configuration for a bucket
func (db *FileBasedKVDB) SetBucketConfig(config *BucketConfig) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	key := config.AccountID + ":" + config.BucketID
	configCopy := *config
	db.bucketConfigs[key] = &configCopy

	return db.save()
}

// GetBucketConfig gets the configuration for a bucket
func (db *FileBasedKVDB) GetBucketConfig(accountID, bucketID string) (*BucketConfig, error) {
	db.mu.RLock()
	defer db.mu.RUnlock()

	key := accountID + ":" + bucketID
	config, exists := db.bucketConfigs[key]
	if !exists {
		return nil, fmt.Errorf("bucket config not found")
	}

	configCopy := *config
	return &configCopy, nil
}

// DeleteBucketConfig deletes the configuration for a bucket
func (db *FileBasedKVDB) DeleteBucketConfig(accountID, bucketID string) error {
	db.mu.Lock()
	defer db.mu.Unlock()

	key := accountID + ":" + bucketID
	if _, exists := db.bucketConfigs[key]; !exists {
		return fmt.Errorf("bucket config not found")
	}

	delete(db.bucketConfigs, key)

	return db.save()
}

// ListBucketConfigs lists all bucket configurations for an account
func (db *FileBasedKVDB) ListBucketConfigs(accountID string) ([]*BucketConfig, error) {
	db.mu.RLock()
	defer db.mu.RUnlock()

	var configs []*BucketConfig
	prefix := accountID + ":"

	for key, config := range db.bucketConfigs {
		if len(key) >= len(prefix) && key[:len(prefix)] == prefix {
			configCopy := *config
			configs = append(configs, &configCopy)
		}
	}

	return configs, nil
}

// ResolveBucketName resolves a bucket name to accountID and bucketID
func (db *FileBasedKVDB) ResolveBucketName(bucketName string) (string, string, error) {
	db.mu.RLock()
	defer db.mu.RUnlock()

	for _, config := range db.bucketConfigs {
		if config.BucketName == bucketName {
			return config.AccountID, config.BucketID, nil
		}
	}

	return "", "", fmt.Errorf("bucket not found: %s", bucketName)
}

// Close closes the database
func (db *FileBasedKVDB) Close() error {
	db.mu.Lock()
	defer db.mu.Unlock()

	return db.save()
}