602 lines
17 KiB
Go
602 lines
17 KiB
Go
package storage
|
|
|
|
import (
|
|
"archive/tar"
|
|
"archive/zip"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
)
|
|
|
|
// Storage handles file storage operations
|
|
type Storage struct {
|
|
basePath string
|
|
}
|
|
|
|
// NewStorage creates a new storage instance
|
|
func NewStorage(basePath string) (*Storage, error) {
|
|
s := &Storage{basePath: basePath}
|
|
if err := s.init(); err != nil {
|
|
return nil, err
|
|
}
|
|
return s, nil
|
|
}
|
|
|
|
// init creates necessary directories
|
|
func (s *Storage) init() error {
|
|
dirs := []string{
|
|
s.basePath,
|
|
s.uploadsPath(),
|
|
s.outputsPath(),
|
|
s.tempPath(),
|
|
}
|
|
|
|
for _, dir := range dirs {
|
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
|
return fmt.Errorf("failed to create directory %s: %w", dir, err)
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// tempPath returns the path for temporary files
|
|
func (s *Storage) tempPath() string {
|
|
return filepath.Join(s.basePath, "temp")
|
|
}
|
|
|
|
// BasePath returns the storage base path (for cleanup tasks)
|
|
func (s *Storage) BasePath() string {
|
|
return s.basePath
|
|
}
|
|
|
|
// TempDir creates a temporary directory under the storage base path
|
|
// Returns the path to the temporary directory
|
|
func (s *Storage) TempDir(pattern string) (string, error) {
|
|
// Ensure temp directory exists
|
|
if err := os.MkdirAll(s.tempPath(), 0755); err != nil {
|
|
return "", fmt.Errorf("failed to create temp directory: %w", err)
|
|
}
|
|
|
|
// Create temp directory under storage base path
|
|
return os.MkdirTemp(s.tempPath(), pattern)
|
|
}
|
|
|
|
// uploadsPath returns the path for uploads
|
|
func (s *Storage) uploadsPath() string {
|
|
return filepath.Join(s.basePath, "uploads")
|
|
}
|
|
|
|
// outputsPath returns the path for outputs
|
|
func (s *Storage) outputsPath() string {
|
|
return filepath.Join(s.basePath, "outputs")
|
|
}
|
|
|
|
// JobPath returns the path for a specific job's files
|
|
func (s *Storage) JobPath(jobID int64) string {
|
|
return filepath.Join(s.basePath, "jobs", fmt.Sprintf("%d", jobID))
|
|
}
|
|
|
|
// SaveUpload saves an uploaded file
|
|
func (s *Storage) SaveUpload(jobID int64, filename string, reader io.Reader) (string, error) {
|
|
jobPath := s.JobPath(jobID)
|
|
if err := os.MkdirAll(jobPath, 0755); err != nil {
|
|
return "", fmt.Errorf("failed to create job directory: %w", err)
|
|
}
|
|
|
|
filePath := filepath.Join(jobPath, filename)
|
|
file, err := os.Create(filePath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to create file: %w", err)
|
|
}
|
|
defer file.Close()
|
|
|
|
if _, err := io.Copy(file, reader); err != nil {
|
|
return "", fmt.Errorf("failed to write file: %w", err)
|
|
}
|
|
|
|
return filePath, nil
|
|
}
|
|
|
|
// SaveOutput saves an output file
|
|
func (s *Storage) SaveOutput(jobID int64, filename string, reader io.Reader) (string, error) {
|
|
outputPath := filepath.Join(s.outputsPath(), fmt.Sprintf("%d", jobID))
|
|
if err := os.MkdirAll(outputPath, 0755); err != nil {
|
|
return "", fmt.Errorf("failed to create output directory: %w", err)
|
|
}
|
|
|
|
filePath := filepath.Join(outputPath, filename)
|
|
file, err := os.Create(filePath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to create file: %w", err)
|
|
}
|
|
defer file.Close()
|
|
|
|
if _, err := io.Copy(file, reader); err != nil {
|
|
return "", fmt.Errorf("failed to write file: %w", err)
|
|
}
|
|
|
|
return filePath, nil
|
|
}
|
|
|
|
// GetFile returns a file reader for the given path
|
|
func (s *Storage) GetFile(filePath string) (*os.File, error) {
|
|
return os.Open(filePath)
|
|
}
|
|
|
|
// DeleteFile deletes a file
|
|
func (s *Storage) DeleteFile(filePath string) error {
|
|
return os.Remove(filePath)
|
|
}
|
|
|
|
// DeleteJobFiles deletes all files for a job
|
|
func (s *Storage) DeleteJobFiles(jobID int64) error {
|
|
jobPath := s.JobPath(jobID)
|
|
if err := os.RemoveAll(jobPath); err != nil && !os.IsNotExist(err) {
|
|
return fmt.Errorf("failed to delete job files: %w", err)
|
|
}
|
|
|
|
outputPath := filepath.Join(s.outputsPath(), fmt.Sprintf("%d", jobID))
|
|
if err := os.RemoveAll(outputPath); err != nil && !os.IsNotExist(err) {
|
|
return fmt.Errorf("failed to delete output files: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// FileExists checks if a file exists
|
|
func (s *Storage) FileExists(filePath string) bool {
|
|
_, err := os.Stat(filePath)
|
|
return err == nil
|
|
}
|
|
|
|
// GetFileSize returns the size of a file
|
|
func (s *Storage) GetFileSize(filePath string) (int64, error) {
|
|
info, err := os.Stat(filePath)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return info.Size(), nil
|
|
}
|
|
|
|
// ExtractZip extracts a ZIP file to the destination directory
|
|
// Returns a list of all extracted file paths
|
|
func (s *Storage) ExtractZip(zipPath, destDir string) ([]string, error) {
|
|
log.Printf("Extracting ZIP archive: %s -> %s", zipPath, destDir)
|
|
|
|
// Ensure destination directory exists
|
|
if err := os.MkdirAll(destDir, 0755); err != nil {
|
|
return nil, fmt.Errorf("failed to create destination directory: %w", err)
|
|
}
|
|
|
|
r, err := zip.OpenReader(zipPath)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to open ZIP file: %w", err)
|
|
}
|
|
defer r.Close()
|
|
|
|
var extractedFiles []string
|
|
fileCount := 0
|
|
dirCount := 0
|
|
|
|
log.Printf("ZIP contains %d entries", len(r.File))
|
|
|
|
for _, f := range r.File {
|
|
// Sanitize file path to prevent directory traversal
|
|
destPath := filepath.Join(destDir, f.Name)
|
|
|
|
cleanDestPath := filepath.Clean(destPath)
|
|
cleanDestDir := filepath.Clean(destDir)
|
|
if !strings.HasPrefix(cleanDestPath, cleanDestDir+string(os.PathSeparator)) && cleanDestPath != cleanDestDir {
|
|
log.Printf("ERROR: Invalid file path in ZIP - target: %s, destDir: %s", cleanDestPath, cleanDestDir)
|
|
return nil, fmt.Errorf("invalid file path in ZIP: %s (target: %s, destDir: %s)", f.Name, cleanDestPath, cleanDestDir)
|
|
}
|
|
|
|
// Create directory structure
|
|
if f.FileInfo().IsDir() {
|
|
if err := os.MkdirAll(destPath, 0755); err != nil {
|
|
return nil, fmt.Errorf("failed to create directory: %w", err)
|
|
}
|
|
dirCount++
|
|
continue
|
|
}
|
|
|
|
// Create parent directories
|
|
if err := os.MkdirAll(filepath.Dir(destPath), 0755); err != nil {
|
|
return nil, fmt.Errorf("failed to create parent directory: %w", err)
|
|
}
|
|
|
|
// Extract file
|
|
rc, err := f.Open()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to open file in ZIP: %w", err)
|
|
}
|
|
|
|
outFile, err := os.Create(destPath)
|
|
if err != nil {
|
|
rc.Close()
|
|
return nil, fmt.Errorf("failed to create file: %w", err)
|
|
}
|
|
|
|
_, err = io.Copy(outFile, rc)
|
|
outFile.Close()
|
|
rc.Close()
|
|
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to extract file: %w", err)
|
|
}
|
|
|
|
extractedFiles = append(extractedFiles, destPath)
|
|
fileCount++
|
|
}
|
|
|
|
log.Printf("ZIP extraction complete: %d files, %d directories extracted to %s", fileCount, dirCount, destDir)
|
|
return extractedFiles, nil
|
|
}
|
|
|
|
// findCommonPrefix finds the common leading directory prefix if all paths share the same first-level directory
|
|
// Returns the prefix to strip (with trailing slash) or empty string if no common prefix
|
|
func findCommonPrefix(relPaths []string) string {
|
|
if len(relPaths) == 0 {
|
|
return ""
|
|
}
|
|
|
|
// Get the first path component of each path
|
|
firstComponents := make([]string, 0, len(relPaths))
|
|
for _, path := range relPaths {
|
|
parts := strings.Split(filepath.ToSlash(path), "/")
|
|
if len(parts) > 0 && parts[0] != "" {
|
|
firstComponents = append(firstComponents, parts[0])
|
|
} else {
|
|
// If any path is at root level, no common prefix
|
|
return ""
|
|
}
|
|
}
|
|
|
|
// Check if all first components are the same
|
|
if len(firstComponents) == 0 {
|
|
return ""
|
|
}
|
|
|
|
commonFirst := firstComponents[0]
|
|
for _, comp := range firstComponents {
|
|
if comp != commonFirst {
|
|
// Not all paths share the same first directory
|
|
return ""
|
|
}
|
|
}
|
|
|
|
// All paths share the same first directory - return it with trailing slash
|
|
return commonFirst + "/"
|
|
}
|
|
|
|
// isBlenderSaveFile checks if a filename is a Blender save file (.blend1, .blend2, etc.)
|
|
// Returns true for files like "file.blend1", "file.blend2", but false for "file.blend"
|
|
func isBlenderSaveFile(filename string) bool {
|
|
lower := strings.ToLower(filename)
|
|
// Check if it ends with .blend followed by one or more digits
|
|
// Pattern: *.blend[digits]
|
|
if !strings.HasSuffix(lower, ".blend") {
|
|
// Doesn't end with .blend, check if it ends with .blend + digits
|
|
idx := strings.LastIndex(lower, ".blend")
|
|
if idx == -1 {
|
|
return false
|
|
}
|
|
// Check if there are digits after .blend
|
|
suffix := lower[idx+len(".blend"):]
|
|
if len(suffix) == 0 {
|
|
return false
|
|
}
|
|
// All remaining characters must be digits
|
|
for _, r := range suffix {
|
|
if r < '0' || r > '9' {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
// Ends with .blend exactly - this is a regular blend file, not a save file
|
|
return false
|
|
}
|
|
|
|
// CreateJobContext creates a tar archive containing all job input files
|
|
// Filters out Blender save files (.blend1, .blend2, etc.)
|
|
// Uses temporary directories and streaming to handle large files efficiently
|
|
func (s *Storage) CreateJobContext(jobID int64) (string, error) {
|
|
jobPath := s.JobPath(jobID)
|
|
contextPath := filepath.Join(jobPath, "context.tar")
|
|
|
|
// Create temporary directory for staging
|
|
tmpDir, err := os.MkdirTemp("", "jiggablend-context-*")
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to create temporary directory: %w", err)
|
|
}
|
|
defer os.RemoveAll(tmpDir)
|
|
|
|
// Collect all files from job directory, excluding the context file itself and Blender save files
|
|
var filesToInclude []string
|
|
err = filepath.Walk(jobPath, func(path string, info os.FileInfo, err error) error {
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Skip directories
|
|
if info.IsDir() {
|
|
return nil
|
|
}
|
|
|
|
// Skip the context file itself if it exists
|
|
if path == contextPath {
|
|
return nil
|
|
}
|
|
|
|
// Skip Blender save files
|
|
if isBlenderSaveFile(info.Name()) {
|
|
return nil
|
|
}
|
|
|
|
// Get relative path from job directory
|
|
relPath, err := filepath.Rel(jobPath, path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Sanitize path - ensure it doesn't escape the job directory
|
|
cleanRelPath := filepath.Clean(relPath)
|
|
if strings.HasPrefix(cleanRelPath, "..") {
|
|
return fmt.Errorf("invalid file path: %s", relPath)
|
|
}
|
|
|
|
filesToInclude = append(filesToInclude, path)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to walk job directory: %w", err)
|
|
}
|
|
|
|
if len(filesToInclude) == 0 {
|
|
return "", fmt.Errorf("no files found to include in context")
|
|
}
|
|
|
|
// Create the tar file using streaming
|
|
contextFile, err := os.Create(contextPath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to create context file: %w", err)
|
|
}
|
|
defer contextFile.Close()
|
|
|
|
tarWriter := tar.NewWriter(contextFile)
|
|
defer tarWriter.Close()
|
|
|
|
// Add each file to the tar archive
|
|
for _, filePath := range filesToInclude {
|
|
file, err := os.Open(filePath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to open file %s: %w", filePath, err)
|
|
}
|
|
|
|
info, err := file.Stat()
|
|
if err != nil {
|
|
file.Close()
|
|
return "", fmt.Errorf("failed to stat file %s: %w", filePath, err)
|
|
}
|
|
|
|
// Get relative path for tar header
|
|
relPath, err := filepath.Rel(jobPath, filePath)
|
|
if err != nil {
|
|
file.Close()
|
|
return "", fmt.Errorf("failed to get relative path for %s: %w", filePath, err)
|
|
}
|
|
|
|
// Normalize path separators for tar (use forward slashes)
|
|
tarPath := filepath.ToSlash(relPath)
|
|
|
|
// Create tar header
|
|
header, err := tar.FileInfoHeader(info, "")
|
|
if err != nil {
|
|
file.Close()
|
|
return "", fmt.Errorf("failed to create tar header for %s: %w", filePath, err)
|
|
}
|
|
header.Name = tarPath
|
|
|
|
// Write header
|
|
if err := tarWriter.WriteHeader(header); err != nil {
|
|
file.Close()
|
|
return "", fmt.Errorf("failed to write tar header for %s: %w", filePath, err)
|
|
}
|
|
|
|
// Copy file contents using streaming
|
|
if _, err := io.Copy(tarWriter, file); err != nil {
|
|
file.Close()
|
|
return "", fmt.Errorf("failed to write file %s to tar: %w", filePath, err)
|
|
}
|
|
|
|
file.Close()
|
|
}
|
|
|
|
// Ensure all data is flushed
|
|
if err := tarWriter.Close(); err != nil {
|
|
return "", fmt.Errorf("failed to close tar writer: %w", err)
|
|
}
|
|
if err := contextFile.Close(); err != nil {
|
|
return "", fmt.Errorf("failed to close context file: %w", err)
|
|
}
|
|
|
|
return contextPath, nil
|
|
}
|
|
|
|
// CreateJobContextFromDir creates a context archive (tar) from files in a source directory
|
|
// This is used during upload to immediately create the context archive as the primary artifact
|
|
// excludeFiles is a set of relative paths (from sourceDir) to exclude from the context
|
|
func (s *Storage) CreateJobContextFromDir(sourceDir string, jobID int64, excludeFiles ...string) (string, error) {
|
|
jobPath := s.JobPath(jobID)
|
|
contextPath := filepath.Join(jobPath, "context.tar")
|
|
|
|
// Ensure job directory exists
|
|
if err := os.MkdirAll(jobPath, 0755); err != nil {
|
|
return "", fmt.Errorf("failed to create job directory: %w", err)
|
|
}
|
|
|
|
// Build set of files to exclude (normalize paths)
|
|
excludeSet := make(map[string]bool)
|
|
for _, excludeFile := range excludeFiles {
|
|
// Normalize the exclude path
|
|
excludePath := filepath.Clean(excludeFile)
|
|
excludeSet[excludePath] = true
|
|
// Also add with forward slash for cross-platform compatibility
|
|
excludeSet[filepath.ToSlash(excludePath)] = true
|
|
}
|
|
|
|
// Collect all files from source directory, excluding Blender save files and excluded files
|
|
var filesToInclude []string
|
|
err := filepath.Walk(sourceDir, func(path string, info os.FileInfo, err error) error {
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Skip directories
|
|
if info.IsDir() {
|
|
return nil
|
|
}
|
|
|
|
// Skip Blender save files
|
|
if isBlenderSaveFile(info.Name()) {
|
|
return nil
|
|
}
|
|
|
|
// Get relative path from source directory
|
|
relPath, err := filepath.Rel(sourceDir, path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Sanitize path - ensure it doesn't escape the source directory
|
|
cleanRelPath := filepath.Clean(relPath)
|
|
if strings.HasPrefix(cleanRelPath, "..") {
|
|
return fmt.Errorf("invalid file path: %s", relPath)
|
|
}
|
|
|
|
// Check if this file should be excluded
|
|
if excludeSet[cleanRelPath] || excludeSet[filepath.ToSlash(cleanRelPath)] {
|
|
return nil
|
|
}
|
|
|
|
filesToInclude = append(filesToInclude, path)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to walk source directory: %w", err)
|
|
}
|
|
|
|
if len(filesToInclude) == 0 {
|
|
return "", fmt.Errorf("no files found to include in context archive")
|
|
}
|
|
|
|
// Collect relative paths to find common prefix
|
|
relPaths := make([]string, 0, len(filesToInclude))
|
|
for _, filePath := range filesToInclude {
|
|
relPath, err := filepath.Rel(sourceDir, filePath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to get relative path for %s: %w", filePath, err)
|
|
}
|
|
relPaths = append(relPaths, relPath)
|
|
}
|
|
|
|
// Find and strip common leading directory if all files share one
|
|
commonPrefix := findCommonPrefix(relPaths)
|
|
|
|
// Validate that there's exactly one .blend file at the root level after prefix stripping
|
|
blendFilesAtRoot := 0
|
|
for _, relPath := range relPaths {
|
|
tarPath := filepath.ToSlash(relPath)
|
|
// Strip common prefix if present
|
|
if commonPrefix != "" && strings.HasPrefix(tarPath, commonPrefix) {
|
|
tarPath = strings.TrimPrefix(tarPath, commonPrefix)
|
|
}
|
|
|
|
// Check if it's a .blend file at root (no path separators after prefix stripping)
|
|
if strings.HasSuffix(strings.ToLower(tarPath), ".blend") {
|
|
// Check if it's at root level (no directory separators)
|
|
if !strings.Contains(tarPath, "/") {
|
|
blendFilesAtRoot++
|
|
}
|
|
}
|
|
}
|
|
|
|
if blendFilesAtRoot == 0 {
|
|
return "", fmt.Errorf("no .blend file found at root level in context archive - .blend files must be at the root level of the uploaded archive, not in subdirectories")
|
|
}
|
|
if blendFilesAtRoot > 1 {
|
|
return "", fmt.Errorf("multiple .blend files found at root level in context archive (found %d, expected 1)", blendFilesAtRoot)
|
|
}
|
|
|
|
// Create the tar file using streaming
|
|
contextFile, err := os.Create(contextPath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to create context file: %w", err)
|
|
}
|
|
defer contextFile.Close()
|
|
|
|
tarWriter := tar.NewWriter(contextFile)
|
|
defer tarWriter.Close()
|
|
|
|
// Add each file to the tar archive
|
|
for i, filePath := range filesToInclude {
|
|
file, err := os.Open(filePath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to open file %s: %w", filePath, err)
|
|
}
|
|
|
|
info, err := file.Stat()
|
|
if err != nil {
|
|
file.Close()
|
|
return "", fmt.Errorf("failed to stat file %s: %w", filePath, err)
|
|
}
|
|
|
|
// Get relative path and strip common prefix if present
|
|
relPath := relPaths[i]
|
|
tarPath := filepath.ToSlash(relPath)
|
|
|
|
// Strip common prefix if found
|
|
if commonPrefix != "" && strings.HasPrefix(tarPath, commonPrefix) {
|
|
tarPath = strings.TrimPrefix(tarPath, commonPrefix)
|
|
}
|
|
|
|
// Create tar header
|
|
header, err := tar.FileInfoHeader(info, "")
|
|
if err != nil {
|
|
file.Close()
|
|
return "", fmt.Errorf("failed to create tar header for %s: %w", filePath, err)
|
|
}
|
|
header.Name = tarPath
|
|
|
|
// Write header
|
|
if err := tarWriter.WriteHeader(header); err != nil {
|
|
file.Close()
|
|
return "", fmt.Errorf("failed to write tar header for %s: %w", filePath, err)
|
|
}
|
|
|
|
// Copy file contents using streaming
|
|
if _, err := io.Copy(tarWriter, file); err != nil {
|
|
file.Close()
|
|
return "", fmt.Errorf("failed to write file %s to tar: %w", filePath, err)
|
|
}
|
|
|
|
file.Close()
|
|
}
|
|
|
|
// Ensure all data is flushed
|
|
if err := tarWriter.Close(); err != nil {
|
|
return "", fmt.Errorf("failed to close tar writer: %w", err)
|
|
}
|
|
if err := contextFile.Close(); err != nil {
|
|
return "", fmt.Errorf("failed to close context file: %w", err)
|
|
}
|
|
|
|
return contextPath, nil
|
|
}
|
|
|