sub2api-cn-relay-manager/internal/pack/source_loader.go

package pack

import (
	"archive/zip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

const (
	maxArchiveEntries   = 256
	maxArchiveFileSize  = 5 << 20
	maxArchiveTotalSize = 20 << 20
)

func LoadPath(path string) (LoadedPack, error) {
	trimmed := strings.TrimSpace(path)
	if trimmed == "" {
		return LoadedPack{}, fmt.Errorf("pack path is required")
	}

	info, err := os.Stat(trimmed)
	if err != nil {
		return LoadedPack{}, fmt.Errorf("stat pack path: %w", err)
	}
	if info.IsDir() {
		return LoadDir(trimmed)
	}
	if strings.EqualFold(filepath.Ext(info.Name()), ".zip") {
		return LoadArchive(trimmed)
	}
	return LoadedPack{}, fmt.Errorf("pack path %q must be a directory or .zip archive", trimmed)
}

func LoadArchive(path string) (LoadedPack, error) {
	root, cleanup, err := extractZipToTemp(path)
	if err != nil {
		return LoadedPack{}, err
	}
	defer cleanup()

	loaded, err := LoadDir(root)
	if err != nil {
		return LoadedPack{}, err
	}
	loaded.Dir = strings.TrimSpace(path)
	return loaded, nil
}

func extractZipToTemp(path string) (string, func(), error) {
	reader, err := zip.OpenReader(strings.TrimSpace(path))
	if err != nil {
		return "", nil, fmt.Errorf("open pack archive: %w", err)
	}
	defer reader.Close()

	if len(reader.File) == 0 {
		return "", nil, fmt.Errorf("pack archive is empty")
	}
	if len(reader.File) > maxArchiveEntries {
		return "", nil, fmt.Errorf("pack archive has too many entries: %d", len(reader.File))
	}

	tempDir, err := os.MkdirTemp("", "relay-pack-*")
	if err != nil {
		return "", nil, fmt.Errorf("create temp dir for pack archive: %w", err)
	}
	cleanup := func() { _ = os.RemoveAll(tempDir) }

	var totalSize uint64
	for _, file := range reader.File {
		cleanName := filepath.Clean(file.Name)
		if cleanName == "." || cleanName == "" {
			continue
		}
		if filepath.IsAbs(cleanName) || cleanName == ".." || strings.HasPrefix(cleanName, ".."+string(filepath.Separator)) {
			cleanup()
			return "", nil, fmt.Errorf("pack archive contains invalid path %q", file.Name)
		}
		if file.FileInfo().Mode()&os.ModeSymlink != 0 {
			cleanup()
			return "", nil, fmt.Errorf("pack archive contains unsupported symlink entry %q", file.Name)
		}
		if file.UncompressedSize64 > maxArchiveFileSize {
			cleanup()
			return "", nil, fmt.Errorf("pack archive entry %q exceeds size limit", file.Name)
		}
		totalSize += file.UncompressedSize64
		if totalSize > maxArchiveTotalSize {
			cleanup()
			return "", nil, fmt.Errorf("pack archive exceeds total size limit")
		}

		targetPath := filepath.Join(tempDir, cleanName)
		relativeTarget, err := filepath.Rel(tempDir, targetPath)
		if err != nil || relativeTarget == ".." || strings.HasPrefix(relativeTarget, ".."+string(filepath.Separator)) {
			cleanup()
			return "", nil, fmt.Errorf("pack archive entry %q escapes extraction root", file.Name)
		}

		if file.FileInfo().IsDir() {
			if err := os.MkdirAll(targetPath, 0o755); err != nil {
				cleanup()
				return "", nil, fmt.Errorf("create archive dir %q: %w", file.Name, err)
			}
			continue
		}

		if err := os.MkdirAll(filepath.Dir(targetPath), 0o755); err != nil {
			cleanup()
			return "", nil, fmt.Errorf("create archive parent dir %q: %w", file.Name, err)
		}

		src, err := file.Open()
		if err != nil {
			cleanup()
			return "", nil, fmt.Errorf("open archive entry %q: %w", file.Name, err)
		}
		dst, err := os.OpenFile(targetPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
		if err != nil {
			src.Close()
			cleanup()
			return "", nil, fmt.Errorf("create archive file %q: %w", file.Name, err)
		}
		_, copyErr := io.Copy(dst, src)
		closeErr := dst.Close()
		srcErr := src.Close()
		if copyErr != nil {
			cleanup()
			return "", nil, fmt.Errorf("extract archive entry %q: %w", file.Name, copyErr)
		}
		if closeErr != nil {
			cleanup()
			return "", nil, fmt.Errorf("close archive file %q: %w", file.Name, closeErr)
		}
		if srcErr != nil {
			cleanup()
			return "", nil, fmt.Errorf("close archive entry %q: %w", file.Name, srcErr)
		}
	}

	root, err := resolvePackRoot(tempDir)
	if err != nil {
		cleanup()
		return "", nil, err
	}
	return root, cleanup, nil
}

func resolvePackRoot(extractDir string) (string, error) {
	manifestPath := filepath.Join(extractDir, "pack.json")
	if _, err := os.Stat(manifestPath); err == nil {
		return extractDir, nil
	}

	entries, err := os.ReadDir(extractDir)
	if err != nil {
		return "", fmt.Errorf("read extracted archive root: %w", err)
	}
	if len(entries) != 1 || !entries[0].IsDir() {
		return "", fmt.Errorf("pack archive must contain pack.json at root or a single top-level directory")
	}

	root := filepath.Join(extractDir, entries[0].Name())
	if _, err := os.Stat(filepath.Join(root, "pack.json")); err != nil {
		return "", fmt.Errorf("pack archive root does not contain pack.json")
	}
	return root, nil
}