ygggo/internal/storage/filestore/filestore.go

/*
 *  Copyright (c) 2024 JB-SelfCompany / Tyr Contributors
 *
 *  https://github.com/JB-SelfCompany/yggmail
 *
 *  This Source Code Form is subject to the terms of the Mozilla Public
 *  License, v. 2.0. If a copy of the MPL was not distributed with this
 *  file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */
package filestore

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// FileStore manages file-based storage for large email messages
type FileStore struct {
	basePath string
	mu       sync.RWMutex
}

// BasePath returns the base path of the file store
func (fs *FileStore) BasePath() string {
	fs.mu.RLock()
	defer fs.mu.RUnlock()
	return fs.basePath
}

// sanitizeMailboxName removes potentially dangerous characters from mailbox names
func sanitizeMailboxName(mailbox string) string {
	// Replace path separators and other dangerous characters
	sanitized := strings.ReplaceAll(mailbox, "/", "_")
	sanitized = strings.ReplaceAll(sanitized, "\\", "_")
	sanitized = strings.ReplaceAll(sanitized, "..", "_")
	sanitized = strings.ReplaceAll(sanitized, ":", "_")
	sanitized = strings.TrimSpace(sanitized)

	if sanitized == "" {
		sanitized = "default"
	}

	return sanitized
}

// NewFileStore creates a new FileStore and initializes directory structure
// Creates subdirectories for different mailboxes (INBOX, Outbox, etc.)
func NewFileStore(basePath string) (*FileStore, error) {
	if basePath == "" {
		return nil, fmt.Errorf("basePath cannot be empty")
	}
	if err := os.MkdirAll(basePath, 0755); err != nil {
		return nil, fmt.Errorf("failed to create base directory: %w", err)
	}
	return &FileStore{basePath: basePath}, nil
}

func genName(aPk string) (ts int64, hs string) {
	ts = time.Now().UnixNano()
	h := sha256.New()
	h.Write(fmt.Appendf([]byte{}, "%d%s", int(ts), aPk))
	hs = hex.EncodeToString(h.Sum(nil))[:16]
	return ts, hs
}

// StoreMail stores an email message from a reader to a file
// Returns the file path where the message was stored and the actual bytes written
// Uses atomic write (temp file + rename) to prevent corruption
func (fs *FileStore) StoreMail(aPk string, mailbox string, reader io.Reader) (string, int64, error) {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	sanitized := sanitizeMailboxName(mailbox)
	mailboxPath := filepath.Join(fs.basePath, aPk, sanitized)

	if err := os.MkdirAll(mailboxPath, 0755); err != nil {
		return "", 0, fmt.Errorf("failed to create user mailbox directory: %w", err)
	}

	ts, hs := genName(aPk)

	// Create temp file in same directory (for atomic rename)
	tempFile, err := os.CreateTemp(mailboxPath, fmt.Sprintf(".tmp_%d_%s_*", ts, hs))
	if err != nil {
		return "", 0, fmt.Errorf("failed to create temp file: %w", err)
	}
	tempPath := tempFile.Name()

	//Clean up temp file on error
	defer func() {
		if tempFile != nil {
			tempFile.Close()
			os.Remove(tempPath)
		}
	}()

	// Stream data to temp file with buffering
	written, err := io.Copy(tempFile, reader)
	if err != nil {
		return "", 0, fmt.Errorf("failed to write data: %w", err)
	}

	// Log size for debugging
	// fmt.Printf("[FileStore] StoreMail: mailID=%d, written=%d bytes (%.2f MB)\n",
	// 	0, written, float64(written)/(1024*1024))
	// if written == 0 {
	// 	log.Printf("WARNING: empty file written for id=%d", 0)
	// }

	// Sync to disk
	if err := tempFile.Sync(); err != nil {
		return "", 0, fmt.Errorf("failed to sync file: %w", err)
	}

	// Close before rename
	if err := tempFile.Close(); err != nil {
		return "", 0, fmt.Errorf("failed to close temp file: %w", err)
	}
	tempFile = nil // Prevent defer cleanup

	// Final file path
	finalPath := filepath.Join(mailboxPath, fmt.Sprintf("%d_%s.eml", ts, hs))

	// Atomic rename
	if err := os.Rename(tempPath, finalPath); err != nil {
		return "", 0, fmt.Errorf("failed to rename temp file: %w", err)
	}

	// Return relative path from basePath
	relPath, err := filepath.Rel(fs.basePath, finalPath)
	if err != nil {
		return "", 0, fmt.Errorf("failed to get relative path: %w", err)
	}

	return relPath, written, nil
}

// ReadMail opens and returns a reader for the specified mail file
// Caller is responsible for closing the returned ReadCloser
func (fs *FileStore) ReadMail(location string) (io.ReadCloser, error) {
	fs.mu.RLock()
	defer fs.mu.RUnlock()

	fullPath := filepath.Join(fs.basePath, location)

	if !strings.HasPrefix(filepath.Clean(fullPath), filepath.Clean(fs.basePath)) {
		return nil, fmt.Errorf("invalid path traversal attempt")
	}

	return os.Open(fullPath)
}

// DeleteMail deletes the specified mail file
func (fs *FileStore) DeleteMail(location string) error {
	// log.Println("Delete: ", location)
	fs.mu.Lock()
	defer fs.mu.Unlock()

	fullPath := filepath.Join(fs.basePath, location)
	return os.Remove(fullPath)
}

// MoveMail moves a mail file to a different mailbox with a new ID
// Returns the new relative path
func (fs *FileStore) MoveMail(oldLocation string, newMailbox string) (string, error) {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	oldFullPath := filepath.Join(fs.basePath, oldLocation)

	parts := strings.Split(filepath.ToSlash(oldLocation), "/")
	if len(parts) < 3 {
		return "", fmt.Errorf("invalid location format")
	}
	aPk := parts[0]
	fileName := parts[len(parts)-1]

	sanitizedNew := sanitizeMailboxName(newMailbox)
	newRelPath := filepath.Join(aPk, sanitizedNew, fileName)
	newFullPath := filepath.Join(fs.basePath, newRelPath)

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

	if err := os.Rename(oldFullPath, newFullPath); err != nil {
		return "", err
	}
	// log.Printf("Moved file from %s to %s", oldFullPath, newFullPath)

	return newRelPath, nil
}

// GetTotalSize returns the total size of all stored mail files in bytes
func (fs *FileStore) GetTotalSize() (int64, error) {
	fs.mu.RLock()
	defer fs.mu.RUnlock()

	var totalSize int64

	err := filepath.Walk(fs.basePath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && strings.HasSuffix(info.Name(), ".eml") {
			totalSize += info.Size()
		}
		return nil
	})

	if err != nil {
		return 0, fmt.Errorf("failed to calculate total size: %w", err)
	}

	return totalSize, nil
}

// CleanupOrphanedFiles removes mail files that are not referenced in the database
// This requires access to the storage layer to query existing mail records
// Returns the number of files deleted and total bytes freed
func (fs *FileStore) CleanupOrphanedFiles(validPaths map[string]bool) (int, int64, error) {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	deletedCount := 0
	var deletedSize int64

	err := filepath.Walk(fs.basePath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Skip directories and non-.eml files
		if info.IsDir() || !strings.HasSuffix(info.Name(), ".eml") {
			return nil
		}

		// Get relative path
		relPath, err := filepath.Rel(fs.basePath, path)
		if err != nil {
			return err
		}

		// Check if this file is in the valid paths map
		if !validPaths[relPath] {
			// Orphaned file - delete it
			size := info.Size()
			if err := os.Remove(path); err != nil {
				return fmt.Errorf("failed to delete orphaned file %s: %w", relPath, err)
			}
			deletedCount++
			deletedSize += size
		}

		return nil
	})

	if err != nil {
		return deletedCount, deletedSize, fmt.Errorf("cleanup failed: %w", err)
	}

	return deletedCount, deletedSize, nil
}

// GetMailboxSize returns the total size of files in a specific mailbox
func (fs *FileStore) GetMailboxSize(aPk string, mailbox string) (int64, error) {
	fs.mu.RLock()
	defer fs.mu.RUnlock()

	sanitized := sanitizeMailboxName(mailbox)
	mailboxPath := filepath.Join(fs.basePath, aPk, sanitized)

	var totalSize int64
	err := filepath.Walk(mailboxPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && strings.HasSuffix(info.Name(), ".eml") {
			totalSize += info.Size()
		}
		return nil
	})
	return totalSize, err
}