logpile/internal/parser/archive.go

package parser

import (
	"archive/tar"
	"archive/zip"
	"bufio"
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

const maxSingleFileSize = 10 * 1024 * 1024
const maxSingleFileSizeLarge = 1024 * 1024 * 1024
const maxZipArchiveSize = 50 * 1024 * 1024
const maxGzipDecompressedSize = 50 * 1024 * 1024


var supportedArchiveExt = map[string]struct{}{
	".ahs": {},
	".gz":  {},
	".tgz": {},
	".tar": {},
	".sds": {},
	".zip": {},
	".txt": {},
	".log": {},
}

// ExtractedFile represents a file extracted from archive
type ExtractedFile struct {
	Path             string
	Content          []byte
	ModTime          time.Time
	Truncated        bool
	TruncatedMessage string
}

// ExtractArchive extracts tar.gz or zip archive and returns file contents
func ExtractArchive(archivePath string) ([]ExtractedFile, error) {
	if !IsSupportedArchiveFilename(archivePath) {
		return nil, fmt.Errorf("unsupported archive format: %s", strings.ToLower(filepath.Ext(archivePath)))
	}
	ext := strings.ToLower(filepath.Ext(archivePath))

	switch ext {
	case ".ahs":
		return extractSingleFileWithLimit(archivePath, maxSingleFileSizeLarge)
	case ".gz", ".tgz":
		return extractTarGz(archivePath)
	case ".tar", ".sds":
		return extractTar(archivePath)
	case ".zip":
		return extractZip(archivePath)
	case ".txt", ".log":
		return extractSingleFileWithLimit(archivePath, maxSingleFileSize)
	default:
		return nil, fmt.Errorf("unsupported archive format: %s", ext)
	}
}

// ExtractArchiveFromReader extracts archive from reader
func ExtractArchiveFromReader(r io.Reader, filename string) ([]ExtractedFile, error) {
	if !IsSupportedArchiveFilename(filename) {
		return nil, fmt.Errorf("unsupported archive format: %s", strings.ToLower(filepath.Ext(filename)))
	}
	ext := strings.ToLower(filepath.Ext(filename))

	switch ext {
	case ".ahs":
		return extractSingleFileFromReaderWithLimit(r, filename, maxSingleFileSizeLarge)
	case ".gz", ".tgz":
		return extractTarGzFromReader(r, filename)
	case ".tar", ".sds":
		return extractTarFromReader(r)
	case ".zip":
		return extractZipFromReader(r)
	case ".txt", ".log":
		return extractSingleFileFromReaderWithLimit(r, filename, maxSingleFileSize)
	default:
		return nil, fmt.Errorf("unsupported archive format: %s", ext)
	}
}

// IsSupportedArchiveFilename reports whether filename extension is supported by archive extractor.
func IsSupportedArchiveFilename(filename string) bool {
	ext := strings.ToLower(strings.TrimSpace(filepath.Ext(filename)))
	if ext == "" {
		return false
	}
	_, ok := supportedArchiveExt[ext]
	return ok
}

// SupportedArchiveExtensions returns sorted list of archive/file extensions
// accepted by archive extractor.
func SupportedArchiveExtensions() []string {
	out := make([]string, 0, len(supportedArchiveExt))
	for ext := range supportedArchiveExt {
		out = append(out, ext)
	}
	sort.Strings(out)
	return out
}

func extractTarGz(archivePath string) ([]ExtractedFile, error) {
	f, err := os.Open(archivePath)
	if err != nil {
		return nil, fmt.Errorf("open archive: %w", err)
	}
	defer f.Close()

	return extractTarGzFromReader(f, filepath.Base(archivePath))
}

func extractTar(archivePath string) ([]ExtractedFile, error) {
	f, err := os.Open(archivePath)
	if err != nil {
		return nil, fmt.Errorf("open archive: %w", err)
	}
	defer f.Close()

	return extractTarFromReader(f)
}

func extractTarFromReader(r io.Reader) ([]ExtractedFile, error) {
	tr := tar.NewReader(r)
	var files []ExtractedFile

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

		// Skip directories
		if header.Typeflag == tar.TypeDir {
			continue
		}

		// Skip large files (>10MB)
		if header.Size > 10*1024*1024 {
			continue
		}

		content, err := io.ReadAll(tr)
		if err != nil {
			return nil, fmt.Errorf("read file %s: %w", header.Name, err)
		}

		files = append(files, ExtractedFile{
			Path:    header.Name,
			Content: content,
			ModTime: header.ModTime,
		})
	}

	return files, nil
}

func extractTarGzFromReader(r io.Reader, filename string) ([]ExtractedFile, error) {
	gzr, err := gzip.NewReader(r)
	if err != nil {
		return nil, fmt.Errorf("gzip reader: %w", err)
	}
	defer gzr.Close()

	// Peek the first tar block (512 bytes) to decide whether this is a tar
	// archive or a single gzipped file, without consuming/truncating the
	// rest of the stream. A raw byte-level cap on the whole decompressed
	// payload would otherwise cut a tar archive mid-entry and corrupt its
	// structure, turning "large but valid archive" into "unexpected EOF"
	// for the entire file instead of extracting what's actually there.
	buf := bufio.NewReaderSize(gzr, 64*1024)
	head, peekErr := buf.Peek(512)
	looksLikeTar := peekErr == nil || (peekErr == io.EOF && len(head) > 0)
	if looksLikeTar {
		if _, err := tar.NewReader(bytes.NewReader(head)).Next(); err != nil {
			looksLikeTar = false
		}
	}

	if looksLikeTar {
		return extractTarEntriesWithLimit(buf, maxSingleFileSizeLarge)
	}

	// Not a tar archive - treat as a single gzipped file. A hard byte cap
	// is safe here since there's no container structure to corrupt.
	decompressed, err := io.ReadAll(io.LimitReader(buf, maxGzipDecompressedSize+1))
	if err != nil {
		return nil, fmt.Errorf("read gzip content: %w", err)
	}
	gzipTruncated := len(decompressed) > maxGzipDecompressedSize
	if gzipTruncated {
		decompressed = decompressed[:maxGzipDecompressedSize]
	}

	baseName := strings.TrimSuffix(filename, ".gz")
	if gzr.Name != "" {
		baseName = gzr.Name
	}

	file := ExtractedFile{
		Path:    baseName,
		Content: decompressed,
		ModTime: gzr.ModTime,
	}
	if gzipTruncated {
		file.Truncated = true
		file.TruncatedMessage = fmt.Sprintf(
			"decompressed gzip content exceeded %d bytes and was truncated",
			maxGzipDecompressedSize,
		)
	}

	return []ExtractedFile{file}, nil
}

// extractTarEntriesWithLimit streams tar entries from r, skipping individual
// files over 10MB, and stops once the cumulative extracted content reaches
// totalLimit — always at an entry boundary, never mid-file.
func extractTarEntriesWithLimit(r io.Reader, totalLimit int64) ([]ExtractedFile, error) {
	tr := tar.NewReader(r)
	var files []ExtractedFile
	var totalExtracted int64

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

		if header.Typeflag == tar.TypeDir {
			continue
		}
		if header.Size > 10*1024*1024 {
			continue
		}
		if totalExtracted+header.Size > totalLimit {
			break
		}

		content, err := io.ReadAll(tr)
		if err != nil {
			return nil, fmt.Errorf("read file %s: %w", header.Name, err)
		}

		files = append(files, ExtractedFile{
			Path:    header.Name,
			Content: content,
			ModTime: header.ModTime,
		})
		totalExtracted += int64(len(content))
	}

	return files, nil
}

func extractZip(archivePath string) ([]ExtractedFile, error) {
	r, err := zip.OpenReader(archivePath)
	if err != nil {
		return nil, fmt.Errorf("open zip: %w", err)
	}
	defer r.Close()

	var files []ExtractedFile

	for _, f := range r.File {
		if f.FileInfo().IsDir() {
			continue
		}

		// Skip large files (>10MB)
		if f.FileInfo().Size() > 10*1024*1024 {
			continue
		}

		rc, err := f.Open()
		if err != nil {
			return nil, fmt.Errorf("open file %s: %w", f.Name, err)
		}

		content, err := io.ReadAll(rc)
		rc.Close()
		if err != nil {
			return nil, fmt.Errorf("read file %s: %w", f.Name, err)
		}

		files = append(files, ExtractedFile{
			Path:    f.Name,
			Content: content,
			ModTime: f.Modified,
		})
	}

	return files, nil
}

func extractZipFromReader(r io.Reader) ([]ExtractedFile, error) {
	// Read all data into memory with a hard cap
	data, err := io.ReadAll(io.LimitReader(r, maxZipArchiveSize+1))
	if err != nil {
		return nil, fmt.Errorf("read zip data: %w", err)
	}
	if len(data) > maxZipArchiveSize {
		return nil, fmt.Errorf("zip too large: max %d bytes", maxZipArchiveSize)
	}

	// Create a ReaderAt from the byte slice
	readerAt := bytes.NewReader(data)

	// Open the zip archive
	zipReader, err := zip.NewReader(readerAt, int64(len(data)))
	if err != nil {
		return nil, fmt.Errorf("open zip: %w", err)
	}

	var files []ExtractedFile

	for _, f := range zipReader.File {
		if f.FileInfo().IsDir() {
			continue
		}

		// Skip large files (>10MB)
		if f.FileInfo().Size() > 10*1024*1024 {
			continue
		}

		rc, err := f.Open()
		if err != nil {
			return nil, fmt.Errorf("open file %s: %w", f.Name, err)
		}

		content, err := io.ReadAll(rc)
		rc.Close()
		if err != nil {
			return nil, fmt.Errorf("read file %s: %w", f.Name, err)
		}

		files = append(files, ExtractedFile{
			Path:    f.Name,
			Content: content,
			ModTime: f.Modified,
		})
	}

	return files, nil
}

func extractSingleFileWithLimit(path string, limit int64) ([]ExtractedFile, error) {
	info, err := os.Stat(path)
	if err != nil {
		return nil, fmt.Errorf("stat file: %w", err)
	}
	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open file: %w", err)
	}
	defer f.Close()

	files, err := extractSingleFileFromReaderWithLimit(f, filepath.Base(path), limit)
	if err != nil {
		return nil, err
	}
	if len(files) > 0 {
		files[0].ModTime = info.ModTime()
	}
	return files, nil
}

func extractSingleFileFromReaderWithLimit(r io.Reader, filename string, limit int64) ([]ExtractedFile, error) {
	content, err := io.ReadAll(io.LimitReader(r, limit+1))
	if err != nil {
		return nil, fmt.Errorf("read file content: %w", err)
	}
	truncated := int64(len(content)) > limit
	if truncated {
		content = content[:limit]
	}

	file := ExtractedFile{
		Path:    filepath.Base(filename),
		Content: content,
	}
	if truncated {
		file.Truncated = true
		file.TruncatedMessage = fmt.Sprintf(
			"file exceeded %d bytes and was truncated",
			limit,
		)
	}

	return []ExtractedFile{file}, nil
}

// FindFileByPattern finds files matching pattern in extracted files
func FindFileByPattern(files []ExtractedFile, patterns ...string) []ExtractedFile {
	var result []ExtractedFile
	for _, f := range files {
		for _, pattern := range patterns {
			if strings.Contains(strings.ToLower(f.Path), strings.ToLower(pattern)) {
				result = append(result, f)
				break
			}
		}
	}
	return result
}

// FindFileByName finds file by exact name (case-insensitive)
func FindFileByName(files []ExtractedFile, name string) *ExtractedFile {
	for _, f := range files {
		if strings.EqualFold(filepath.Base(f.Path), name) {
			return &f
		}
	}
	return nil
}