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) }