core/server/storage/pvc.go

package storage

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

// ── PVC Store ──────────────────────────────

// PVCStore implements ObjectStore using the local filesystem.
// Used with Kubernetes PersistentVolumeClaims (CephFS, NFS, local PV)
// or plain Docker volume mounts.
type PVCStore struct {
	basePath string

	// Content-type cache: key → content type.
	// Populated on Put, used by Get. PVC doesn't store metadata
	// alongside files, so we keep MIME types in memory. On restart,
	// Get falls back to "application/octet-stream" — the DB has the
	// authoritative content_type for serving to clients.
	mu       sync.RWMutex
	mimeMap  map[string]string
}

// NewPVC creates a PVCStore rooted at basePath.
// Returns an error if the path does not exist or is not writable.
func NewPVC(basePath string) (*PVCStore, error) {
	// Ensure base path exists
	if err := os.MkdirAll(basePath, 0o755); err != nil {
		return nil, fmt.Errorf("storage/pvc: cannot create base path %q: %w", basePath, err)
	}

	// Create well-known subdirectories
	for _, sub := range []string{"files", "processing"} {
		if err := os.MkdirAll(filepath.Join(basePath, sub), 0o755); err != nil {
			return nil, fmt.Errorf("storage/pvc: cannot create %s dir: %w", sub, err)
		}
	}

	s := &PVCStore{
		basePath: basePath,
		mimeMap:  make(map[string]string),
	}

	// Validate writability
	if err := s.Healthy(context.Background()); err != nil {
		return nil, err
	}

	return s, nil
}

func (s *PVCStore) Backend() string { return "pvc" }

// Put writes data to the filesystem at the resolved key path.
func (s *PVCStore) Put(ctx context.Context, key string, r io.Reader, size int64, contentType string) error {
	if err := validateKey(key); err != nil {
		return err
	}

	absPath := s.resolve(key)

	// Ensure parent directory exists
	if err := os.MkdirAll(filepath.Dir(absPath), 0o755); err != nil {
		return fmt.Errorf("storage/pvc: mkdir for %q: %w", key, err)
	}

	// Write to temp file first, then rename (atomic on same filesystem)
	tmp, err := os.CreateTemp(filepath.Dir(absPath), ".upload-*")
	if err != nil {
		return fmt.Errorf("storage/pvc: create temp for %q: %w", key, err)
	}
	tmpName := tmp.Name()
	defer func() {
		tmp.Close()
		// Clean up temp file on failure
		if err != nil {
			os.Remove(tmpName)
		}
	}()

	written, err := io.Copy(tmp, r)
	if err != nil {
		return fmt.Errorf("storage/pvc: write %q: %w", key, err)
	}
	if err = tmp.Close(); err != nil {
		return fmt.Errorf("storage/pvc: close %q: %w", key, err)
	}

	// Verify size if provided (0 means unknown)
	if size > 0 && written != size {
		os.Remove(tmpName)
		return fmt.Errorf("storage/pvc: size mismatch for %q: expected %d, got %d", key, size, written)
	}

	// Atomic rename
	if err = os.Rename(tmpName, absPath); err != nil {
		return fmt.Errorf("storage/pvc: rename %q: %w", key, err)
	}

	// Cache content type
	if contentType != "" {
		s.mu.Lock()
		s.mimeMap[key] = contentType
		s.mu.Unlock()
	}

	return nil
}

// Get returns a reader for the file at key.
func (s *PVCStore) Get(ctx context.Context, key string) (io.ReadCloser, int64, string, error) {
	if err := validateKey(key); err != nil {
		return nil, 0, "", err
	}

	absPath := s.resolve(key)

	info, err := os.Stat(absPath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, 0, "", ErrNotFound
		}
		return nil, 0, "", fmt.Errorf("storage/pvc: stat %q: %w", key, err)
	}

	f, err := os.Open(absPath)
	if err != nil {
		return nil, 0, "", fmt.Errorf("storage/pvc: open %q: %w", key, err)
	}

	// Look up cached content type
	s.mu.RLock()
	ct := s.mimeMap[key]
	s.mu.RUnlock()
	if ct == "" {
		ct = "application/octet-stream"
	}

	return f, info.Size(), ct, nil
}

// Delete removes the file at key.
func (s *PVCStore) Delete(ctx context.Context, key string) error {
	if err := validateKey(key); err != nil {
		return err
	}

	absPath := s.resolve(key)

	err := os.Remove(absPath)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("storage/pvc: delete %q: %w", key, err)
	}

	// Clean up cache
	s.mu.Lock()
	delete(s.mimeMap, key)
	s.mu.Unlock()

	return nil
}

// DeletePrefix removes all files and directories under the given prefix.
func (s *PVCStore) DeletePrefix(ctx context.Context, prefix string) error {
	if err := validateKey(prefix); err != nil {
		return err
	}

	absPath := s.resolve(prefix)

	// Check if the path exists
	if _, err := os.Stat(absPath); os.IsNotExist(err) {
		return nil // Nothing to delete
	}

	if err := os.RemoveAll(absPath); err != nil {
		return fmt.Errorf("storage/pvc: delete prefix %q: %w", prefix, err)
	}

	// Clean up cache entries with this prefix
	s.mu.Lock()
	for k := range s.mimeMap {
		if strings.HasPrefix(k, prefix) {
			delete(s.mimeMap, k)
		}
	}
	s.mu.Unlock()

	return nil
}

// Exists checks if a file exists at key.
func (s *PVCStore) Exists(ctx context.Context, key string) (bool, error) {
	if err := validateKey(key); err != nil {
		return false, err
	}

	_, err := os.Stat(s.resolve(key))
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, fmt.Errorf("storage/pvc: exists %q: %w", key, err)
}

// List returns objects under the given prefix, up to limit entries.
func (s *PVCStore) List(ctx context.Context, prefix string, limit int) ([]ObjectEntry, error) {
	if prefix != "" {
		if err := validateKey(prefix); err != nil {
			return nil, err
		}
	}
	if limit <= 0 {
		limit = 100
	}

	absPrefix := s.resolve(prefix)

	// If the prefix path doesn't exist, return empty.
	info, err := os.Stat(absPrefix)
	if os.IsNotExist(err) {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("storage/pvc: list %q: %w", prefix, err)
	}

	var entries []ObjectEntry

	if !info.IsDir() {
		// Prefix points to a single file — return it.
		s.mu.RLock()
		ct := s.mimeMap[prefix]
		s.mu.RUnlock()
		if ct == "" {
			ct = "application/octet-stream"
		}
		return []ObjectEntry{{Key: prefix, Size: info.Size(), ContentType: ct}}, nil
	}

	// Walk the directory tree under the prefix.
	walkRoot := absPrefix
	err = filepath.Walk(walkRoot, func(path string, fi os.FileInfo, walkErr error) error {
		if walkErr != nil {
			return nil // skip inaccessible entries
		}
		if fi.IsDir() {
			return nil
		}
		if len(entries) >= limit {
			return filepath.SkipAll
		}

		// Recover the storage key by stripping basePath + "/".
		rel, relErr := filepath.Rel(s.basePath, path)
		if relErr != nil {
			return nil
		}
		key := filepath.ToSlash(rel)

		s.mu.RLock()
		ct := s.mimeMap[key]
		s.mu.RUnlock()
		if ct == "" {
			ct = "application/octet-stream"
		}

		entries = append(entries, ObjectEntry{Key: key, Size: fi.Size(), ContentType: ct})
		return nil
	})
	if err != nil {
		return entries, fmt.Errorf("storage/pvc: list walk %q: %w", prefix, err)
	}

	return entries, nil
}

// Healthy verifies the storage path is writable.
func (s *PVCStore) Healthy(ctx context.Context) error {
	probe := filepath.Join(s.basePath, ".health-probe")
	if err := os.WriteFile(probe, []byte("ok"), 0o644); err != nil {
		return fmt.Errorf("storage/pvc: not writable: %w", err)
	}
	os.Remove(probe)
	return nil
}

// Stats returns file count and total size under the base path.
func (s *PVCStore) Stats(ctx context.Context) (*StorageStats, error) {
	stats := &StorageStats{
		Backend:    "pvc",
		Path:       s.basePath,
		Configured: true,
	}

	// Check health
	if err := s.Healthy(ctx); err != nil {
		stats.Healthy = false
		return stats, nil
	}
	stats.Healthy = true

	// Walk files directory for counts
	fileDir := filepath.Join(s.basePath, "files")
	err := filepath.Walk(fileDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil // skip errors (e.g. permission denied)
		}
		if !info.IsDir() {
			stats.TotalFiles++
			stats.TotalBytes += info.Size()
		}
		return nil
	})
	if err != nil && !os.IsNotExist(err) {
		return stats, fmt.Errorf("storage/pvc: stats walk: %w", err)
	}

	return stats, nil
}

// ── Helpers ────────────────────────────────

// resolve converts a storage key to an absolute filesystem path.
func (s *PVCStore) resolve(key string) string {
	return filepath.Join(s.basePath, filepath.FromSlash(key))
}

// validateKey ensures the key is safe (no path traversal, no absolute paths).
func validateKey(key string) error {
	if key == "" {
		return fmt.Errorf("storage: empty key")
	}
	if strings.HasPrefix(key, "/") || strings.HasPrefix(key, "\\") {
		return fmt.Errorf("storage: absolute key not allowed: %q", key)
	}
	cleaned := filepath.Clean(key)
	if strings.Contains(cleaned, "..") {
		return fmt.Errorf("storage: path traversal not allowed: %q", key)
	}
	return nil
}