alanyeung-s3/aws/internal/storage/user_aware.go

package storage

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

// UserAwareStorage wraps Storage with user isolation using accountID/bucketID structure
type UserAwareStorage struct {
	baseDir   string
	multipart *MultipartStorage
	mutex     sync.RWMutex
}

// NewUserAwareStorage creates a storage with user isolation
func NewUserAwareStorage(baseDir string) (*UserAwareStorage, error) {
	if err := os.MkdirAll(baseDir, 0755); err != nil {
		return nil, err
	}

	return &UserAwareStorage{
		baseDir:   baseDir,
		multipart: NewMultipartStorage(),
	}, nil
}

// getUserDir returns the directory for a specific user
func (s *UserAwareStorage) getUserDir(accountID string) string {
	return filepath.Join(s.baseDir, accountID)
}

// getBucketID generates a bucket ID from bucket name
func (s *UserAwareStorage) getBucketID(bucketName string) string {
	// Use MD5 hash of bucket name as bucket ID
	// This ensures consistent bucket IDs for same bucket names
	return bucketName
}

// initUserStorage initializes storage for a user
func (s *UserAwareStorage) initUserStorage(accountID string) error {
	userDir := s.getUserDir(accountID)
	return os.MkdirAll(userDir, 0755)
}

// CreateBucketForUser creates a bucket for a specific user
func (s *UserAwareStorage) CreateBucketForUser(accountID, bucketName string) (string, error) {
	if err := s.initUserStorage(accountID); err != nil {
		return "", err
	}

	bucketID := s.getBucketID(bucketName)
	bucketPath := filepath.Join(s.getUserDir(accountID), bucketID)

	if _, err := os.Stat(bucketPath); err == nil {
		return "", ErrBucketAlreadyExists
	}

	if err := os.MkdirAll(bucketPath, 0755); err != nil {
		return "", err
	}

	return bucketID, nil
}

// DeleteBucketForUser deletes a bucket for a specific user
func (s *UserAwareStorage) DeleteBucketForUser(accountID, bucketName string) error {
	bucketID := s.getBucketID(bucketName)
	bucketPath := filepath.Join(s.getUserDir(accountID), bucketID)
	return os.RemoveAll(bucketPath)
}

// BucketExistsForUser checks if a bucket exists for a specific user
func (s *UserAwareStorage) BucketExistsForUser(accountID, bucketName string) (bool, error) {
	bucketID := s.getBucketID(bucketName)
	bucketPath := filepath.Join(s.getUserDir(accountID), bucketID)
	_, err := os.Stat(bucketPath)
	if os.IsNotExist(err) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return true, nil
}

// GetBucketIDForUser gets the bucket ID for a bucket name
func (s *UserAwareStorage) GetBucketIDForUser(accountID, bucketName string) (string, error) {
	bucketID := s.getBucketID(bucketName)
	bucketPath := filepath.Join(s.getUserDir(accountID), bucketID)

	if _, err := os.Stat(bucketPath); os.IsNotExist(err) {
		return "", ErrBucketNotFound
	}

	return bucketID, nil
}

// ListBucketsForUser lists all buckets for a specific user
// This method doesn't work well with the new structure since we only have bucket IDs
// We need to rely on the KV database for bucket name mappings
func (s *UserAwareStorage) ListBucketsForUser(accountID string) ([]BucketInfo, error) {
	userDir := s.getUserDir(accountID)

	// If user directory doesn't exist, return empty list
	if _, err := os.Stat(userDir); os.IsNotExist(err) {
		return []BucketInfo{}, nil
	}

	entries, err := os.ReadDir(userDir)
	if err != nil {
		return nil, err
	}

	var buckets []BucketInfo
	for _, entry := range entries {
		if entry.IsDir() {
			info, err := entry.Info()
			if err != nil {
				continue
			}
			// entry.Name() is the bucket ID (hash)
			buckets = append(buckets, BucketInfo{
				Name:         entry.Name(), // This is bucket ID, not bucket name
				CreationDate: info.ModTime(),
			})
		}
	}

	return buckets, nil
}

// PutObjectForUser stores an object for a specific user
func (s *UserAwareStorage) PutObjectForUser(accountID, bucketName, key string, reader io.Reader) error {
	bucketID := s.getBucketID(bucketName)
	bucketPath := filepath.Join(s.getUserDir(accountID), bucketID)

	if _, err := os.Stat(bucketPath); os.IsNotExist(err) {
		return ErrBucketNotFound
	}

	objectPath := filepath.Join(bucketPath, key)

	objectDir := filepath.Dir(objectPath)
	if err := os.MkdirAll(objectDir, 0755); err != nil {
		return err
	}

	file, err := os.Create(objectPath)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, reader)
	return err
}

// GetObjectForUser retrieves an object for a specific user
func (s *UserAwareStorage) GetObjectForUser(accountID, bucketName, key string) (io.ReadCloser, *ObjectInfo, error) {
	bucketID := s.getBucketID(bucketName)
	objectPath := filepath.Join(s.getUserDir(accountID), bucketID, key)

	fileInfo, err := os.Stat(objectPath)
	if os.IsNotExist(err) {
		return nil, nil, ErrObjectNotFound
	}
	if err != nil {
		return nil, nil, err
	}

	file, err := os.Open(objectPath)
	if err != nil {
		return nil, nil, err
	}

	info := &ObjectInfo{
		Key:          key,
		Size:         fileInfo.Size(),
		LastModified: fileInfo.ModTime(),
		ETag:         generateETag(fileInfo),
	}

	return file, info, nil
}

// GetObjectByBucketIDForUser retrieves an object using bucket ID directly
func (s *UserAwareStorage) GetObjectByBucketIDForUser(accountID, bucketID, key string) (io.ReadCloser, *ObjectInfo, error) {
	objectPath := filepath.Join(s.getUserDir(accountID), bucketID, key)

	fileInfo, err := os.Stat(objectPath)
	if os.IsNotExist(err) {
		return nil, nil, ErrObjectNotFound
	}
	if err != nil {
		return nil, nil, err
	}

	file, err := os.Open(objectPath)
	if err != nil {
		return nil, nil, err
	}

	info := &ObjectInfo{
		Key:          key,
		Size:         fileInfo.Size(),
		LastModified: fileInfo.ModTime(),
		ETag:         generateETag(fileInfo),
	}

	return file, info, nil
}

// DeleteObjectForUser removes an object for a specific user
func (s *UserAwareStorage) DeleteObjectForUser(accountID, bucketName, key string) error {
	bucketID := s.getBucketID(bucketName)
	objectPath := filepath.Join(s.getUserDir(accountID), bucketID, key)
	err := os.Remove(objectPath)
	if os.IsNotExist(err) {
		return nil
	}
	return err
}

// ObjectExistsForUser checks if an object exists for a specific user
func (s *UserAwareStorage) ObjectExistsForUser(accountID, bucketName, key string) (bool, error) {
	bucketID := s.getBucketID(bucketName)
	objectPath := filepath.Join(s.getUserDir(accountID), bucketID, key)
	_, err := os.Stat(objectPath)
	if os.IsNotExist(err) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return true, nil
}

// GetObjectInfoForUser retrieves object metadata for a specific user
func (s *UserAwareStorage) GetObjectInfoForUser(accountID, bucketName, key string) (*ObjectInfo, error) {
	bucketID := s.getBucketID(bucketName)
	objectPath := filepath.Join(s.getUserDir(accountID), bucketID, key)

	fileInfo, err := os.Stat(objectPath)
	if os.IsNotExist(err) {
		return nil, ErrObjectNotFound
	}
	if err != nil {
		return nil, err
	}

	return &ObjectInfo{
		Key:          key,
		Size:         fileInfo.Size(),
		LastModified: fileInfo.ModTime(),
		ETag:         generateETag(fileInfo),
	}, nil
}

// ListObjectsForUser lists all objects in a bucket for a specific user
func (s *UserAwareStorage) ListObjectsForUser(accountID, bucketName string) ([]ObjectInfo, error) {
	bucketID := s.getBucketID(bucketName)
	bucketPath := filepath.Join(s.getUserDir(accountID), bucketID)

	if _, err := os.Stat(bucketPath); os.IsNotExist(err) {
		return nil, ErrBucketNotFound
	}

	var objects []ObjectInfo

	// Walk the bucket directory recursively
	err := filepath.Walk(bucketPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil // Skip errors
		}

		if !info.IsDir() {
			// Get relative path from bucket
			relPath, err := filepath.Rel(bucketPath, path)
			if err != nil {
				return nil
			}

			// Convert path separators to forward slashes for S3 compatibility
			key := filepath.ToSlash(relPath)

			objects = append(objects, ObjectInfo{
				Key:          key,
				Size:         info.Size(),
				LastModified: info.ModTime(),
				ETag:         generateETag(info),
			})
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	return objects, nil
}

// ListObjectsByBucketIDForUser lists all objects using bucket ID directly
func (s *UserAwareStorage) ListObjectsByBucketIDForUser(accountID, bucketID string) ([]ObjectInfo, error) {
	bucketPath := filepath.Join(s.getUserDir(accountID), bucketID)

	if _, err := os.Stat(bucketPath); os.IsNotExist(err) {
		return nil, ErrBucketNotFound
	}

	var objects []ObjectInfo

	// Walk the bucket directory recursively
	err := filepath.Walk(bucketPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil // Skip errors
		}

		if !info.IsDir() {
			// Get relative path from bucket
			relPath, err := filepath.Rel(bucketPath, path)
			if err != nil {
				return nil
			}

			// Convert path separators to forward slashes for S3 compatibility
			key := filepath.ToSlash(relPath)

			objects = append(objects, ObjectInfo{
				Key:          key,
				Size:         info.Size(),
				LastModified: info.ModTime(),
				ETag:         generateETag(info),
			})
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	return objects, nil
}

// GetStorageStatsForUser returns storage statistics for a user
func (s *UserAwareStorage) GetStorageStatsForUser(accountID string) (*StorageStats, error) {
	userDir := s.getUserDir(accountID)

	stats := &StorageStats{
		AccountID: accountID,
	}

	if _, err := os.Stat(userDir); os.IsNotExist(err) {
		return stats, nil
	}

	// Walk the user directory to calculate stats
	entries, err := os.ReadDir(userDir)
	if err != nil {
		return stats, err
	}

	for _, entry := range entries {
		if entry.IsDir() {
			stats.BucketCount++

			// Count objects in this bucket
			bucketPath := filepath.Join(userDir, entry.Name())
			filepath.Walk(bucketPath, func(path string, info os.FileInfo, err error) error {
				if err != nil || info.IsDir() {
					return nil
				}
				stats.ObjectCount++
				stats.TotalSize += info.Size()
				return nil
			})
		}
	}

	return stats, nil
}

// Helper function to sanitize paths and prevent directory traversal
func sanitizePath(path string) string {
	// Remove any ../ or ./ patterns
	path = strings.ReplaceAll(path, "../", "")
	path = strings.ReplaceAll(path, "./", "")
	return path
}

// Add these methods to UserAwareStorage

// CreateMultipartUpload initiates a multipart upload
func (s *UserAwareStorage) CreateMultipartUpload(accountID, bucket, key string) (string, error) {
	return s.multipart.CreateMultipartUpload(accountID, bucket, key)
}

// UploadPart uploads a part for a multipart upload
func (s *UserAwareStorage) UploadPart(uploadID string, partNumber int, data io.Reader) (string, error) {
	return s.multipart.UploadPart(uploadID, partNumber, data)
}

// ListParts lists parts of a multipart upload
func (s *UserAwareStorage) ListParts(uploadID string, maxParts, partNumberMarker int) ([]PartInfo, int, bool, error) {
	return s.multipart.ListParts(uploadID, maxParts, partNumberMarker)
}

// GetMultipartUpload retrieves multipart upload info
func (s *UserAwareStorage) GetMultipartUpload(uploadID string) (*MultipartUploadInfo, error) {
	return s.multipart.GetUpload(uploadID)
}

// AbortMultipartUpload cancels a multipart upload
func (s *UserAwareStorage) AbortMultipartUpload(uploadID string) error {
	return s.multipart.AbortMultipartUpload(uploadID)
}

// CompleteMultipartUpload finalizes a multipart upload
func (s *UserAwareStorage) CompleteMultipartUpload(uploadID string, parts []int) error {
	// CRITICAL FIX: Get upload info BEFORE completing (which deletes it from active uploads)
	upload, err := s.multipart.GetUpload(uploadID)
	if err != nil {
		return fmt.Errorf("upload not found: %s", uploadID)
	}

	// Get combined data (this also deletes the upload from active uploads map)
	data, etag, err := s.multipart.CompleteMultipartUpload(uploadID, parts)
	if err != nil {
		return err
	}

	// Store the final object
	reader := bytes.NewReader(data)
	if err := s.PutObjectForUser(upload.AccountID, upload.Bucket, upload.Key, reader); err != nil {
		return err
	}

	_ = etag // ETag could be stored in object metadata if needed

	return nil
}

// ListMultipartUploads lists in-progress multipart uploads
func (s *UserAwareStorage) ListMultipartUploads(accountID, bucket string, maxUploads int, keyMarker, uploadIDMarker string) ([]MultipartUploadInfo, string, string, bool, error) {
	return s.multipart.ListMultipartUploads(accountID, bucket, maxUploads, keyMarker, uploadIDMarker)
}