From 82d900220b77f410b981fd82b824f70c399046b3 Mon Sep 17 00:00:00 2001
From: Ville Aikas <vaikas@chainguard.dev>
Date: Sun, 25 Sep 2022 15:58:48 -0700
Subject: [PATCH] better tests, add CompressFS that uses FS for walking.

Signed-off-by: Ville Aikas <vaikas@chainguard.dev>
---
 pkg/repo/repo.go      | 53 +++++++++++++++++++++++++++++++++++++++++++
 pkg/repo/repo_test.go | 50 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 102 insertions(+), 1 deletion(-)

diff --git a/pkg/repo/repo.go b/pkg/repo/repo.go
index ad7ad36a1..027601faa 100644
--- a/pkg/repo/repo.go
+++ b/pkg/repo/repo.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"io/fs"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -156,6 +157,58 @@ func Compress(src string, buf io.Writer) error {
 	return zr.Close()
 }
 
+// CompressFS archives a TUF repository so that it can be written to Secret
+// for later use.
+func CompressFS(fsys fs.FS, buf io.Writer, skipDirs map[string]bool) error {
+	// tar > gzip > buf
+	zr := gzip.NewWriter(buf)
+	tw := tar.NewWriter(zr)
+
+	err := fs.WalkDir(fsys, "repository", func(file string, d fs.DirEntry, err error) error {
+		// Skip the 'keys' and 'staged' directory
+		if d.IsDir() && skipDirs[d.Name()] {
+			return filepath.SkipDir
+		}
+
+		// Stat the file to get the details of it.
+		fi, err := fs.Stat(fsys, file)
+		if err != nil {
+			return fmt.Errorf("fs.Stat %s: %w", file, err)
+		}
+		header, err := tar.FileInfoHeader(fi, file)
+		if err != nil {
+			return fmt.Errorf("FileInfoHeader %s: %w", file, err)
+		}
+		header.Name = filepath.ToSlash(file)
+		if err := tw.WriteHeader(header); err != nil {
+			return err
+		}
+		// For files, write the contents.
+		if !d.IsDir() {
+			data, err := fsys.Open(file)
+			if err != nil {
+				return fmt.Errorf("opening %s: %w", file, err)
+			}
+			if _, err := io.Copy(tw, data); err != nil {
+				return fmt.Errorf("copying %s: %w", file, err)
+			}
+		}
+		return nil
+	})
+
+	if err != nil {
+		tw.Close()
+		zr.Close()
+		return fmt.Errorf("WalkDir: %w", err)
+	}
+
+	if err := tw.Close(); err != nil {
+		zr.Close()
+		return fmt.Errorf("tar.NewWriter Close(): %w", err)
+	}
+	return zr.Close()
+}
+
 // check for path traversal and correct forward slashes
 func validRelPath(p string) bool {
 	if p == "" || strings.Contains(p, `\`) || strings.HasPrefix(p, "/") || strings.Contains(p, "../") {
diff --git a/pkg/repo/repo_test.go b/pkg/repo/repo_test.go
index 0717657e5..fd7394248 100644
--- a/pkg/repo/repo_test.go
+++ b/pkg/repo/repo_test.go
@@ -97,7 +97,55 @@ func TestCompressUncompress(t *testing.T) {
 		t.Fatalf("Failed to uncompress: %v", err)
 	}
 	// Then check that files have been uncompressed there.
-	// TODO(vaikas): Do a full diff?
+	meta, err := repo.GetMeta()
+	if err != nil {
+		t.Errorf("Failed to GetMeta: %s", err)
+	}
+	root := meta["root.json"]
+
+	// This should have roundtripped to the new directory.
+	rtRoot, err := os.ReadFile(filepath.Join(dstDir, "repository", "root.json"))
+	if err != nil {
+		t.Errorf("Failed to read the roundtripped root %v", err)
+	}
+	if bytes.Compare(root, rtRoot) != 0 {
+		t.Errorf("Roundtripped root differs:\n%s\n%s", string(root), string(rtRoot))
+	}
+
+	// As well as, say rekor.pub under targets dir
+	rtRekor, err := os.ReadFile(filepath.Join(dstDir, "repository", "targets", "rekor.pub"))
+	if err != nil {
+		t.Errorf("Failed to read the roundtripped rekor %v", err)
+	}
+	if bytes.Compare(files["rekor.pub"], rtRekor) != 0 {
+		t.Errorf("Roundtripped rekor differs:\n%s\n%s", rekorPublicKey, string(rtRekor))
+	}
+
+}
+
+func TestCompressUncompressFS(t *testing.T) {
+	files := map[string][]byte{
+		"fulcio_v1.crt.pem": []byte(fulcioRootCert),
+		"ctfe.pub":          []byte(ctlogPublicKey),
+		"rekor.pub":         []byte(rekorPublicKey),
+	}
+	repo, dir, err := CreateRepo(context.Background(), files)
+	if err != nil {
+		t.Fatalf("Failed to CreateRepo: %s", err)
+	}
+	defer os.RemoveAll(dir)
+
+	var buf bytes.Buffer
+	fsys := os.DirFS(dir)
+	if err = CompressFS(fsys, &buf, map[string]bool{"keys": true, "staged": true}); err != nil {
+		t.Fatalf("Failed to compress: %v", err)
+	}
+	os.WriteFile("/tmp/newcompressed", buf.Bytes(), os.ModePerm)
+	dstDir := t.TempDir()
+	if err = Uncompress(&buf, dstDir); err != nil {
+		t.Fatalf("Failed to uncompress: %v", err)
+	}
+	// Then check that files have been uncompressed there.
 	meta, err := repo.GetMeta()
 	if err != nil {
 		t.Errorf("Failed to GetMeta: %s", err)