core/server/workspace/archive.go

package workspace

import (
	"archive/tar"
	"archive/zip"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	"switchboard-core/models"
)

// ── Archive Configuration ───────────────────

const (
	// MaxArchiveFiles is the max number of files to extract from an archive.
	MaxArchiveFiles = 10000

	// MaxArchiveFileSize is the max size of a single file within an archive.
	MaxArchiveFileSize int64 = 100 * 1024 * 1024 // 100 MB
)

// ── Extract ─────────────────────────────────

// ExtractArchive extracts a zip or tar.gz archive into the workspace.
// Returns the number of files extracted. Respects workspace quota.
// format must be "zip" or "tar.gz".
func (fs *FS) ExtractArchive(ctx context.Context, w *models.Workspace, archivePath, format string) (int, error) {
	var count int
	var err error
	switch format {
	case "zip":
		count, err = fs.extractZip(ctx, w, archivePath)
	case "tar.gz", "tgz":
		count, err = fs.extractTarGz(ctx, w, archivePath)
	default:
		return 0, fmt.Errorf("workspace: unsupported archive format: %s", format)
	}

	// Trigger batch indexing for all extracted files (v0.21.2)
	if err == nil && fs.indexer != nil && count > 0 {
		files, listErr := fs.store.ListFiles(ctx, w.ID, "", true)
		if listErr == nil {
			var teamID *string
			if w.OwnerType == models.WorkspaceOwnerTeam {
				teamID = &w.OwnerID
			}
			fs.indexer.IndexBatch(w, files, w.OwnerID, teamID)
		}
	}

	return count, err
}

func (fs *FS) extractZip(ctx context.Context, w *models.Workspace, archivePath string) (int, error) {
	r, err := zip.OpenReader(archivePath)
	if err != nil {
		return 0, fmt.Errorf("workspace: open zip: %w", err)
	}
	defer r.Close()

	if len(r.File) > MaxArchiveFiles {
		return 0, fmt.Errorf("workspace: archive contains %d files (max %d)", len(r.File), MaxArchiveFiles)
	}

	// Detect common root prefix (e.g. "project-name/") and strip it
	prefix := detectCommonPrefix(zipFileNames(r.File))

	root := fs.filesDir(w)
	count := 0
	var totalBytes int64

	for _, f := range r.File {
		if err := ctx.Err(); err != nil {
			return count, err
		}

		name := stripPrefix(f.Name, prefix)
		if name == "" || name == "." {
			continue
		}

		// Security: reject absolute paths and traversal
		if isUnsafePath(name) {
			log.Printf("workspace: skipping unsafe zip entry: %s", f.Name)
			continue
		}

		destPath := filepath.Join(root, filepath.FromSlash(name))

		if f.FileInfo().IsDir() {
			if err := os.MkdirAll(destPath, 0750); err != nil {
				return count, err
			}
			fs.store.UpsertFile(ctx, &models.WorkspaceFile{
				WorkspaceID: w.ID,
				Path:        name,
				IsDirectory: true,
			})
			continue
		}

		// Size check
		if f.UncompressedSize64 > uint64(MaxArchiveFileSize) {
			log.Printf("workspace: skipping oversized file: %s (%d bytes)", name, f.UncompressedSize64)
			continue
		}

		// Quota check
		quota := DefaultMaxBytes
		if w.MaxBytes != nil {
			quota = *w.MaxBytes
		}
		if totalBytes+int64(f.UncompressedSize64) > quota {
			return count, fmt.Errorf("workspace: quota exceeded during extraction (%d bytes)", quota)
		}

		// Extract file
		rc, err := f.Open()
		if err != nil {
			return count, fmt.Errorf("workspace: open zip entry %s: %w", name, err)
		}

		n, hash, err := extractToFile(destPath, rc, MaxArchiveFileSize)
		rc.Close()
		if err != nil {
			return count, fmt.Errorf("workspace: extract %s: %w", name, err)
		}

		totalBytes += n
		count++

		contentType := detectContentType(name, destPath)
		fs.store.UpsertFile(ctx, &models.WorkspaceFile{
			WorkspaceID: w.ID,
			Path:        name,
			IsDirectory: false,
			ContentType: contentType,
			SizeBytes:   n,
			SHA256:      hash,
		})
	}

	log.Printf("workspace: extracted %d files from zip (%d bytes)", count, totalBytes)
	return count, nil
}

func (fs *FS) extractTarGz(ctx context.Context, w *models.Workspace, archivePath string) (int, error) {
	f, err := os.Open(archivePath)
	if err != nil {
		return 0, fmt.Errorf("workspace: open tar.gz: %w", err)
	}
	defer f.Close()

	gz, err := gzip.NewReader(f)
	if err != nil {
		return 0, fmt.Errorf("workspace: gzip reader: %w", err)
	}
	defer gz.Close()

	tr := tar.NewReader(gz)
	root := fs.filesDir(w)
	count := 0
	var totalBytes int64

	// First pass: collect names to detect common prefix
	// Since tar is streaming, we can't do two passes easily.
	// We'll detect prefix from the first entry's directory.
	var prefix string
	prefixDetected := false

	for {
		if err := ctx.Err(); err != nil {
			return count, err
		}

		header, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return count, fmt.Errorf("workspace: tar read: %w", err)
		}

		if count >= MaxArchiveFiles {
			return count, fmt.Errorf("workspace: archive contains more than %d files", MaxArchiveFiles)
		}

		// Detect prefix from first entry
		if !prefixDetected {
			if idx := strings.IndexByte(header.Name, '/'); idx > 0 {
				prefix = header.Name[:idx+1]
			}
			prefixDetected = true
		}

		name := stripPrefix(header.Name, prefix)
		if name == "" || name == "." {
			continue
		}

		if isUnsafePath(name) {
			log.Printf("workspace: skipping unsafe tar entry: %s", header.Name)
			continue
		}

		destPath := filepath.Join(root, filepath.FromSlash(name))

		switch header.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(destPath, 0750); err != nil {
				return count, err
			}
			fs.store.UpsertFile(ctx, &models.WorkspaceFile{
				WorkspaceID: w.ID,
				Path:        name,
				IsDirectory: true,
			})

		case tar.TypeReg:
			if header.Size > MaxArchiveFileSize {
				log.Printf("workspace: skipping oversized file: %s (%d bytes)", name, header.Size)
				continue
			}

			quota := DefaultMaxBytes
			if w.MaxBytes != nil {
				quota = *w.MaxBytes
			}
			if totalBytes+header.Size > quota {
				return count, fmt.Errorf("workspace: quota exceeded during extraction (%d bytes)", quota)
			}

			n, hash, extractErr := extractToFile(destPath, tr, MaxArchiveFileSize)
			if extractErr != nil {
				return count, fmt.Errorf("workspace: extract %s: %w", name, extractErr)
			}

			totalBytes += n
			count++

			contentType := detectContentType(name, destPath)
			fs.store.UpsertFile(ctx, &models.WorkspaceFile{
				WorkspaceID: w.ID,
				Path:        name,
				IsDirectory: false,
				ContentType: contentType,
				SizeBytes:   n,
				SHA256:      hash,
			})

		default:
			// Skip symlinks, devices, etc.
			continue
		}
	}

	log.Printf("workspace: extracted %d files from tar.gz (%d bytes)", count, totalBytes)
	return count, nil
}

// ── Create Archive ──────────────────────────

// CreateArchive packages the workspace into a zip or tar.gz archive.
// Returns a path to the temporary archive file. Caller must remove it when done.
func (fs *FS) CreateArchive(ctx context.Context, w *models.Workspace, format string) (string, error) {
	switch format {
	case "zip":
		return fs.createZip(ctx, w)
	case "tar.gz", "tgz":
		return fs.createTarGz(ctx, w)
	default:
		return "", fmt.Errorf("workspace: unsupported archive format: %s", format)
	}
}

func (fs *FS) createZip(ctx context.Context, w *models.Workspace) (string, error) {
	root := fs.filesDir(w)

	tmp, err := os.CreateTemp("", "ws-*.zip")
	if err != nil {
		return "", err
	}
	tmpName := tmp.Name()

	zw := zip.NewWriter(tmp)

	err = filepath.Walk(root, func(abs string, info os.FileInfo, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}
		if err := ctx.Err(); err != nil {
			return err
		}

		rel, err := filepath.Rel(root, abs)
		if err != nil {
			return err
		}
		if rel == "." {
			return nil
		}
		rel = filepath.ToSlash(rel)

		if info.IsDir() {
			_, err := zw.Create(rel + "/")
			return err
		}

		header, err := zip.FileInfoHeader(info)
		if err != nil {
			return err
		}
		header.Name = rel
		header.Method = zip.Deflate

		writer, err := zw.CreateHeader(header)
		if err != nil {
			return err
		}

		f, err := os.Open(abs)
		if err != nil {
			return err
		}
		defer f.Close()

		_, err = io.Copy(writer, f)
		return err
	})

	if closeErr := zw.Close(); closeErr != nil && err == nil {
		err = closeErr
	}
	if closeErr := tmp.Close(); closeErr != nil && err == nil {
		err = closeErr
	}

	if err != nil {
		os.Remove(tmpName)
		return "", fmt.Errorf("workspace: create zip: %w", err)
	}

	return tmpName, nil
}

func (fs *FS) createTarGz(ctx context.Context, w *models.Workspace) (string, error) {
	root := fs.filesDir(w)

	tmp, err := os.CreateTemp("", "ws-*.tar.gz")
	if err != nil {
		return "", err
	}
	tmpName := tmp.Name()

	gw := gzip.NewWriter(tmp)
	tw := tar.NewWriter(gw)

	err = filepath.Walk(root, func(abs string, info os.FileInfo, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}
		if err := ctx.Err(); err != nil {
			return err
		}

		rel, err := filepath.Rel(root, abs)
		if err != nil {
			return err
		}
		if rel == "." {
			return nil
		}
		rel = filepath.ToSlash(rel)

		header, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return err
		}
		header.Name = rel

		if err := tw.WriteHeader(header); err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		f, err := os.Open(abs)
		if err != nil {
			return err
		}
		defer f.Close()

		_, err = io.Copy(tw, f)
		return err
	})

	if closeErr := tw.Close(); closeErr != nil && err == nil {
		err = closeErr
	}
	if closeErr := gw.Close(); closeErr != nil && err == nil {
		err = closeErr
	}
	if closeErr := tmp.Close(); closeErr != nil && err == nil {
		err = closeErr
	}

	if err != nil {
		os.Remove(tmpName)
		return "", fmt.Errorf("workspace: create tar.gz: %w", err)
	}

	return tmpName, nil
}

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

// extractToFile writes content from a reader to a file, creating parent dirs.
// Returns bytes written and SHA256 hash.
func extractToFile(destPath string, r io.Reader, maxSize int64) (int64, string, error) {
	if err := os.MkdirAll(filepath.Dir(destPath), 0750); err != nil {
		return 0, "", err
	}

	f, err := os.Create(destPath)
	if err != nil {
		return 0, "", err
	}
	defer f.Close()

	h := sha256.New()
	tee := io.TeeReader(io.LimitReader(r, maxSize+1), h)

	n, err := io.Copy(f, tee)
	if err != nil {
		return n, "", err
	}
	if n > maxSize {
		os.Remove(destPath)
		return 0, "", fmt.Errorf("file exceeds max size (%d bytes)", maxSize)
	}

	return n, hex.EncodeToString(h.Sum(nil)), nil
}

// isUnsafePath returns true if the path contains traversal or unsafe patterns.
func isUnsafePath(p string) bool {
	if strings.HasPrefix(p, "/") || strings.HasPrefix(p, "\\") {
		return true
	}
	if strings.Contains(p, "..") {
		return true
	}
	if strings.HasPrefix(p, ".") && !strings.HasPrefix(p, "./") {
		// Allow dotfiles like .gitignore, but not .. or hidden dirs as root
		// Actually, dotfiles are fine. Let them through.
		return false
	}
	return false
}

// detectCommonPrefix finds a shared directory prefix among file names.
// E.g. ["project/src/a.go", "project/src/b.go"] → "project/"
func detectCommonPrefix(names []string) string {
	if len(names) == 0 {
		return ""
	}

	// Check if all files share a common first directory component
	var prefix string
	for _, name := range names {
		idx := strings.IndexByte(name, '/')
		if idx < 0 {
			// File at root level — no common prefix
			return ""
		}
		dir := name[:idx+1]
		if prefix == "" {
			prefix = dir
		} else if dir != prefix {
			return ""
		}
	}
	return prefix
}

// stripPrefix removes the common archive prefix from a file name.
func stripPrefix(name, prefix string) string {
	if prefix != "" {
		name = strings.TrimPrefix(name, prefix)
	}
	name = strings.TrimSuffix(name, "/")
	name = filepath.ToSlash(name)
	return cleanPath(name)
}

// zipFileNames extracts file names from zip entries.
func zipFileNames(files []*zip.File) []string {
	names := make([]string, len(files))
	for i, f := range files {
		names[i] = f.Name
	}
	return names
}