HexaHost-GameCloud/apps/node-agent/internal/archive/archive.go

package archive

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

// CreateTarGz archives sourceDir into outputPath and returns checksum and size.
func CreateTarGz(sourceDir, outputPath string) (sha256sum string, size int64, err error) {
	sourceDir = filepath.Clean(sourceDir)

	info, err := os.Stat(sourceDir)
	if err != nil {
		return "", 0, err
	}
	if !info.IsDir() {
		return "", 0, fmt.Errorf("source is not a directory")
	}

	if err := os.MkdirAll(filepath.Dir(outputPath), 0o755); err != nil {
		return "", 0, err
	}

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

	hasher := sha256.New()
	writer := io.MultiWriter(file, hasher)
	gzipWriter := gzip.NewWriter(writer)
	tarWriter := tar.NewWriter(gzipWriter)

	err = filepath.Walk(sourceDir, func(path string, entry os.FileInfo, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}

		rel, err := filepath.Rel(sourceDir, path)
		if err != nil {
			return err
		}
		if rel == "." {
			return nil
		}

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

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

		if entry.IsDir() {
			return nil
		}

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

		_, err = io.Copy(tarWriter, sourceFile)
		return err
	})
	if err != nil {
		return "", 0, err
	}

	if err := tarWriter.Close(); err != nil {
		return "", 0, err
	}
	if err := gzipWriter.Close(); err != nil {
		return "", 0, err
	}

	stat, err := file.Stat()
	if err != nil {
		return "", 0, err
	}

	return hex.EncodeToString(hasher.Sum(nil)), stat.Size(), nil
}

// CreateTarGzFromPaths archives one or more paths under dataRoot into outputPath.
func CreateTarGzFromPaths(dataRoot string, relativePaths []string, outputPath string) (sha256sum string, size int64, err error) {
	dataRoot = filepath.Clean(dataRoot)
	if err := os.MkdirAll(filepath.Dir(outputPath), 0o755); err != nil {
		return "", 0, err
	}

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

	hasher := sha256.New()
	writer := io.MultiWriter(file, hasher)
	gzipWriter := gzip.NewWriter(writer)
	tarWriter := tar.NewWriter(gzipWriter)

	for _, rel := range relativePaths {
		abs := filepath.Join(dataRoot, filepath.FromSlash(rel))
		abs = filepath.Clean(abs)
		if !strings.HasPrefix(abs, dataRoot+string(os.PathSeparator)) && abs != dataRoot {
			return "", 0, fmt.Errorf("path escapes data root: %s", rel)
		}

		info, err := os.Stat(abs)
		if err != nil {
			return "", 0, err
		}

		baseName := filepath.Base(abs)
		if info.IsDir() {
			err = filepath.Walk(abs, func(path string, entry os.FileInfo, walkErr error) error {
				if walkErr != nil {
					return walkErr
				}
				innerRel, err := filepath.Rel(abs, path)
				if err != nil {
					return err
				}
				if innerRel == "." {
					return nil
				}
				headerName := filepath.ToSlash(filepath.Join(baseName, innerRel))
				return writeTarEntry(tarWriter, path, entry, headerName)
			})
			if err != nil {
				return "", 0, err
			}
			continue
		}

		if err := writeTarEntry(tarWriter, abs, info, baseName); err != nil {
			return "", 0, err
		}
	}

	if err := tarWriter.Close(); err != nil {
		return "", 0, err
	}
	if err := gzipWriter.Close(); err != nil {
		return "", 0, err
	}

	stat, err := file.Stat()
	if err != nil {
		return "", 0, err
	}

	return hex.EncodeToString(hasher.Sum(nil)), stat.Size(), nil
}

func writeTarEntry(tarWriter *tar.Writer, path string, entry os.FileInfo, headerName string) error {
	header, err := tar.FileInfoHeader(entry, "")
	if err != nil {
		return err
	}
	header.Name = headerName
	if err := tarWriter.WriteHeader(header); err != nil {
		return err
	}
	if entry.IsDir() {
		return nil
	}
	sourceFile, err := os.Open(path)
	if err != nil {
		return err
	}
	defer sourceFile.Close()
	_, err = io.Copy(tarWriter, sourceFile)
	return err
}

// ExtractTarGz extracts archivePath into destDir.
func ExtractTarGz(archivePath, destDir string) error {
	destDir = filepath.Clean(destDir)
	if err := os.MkdirAll(destDir, 0o755); err != nil {
		return err
	}

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

	gzipReader, err := gzip.NewReader(file)
	if err != nil {
		return err
	}
	defer gzipReader.Close()

	tarReader := tar.NewReader(gzipReader)

	for {
		header, err := tarReader.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		target := filepath.Join(destDir, filepath.FromSlash(header.Name))
		if !strings.HasPrefix(filepath.Clean(target), destDir+string(os.PathSeparator)) && filepath.Clean(target) != destDir {
			return fmt.Errorf("archive entry escapes destination: %s", header.Name)
		}

		switch header.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(target, 0o755); err != nil {
				return err
			}
		case tar.TypeReg:
			if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
				return err
			}

			out, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.FileMode(header.Mode))
			if err != nil {
				return err
			}

			if _, err := io.Copy(out, tarReader); err != nil {
				out.Close()
				return err
			}
			out.Close()
		default:
			continue
		}
	}

	return nil
}