package volsnap

import (
	"archive/tar"
	"compress/gzip"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

type tentry struct {
	name     string
	typeflag byte
	body     string
	linkname string
}

// buildTarGz writes an arbitrary (possibly hostile) tar.gz so the extractor's
// untrusted-input hardening can be exercised — writeArchive only emits well-
// formed reg/dir entries, which is the wrong shape for these tests.
func buildTarGz(t *testing.T, entries []tentry) string {
	t.Helper()
	dest := filepath.Join(t.TempDir(), "snap.tar.gz")
	f, err := os.Create(dest)
	if err != nil {
		t.Fatal(err)
	}
	gz := gzip.NewWriter(f)
	tw := tar.NewWriter(gz)
	for _, e := range entries {
		hdr := &tar.Header{Name: e.name, Typeflag: e.typeflag, Mode: 0o600, Linkname: e.linkname}
		switch e.typeflag {
		case tar.TypeReg:
			hdr.Size = int64(len(e.body))
		case tar.TypeChar, tar.TypeBlock:
			hdr.Devmajor, hdr.Devminor = 1, 1
		}
		if err := tw.WriteHeader(hdr); err != nil {
			t.Fatal(err)
		}
		if e.typeflag == tar.TypeReg {
			if _, err := tw.Write([]byte(e.body)); err != nil {
				t.Fatal(err)
			}
		}
	}
	if err := tw.Close(); err != nil {
		t.Fatal(err)
	}
	if err := gz.Close(); err != nil {
		t.Fatal(err)
	}
	if err := f.Close(); err != nil {
		t.Fatal(err)
	}
	return dest
}

func TestSafeExtractIndex_RoundTrip(t *testing.T) {
	arc := buildTarGz(t, []tentry{
		{name: "0/", typeflag: tar.TypeDir},
		{name: "0/a.txt", typeflag: tar.TypeReg, body: "hello"},
		{name: "0/sub/", typeflag: tar.TypeDir},
		{name: "0/sub/b.txt", typeflag: tar.TypeReg, body: "world"},
		{name: "manifest.json", typeflag: tar.TypeReg, body: "[]"},
	})
	dest := filepath.Join(t.TempDir(), "out")
	n, err := safeExtractIndex(arc, 0, dest, maxRestoreUncompressedBytes)
	if err != nil {
		t.Fatalf("extract: %v", err)
	}
	if n != int64(len("hello")+len("world")) {
		t.Errorf("written = %d, want %d", n, len("hello")+len("world"))
	}
	if got, _ := os.ReadFile(filepath.Join(dest, "a.txt")); string(got) != "hello" {
		t.Errorf("a.txt = %q", got)
	}
	if got, _ := os.ReadFile(filepath.Join(dest, "sub", "b.txt")); string(got) != "world" {
		t.Errorf("sub/b.txt = %q", got)
	}
	if _, err := os.Stat(filepath.Join(dest, "manifest.json")); !os.IsNotExist(err) {
		t.Error("manifest.json must not be extracted into the volume")
	}
}

func TestSafeExtractIndex_IsolatesIndex(t *testing.T) {
	arc := buildTarGz(t, []tentry{
		{name: "1/keep.txt", typeflag: tar.TypeReg, body: "one"},
		{name: "10/other.txt", typeflag: tar.TypeReg, body: "ten"},
		{name: "2/nope.txt", typeflag: tar.TypeReg, body: "two"},
	})
	dest := filepath.Join(t.TempDir(), "out")
	if _, err := safeExtractIndex(arc, 1, dest, maxRestoreUncompressedBytes); err != nil {
		t.Fatalf("extract: %v", err)
	}
	if _, err := os.Stat(filepath.Join(dest, "keep.txt")); err != nil {
		t.Errorf("index 1 file missing: %v", err)
	}
	// "10/" must not bleed into index 1 (prefix boundary), nor "2/".
	for _, leaked := range []string{"other.txt", "nope.txt"} {
		if _, err := os.Stat(filepath.Join(dest, leaked)); !os.IsNotExist(err) {
			t.Errorf("index 1 extraction leaked %q from another index", leaked)
		}
	}
}

func TestSafeExtractIndex_RejectsDisallowedTypes(t *testing.T) {
	cases := map[string]tentry{
		"symlink":  {name: "0/link", typeflag: tar.TypeSymlink, linkname: "/etc/passwd"},
		"hardlink": {name: "0/hard", typeflag: tar.TypeLink, linkname: "0/real"},
		"chardev":  {name: "0/cdev", typeflag: tar.TypeChar},
		"blockdev": {name: "0/bdev", typeflag: tar.TypeBlock},
		"fifo":     {name: "0/fifo", typeflag: tar.TypeFifo},
		"sparse":   {name: "0/sparse", typeflag: tar.TypeCont}, // GNU sparse / contiguous
	}
	for name, ent := range cases {
		t.Run(name, func(t *testing.T) {
			arc := buildTarGz(t, []tentry{ent})
			dest := filepath.Join(t.TempDir(), "out")
			if _, err := safeExtractIndex(arc, 0, dest, maxRestoreUncompressedBytes); err == nil {
				t.Fatalf("expected %s entry to be rejected", name)
			}
		})
	}
}

func TestSafeExtractIndex_RejectsBomb(t *testing.T) {
	arc := buildTarGz(t, []tentry{
		{name: "0/big.bin", typeflag: tar.TypeReg, body: strings.Repeat("x", 4096)},
	})
	dest := filepath.Join(t.TempDir(), "out")
	if _, err := safeExtractIndex(arc, 0, dest, 1024); err == nil {
		t.Fatal("expected extraction to abort past the decompression cap")
	}
}

func TestSafeExtractIndex_NoEscapeOutsideDest(t *testing.T) {
	// A "../" climb and an absolute member must never materialize a file
	// outside dest, regardless of which guard catches it. Note the backslash
	// case is platform-split: on Windows `..\winslip.txt` is a real climb that
	// withinDir rejects; on Linux it is a literal one-segment filename that
	// stays harmlessly inside dest (no guard fires). Both satisfy containment.
	arc := buildTarGz(t, []tentry{
		{name: "0/../../escape.txt", typeflag: tar.TypeReg, body: "pwned"},
		{name: "/abs-escape.txt", typeflag: tar.TypeReg, body: "pwned"},
		{name: `0/..\winslip.txt`, typeflag: tar.TypeReg, body: "pwned"},
		{name: "0/ok.txt", typeflag: tar.TypeReg, body: "fine"},
	})
	outParent := t.TempDir()
	dest := filepath.Join(outParent, "out")
	// May or may not error depending on platform/guard; the invariant is that
	// nothing escapes dest.
	_, _ = safeExtractIndex(arc, 0, dest, maxRestoreUncompressedBytes)

	for _, escaped := range []string{
		filepath.Join(outParent, "escape.txt"),
		filepath.Join(filepath.Dir(outParent), "escape.txt"),
		filepath.Join(outParent, "abs-escape.txt"),
		filepath.Join(outParent, "winslip.txt"),
	} {
		if _, err := os.Stat(escaped); err == nil {
			t.Errorf("zip-slip escaped to %s", escaped)
		}
	}
}