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

package kvdb

import (
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"go.etcd.io/bbolt"
)

var (
	usersBucket   = []byte("users")
	bucketsBucket = []byte("buckets")
)

// Ensure BoltKVDB implements KVDB interface
var _ KVDB = (*BoltKVDB)(nil)

// BucketConfig stores bucket-level configuration
type BucketConfig struct {
	BucketID      string `json:"bucket_id"`
	AccountID     string `json:"account_id"`
	BucketName    string `json:"bucket_name"`
	PublicViewing bool   `json:"public_viewing"`
	CreatedAt     string `json:"created_at"`
}

// BoltKVDB is a BoltDB-based key-value database
type BoltKVDB struct {
	db *bbolt.DB
}

// NewBoltKVDB creates a new BoltDB-based key-value database
func NewBoltKVDB(path string) (*BoltKVDB, error) {
	db, err := bbolt.Open(path, 0600, &bbolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		return nil, err
	}

	// Create buckets if they don't exist
	err = db.Update(func(tx *bbolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(usersBucket)
		if err != nil {
			return err
		}
		_, err = tx.CreateBucketIfNotExists(bucketsBucket)
		return err
	})
	if err != nil {
		db.Close()
		return nil, err
	}

	return &BoltKVDB{db: db}, nil
}

// GetUser retrieves a user by access key ID
func (db *BoltKVDB) GetUser(accessKeyID string) (*User, error) {
	var user User
	err := db.db.View(func(tx *bbolt.Tx) error {
		b := tx.Bucket(usersBucket)
		data := b.Get([]byte(accessKeyID))
		if data == nil {
			return ErrUserNotFound
		}
		return json.Unmarshal(data, &user)
	})
	if err != nil {
		return nil, err
	}
	return &user, nil
}

// CreateUser creates a new user
func (db *BoltKVDB) CreateUser(user *User) error {
	return db.db.Update(func(tx *bbolt.Tx) error {
		b := tx.Bucket(usersBucket)

		// Check if user already exists
		if b.Get([]byte(user.AccessKeyID)) != nil {
			return ErrUserAlreadyExists
		}

		data, err := json.Marshal(user)
		if err != nil {
			return err
		}

		return b.Put([]byte(user.AccessKeyID), data)
	})
}

// UpdateUser updates an existing user
func (db *BoltKVDB) UpdateUser(user *User) error {
	return db.db.Update(func(tx *bbolt.Tx) error {
		b := tx.Bucket(usersBucket)

		// Check if user exists
		if b.Get([]byte(user.AccessKeyID)) == nil {
			return ErrUserNotFound
		}

		data, err := json.Marshal(user)
		if err != nil {
			return err
		}

		return b.Put([]byte(user.AccessKeyID), data)
	})
}

// DeleteUser deletes a user
func (db *BoltKVDB) DeleteUser(accessKeyID string) error {
	return db.db.Update(func(tx *bbolt.Tx) error {
		b := tx.Bucket(usersBucket)

		// Check if user exists
		if b.Get([]byte(accessKeyID)) == nil {
			return ErrUserNotFound
		}

		return b.Delete([]byte(accessKeyID))
	})
}

// ListUsers returns all users
func (db *BoltKVDB) ListUsers() ([]*User, error) {
	var users []*User
	err := db.db.View(func(tx *bbolt.Tx) error {
		b := tx.Bucket(usersBucket)
		return b.ForEach(func(k, v []byte) error {
			var user User
			if err := json.Unmarshal(v, &user); err != nil {
				return err
			}
			users = append(users, &user)
			return nil
		})
	})
	return users, err
}

// ValidateCredentials validates user credentials
func (db *BoltKVDB) 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
}

// Close closes the database
func (db *BoltKVDB) Close() error {
	return db.db.Close()
}

// Bucket configuration methods

// SetBucketConfig sets the configuration for a bucket
func (db *BoltKVDB) SetBucketConfig(config *BucketConfig) error {
	return db.db.Update(func(tx *bbolt.Tx) error {
		b := tx.Bucket(bucketsBucket)

		// Key format: accountID:bucketID
		key := config.AccountID + ":" + config.BucketID

		data, err := json.Marshal(config)
		if err != nil {
			return err
		}

		return b.Put([]byte(key), data)
	})
}

// GetBucketConfig gets the configuration for a bucket
func (db *BoltKVDB) GetBucketConfig(accountID, bucketID string) (*BucketConfig, error) {
	var config BucketConfig

	err := db.db.View(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket(bucketsBucket)
		if bucket == nil {
			return fmt.Errorf("buckets bucket not found")
		}

		key := fmt.Sprintf("%s:%s", accountID, bucketID)
		data := bucket.Get([]byte(key))
		if data == nil {
			return fmt.Errorf("bucket config not found")
		}

		return json.Unmarshal(data, &config)
	})

	if err != nil {
		return nil, err
	}

	return &config, nil
}

// DeleteBucketConfig deletes the configuration for a bucket
func (db *BoltKVDB) DeleteBucketConfig(accountID, bucketID string) error {
	return db.db.Update(func(tx *bbolt.Tx) error {
		b := tx.Bucket(bucketsBucket)
		key := accountID + ":" + bucketID
		return b.Delete([]byte(key))
	})
}

// ListBucketConfigs lists all bucket configurations for an account
func (db *BoltKVDB) ListBucketConfigs(accountID string) ([]*BucketConfig, error) {
	var configs []*BucketConfig
	prefix := []byte(accountID + ":")

	err := db.db.View(func(tx *bbolt.Tx) error {
		b := tx.Bucket(bucketsBucket)
		c := b.Cursor()

		for k, v := c.Seek(prefix); k != nil && len(k) >= len(prefix) && string(k[:len(prefix)]) == string(prefix); k, v = c.Next() {
			var config BucketConfig
			if err := json.Unmarshal(v, &config); err != nil {
				continue
			}
			configs = append(configs, &config)
		}
		return nil
	})

	return configs, err
}

// ResolveBucketName resolves a bucket name to accountID and bucketID
func (db *BoltKVDB) ResolveBucketName(bucketName string) (accountID string, bucketID string, err error) {
	err = db.db.View(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket(bucketsBucket)
		if bucket == nil {
			return fmt.Errorf("buckets bucket not found")
		}

		// Iterate through all bucket configs to find matching bucket name
		cursor := bucket.Cursor()
		for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
			var config BucketConfig
			if err := json.Unmarshal(v, &config); err != nil {
				continue
			}

			if config.BucketName == bucketName {
				accountID = config.AccountID
				bucketID = config.BucketID
				return nil
			}
		}

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

	if err != nil {
		return "", "", err
	}

	return accountID, bucketID, nil
}

// InitializeDefaultUsers initializes hardcoded users
func (db *BoltKVDB) InitializeDefaultUsers() error {
	defaultUsers := []*User{
		{
			AccessKeyID:     "AKIAIOSFODNN7EXAMPLE",
			SecretAccessKey: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
			AccountID:       "123456789012",
			Username:        "alice",
			Email:           "[email protected]",
		},
		{
			AccessKeyID:     "AKIAJSIE27KKMHXI3BJQ",
			SecretAccessKey: "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",
			AccountID:       "987654321098",
			Username:        "bob",
			Email:           "[email protected]",
		},
	}

	for _, user := range defaultUsers {
		// Try to create user, ignore if already exists
		err := db.CreateUser(user)
		if err != nil && err != ErrUserAlreadyExists {
			return err
		}
	}

	return nil
}