fix: move internal/ to pkg/, verify checksum, minor fixes

Author: jonstjohn

README.md

@@ -171,6 +171,7 @@ ls export-contents/*.schema.txt
 - **Schema only** - Table schemas are exported, but **no actual table data** is included
 - **Read-only** - The tool only reads data and makes no modifications to your cluster
 - **Local export** - All data is written to a local zip file under your control
+- **Verified updates** - `workload-exporter update` verifies the SHA256 checksum of the downloaded binary against the checksums published with each GitHub release before installing
 
 ## Requirements
 

cmd/update.go

@@ -6,7 +6,7 @@ import (
 	"os"
 	"time"
 
-	"github.com/cockroachlabs/workload-exporter/internal/update"
+	"github.com/cockroachlabs/workload-exporter/pkg/update"
 	"github.com/spf13/cobra"
 )
 
@@ -27,7 +27,9 @@ func newUpdateCmd() *cobra.Command {
 				return runUpdateCheck(cmd.Context())
 			}
 
-			return update.PerformUpdate(cmd.Context(), os.Stdout, Version)
+			ctx, cancel := context.WithTimeout(cmd.Context(), 5*time.Minute)
+			defer cancel()
+			return update.PerformUpdate(ctx, os.Stdout, Version)
 		},
 	}
 
@@ -53,7 +55,7 @@ func runUpdateCheck(ctx context.Context) error {
 		return fmt.Errorf("update check failed: %w", err)
 	}
 
-	if result.TagName == Version {
+	if !update.IsNewer(result.TagName, Version) {
 		fmt.Println("up to date")
 	} else {
 		fmt.Printf("\nnew version available: %s (current: %s)\n", result.TagName, Version)

docs/DEVELOPMENT.md

@@ -4,7 +4,7 @@ This guide is for developers who want to build, test, or contribute to the workl
 
 ## Prerequisites
 
-- Go 1.18 or later
+- Go 1.23 or later
 - Git
 
 ## Building from Source
@@ -64,11 +64,17 @@ workload-exporter/
 ├── cmd/               # CLI commands
 │   ├── root.go       # Root command
 │   ├── export.go     # Export command
+│   ├── update.go     # Update command
 │   └── version.go    # Version command
 ├── pkg/
-│   └── export/       # Core export functionality
-│       ├── exporter.go       # Main exporter logic
-│       └── exporter_test.go  # Unit tests
+│   ├── export/       # Core export functionality
+│   │   ├── exporter.go       # Main exporter logic
+│   │   └── exporter_test.go  # Unit tests
+│   └── update/       # Self-update functionality
+│       ├── update.go         # Update check, caching, and semver comparison
+│       ├── selfupdate.go     # Binary download, checksum verification, and replacement
+│       ├── update_test.go    # Unit tests for update checking
+│       └── selfupdate_test.go # Unit tests for self-update
 ├── docs/             # Documentation
 ├── Makefile          # Build automation
 └── go.mod            # Go dependencies

internal/update/selfupdate.go pkg/update/selfupdate.gointernal/update/selfupdate.go renamed to pkg/update/selfupdate.go

@@ -6,6 +6,9 @@ import (
 	"bytes"
 	"compress/gzip"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -17,6 +20,11 @@ import (
 	"time"
 )
 
+const (
+	maxBinarySize   = 100 * 1024 * 1024 // 100 MB guard against runaway downloads
+	maxChecksumSize = 1 << 20           // 1 MB, well above any real checksums.txt
+)
+
 // UpdateDeps holds the external dependencies for PerformUpdate.
 // Tests inject fakes; production uses defaultUpdateDeps().
 type UpdateDeps struct {
@@ -42,7 +50,9 @@ func defaultUpdateDeps(version string) UpdateDeps {
 		CheckLatest: func(ctx context.Context) (*ReleaseInfo, error) {
 			return Check(ctx, version)
 		},
-		Download:       defaultDownload,
+		Download: func(ctx context.Context, v string) ([]byte, error) {
+			return defaultDownload(ctx, http.DefaultClient, v)
+		},
 		RunVersion:     defaultRunVersion,
 		CurrentVersion: version,
 	}
@@ -74,30 +84,64 @@ func assetDownloadURL(version string) string {
 	return fmt.Sprintf("https://github.com/%s/releases/download/%s/%s", repo, version, assetName(version))
 }
 
-func defaultDownload(ctx context.Context, version string) ([]byte, error) {
-	downloadURL := assetDownloadURL(version)
-	req, err := http.NewRequestWithContext(ctx, "GET", downloadURL, nil)
+func defaultDownload(ctx context.Context, client HTTPDoer, version string) ([]byte, error) {
+	// Download the release archive.
+	archiveData, err := fetchURL(ctx, client, assetDownloadURL(version), maxBinarySize)
+	if err != nil {
+		return nil, fmt.Errorf("downloading asset: %w", err)
+	}
+
+	// Download checksums.txt and verify the archive before extracting.
+	checksumURL := fmt.Sprintf("https://github.com/%s/releases/download/%s/checksums.txt", repo, version)
+	checksumData, err := fetchURL(ctx, client, checksumURL, maxChecksumSize)
 	if err != nil {
+		return nil, fmt.Errorf("downloading checksums: %w", err)
+	}
+	if err := verifyChecksum(archiveData, assetName(version), checksumData); err != nil {
 		return nil, err
 	}
 
-	resp, err := http.DefaultClient.Do(req)
+	// Extract binary from archive.
+	return extractBinary(archiveData, version)
+}
+
+// fetchURL performs a GET request and returns the response body, bounded by limit bytes.
+func fetchURL(ctx context.Context, client HTTPDoer, url string, limit int64) ([]byte, error) {
+	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
 	if err != nil {
-		return nil, fmt.Errorf("downloading asset: %w", err)
+		return nil, err
+	}
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, err
 	}
 	defer func() { _ = resp.Body.Close() }()
-
 	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("asset download returned %d", resp.StatusCode)
+		return nil, fmt.Errorf("HTTP %d for %s", resp.StatusCode, url)
 	}
+	return io.ReadAll(io.LimitReader(resp.Body, limit))
+}
 
-	archiveData, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return nil, fmt.Errorf("reading asset: %w", err)
+// verifyChecksum checks the SHA256 of data against the entry for filename in
+// a checksums.txt file (sha256sum format: "<hash>  <filename>").
+func verifyChecksum(data []byte, filename string, checksumFile []byte) error {
+	for _, line := range strings.Split(string(checksumFile), "\n") {
+		line = strings.TrimSpace(line)
+		if line == "" {
+			continue
+		}
+		parts := strings.Fields(line)
+		if len(parts) != 2 || parts[1] != filename {
+			continue
+		}
+		sum := sha256.Sum256(data)
+		actual := hex.EncodeToString(sum[:])
+		if actual != parts[0] {
+			return fmt.Errorf("checksum mismatch for %s: got %s, want %s", filename, actual, parts[0])
+		}
+		return nil
 	}
-
-	// Extract binary from archive.
-	return extractBinary(archiveData, version)
+	return fmt.Errorf("no checksum entry found for %s in checksums.txt", filename)
 }
 
 // extractBinary extracts the workload-exporter binary from a .tar.gz or .zip archive.
@@ -120,14 +164,14 @@ func extractFromTarGz(data []byte, wantName string) ([]byte, error) {
 	tr := tar.NewReader(gz)
 	for {
 		hdr, err := tr.Next()
-		if err == io.EOF {
+		if errors.Is(err, io.EOF) {
 			break
 		}
 		if err != nil {
 			return nil, fmt.Errorf("reading tar: %w", err)
 		}
 		if hdr.Name == wantName || filepath.Base(hdr.Name) == wantName {
-			return io.ReadAll(tr)
+			return io.ReadAll(io.LimitReader(tr, maxBinarySize))
 		}
 	}
 	return nil, fmt.Errorf("binary %s not found in archive", wantName)
@@ -145,7 +189,7 @@ func extractFromZip(data []byte, wantName string) ([]byte, error) {
 				return nil, err
 			}
 			defer func() { _ = rc.Close() }()
-			return io.ReadAll(rc)
+			return io.ReadAll(io.LimitReader(rc, maxBinarySize))
 		}
 	}
 	return nil, fmt.Errorf("binary %s not found in archive", wantName)
@@ -180,7 +224,7 @@ func performUpdate(ctx context.Context, w io.Writer, deps UpdateDeps) error {
 		return nil
 	}
 
-	if release.TagName == deps.CurrentVersion {
+	if !semverGreater(release.TagName, deps.CurrentVersion) {
 		_, _ = fmt.Fprintf(w, "already up to date (%s)\n", deps.CurrentVersion)
 		return nil
 	}
@@ -232,6 +276,9 @@ func performUpdate(ctx context.Context, w io.Writer, deps UpdateDeps) error {
 	if !strings.Contains(newVersion, "workload-exporter") {
 		return fmt.Errorf("sanity check failed: unexpected version output: %s", newVersion)
 	}
+	if !strings.Contains(newVersion, release.TagName) {
+		return fmt.Errorf("sanity check failed: version output %q does not contain expected %s", newVersion, release.TagName)
+	}
 
 	_, _ = fmt.Fprintf(w, "verified: %s\n", newVersion)
 
@@ -265,9 +312,9 @@ func copyFile(src, dst string) error {
 	if err != nil {
 		return err
 	}
-	defer func() { _ = out.Close() }()
 
 	if _, err := io.Copy(out, in); err != nil {
+		_ = out.Close()
 		return err
 	}
 	return out.Close()

internal/update/selfupdate_test.go pkg/update/selfupdate_test.gointernal/update/selfupdate_test.go renamed to pkg/update/selfupdate_test.go

@@ -3,6 +3,8 @@ package update
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -162,6 +164,106 @@ func TestPerformUpdate_BadVersionOutput(t *testing.T) {
 	}
 }
 
+func TestPerformUpdate_VersionMismatch(t *testing.T) {
+	_, binaryPath := setupBinary(t)
+	deps := fakeDeps(t, binaryPath)
+	// RunVersion returns the wrong version tag.
+	deps.RunVersion = func(_ context.Context, _ string) (string, error) {
+		return "workload-exporter version v1.0.0", nil
+	}
+
+	var buf bytes.Buffer
+	err := performUpdate(context.Background(), &buf, deps)
+	if err == nil {
+		t.Fatal("expected error for version mismatch, got nil")
+	}
+	if !strings.Contains(err.Error(), "sanity check failed") {
+		t.Errorf("error = %q, want it to contain 'sanity check failed'", err)
+	}
+}
+
+func TestPerformUpdate_Downgrade(t *testing.T) {
+	_, binaryPath := setupBinary(t)
+	deps := fakeDeps(t, binaryPath)
+	// Simulate a case where the installed binary is already newer than "latest".
+	deps.CurrentVersion = "v3.0.0"
+	deps.CheckLatest = func(_ context.Context) (*ReleaseInfo, error) {
+		return &ReleaseInfo{TagName: "v2.0.0"}, nil
+	}
+
+	var buf bytes.Buffer
+	err := performUpdate(context.Background(), &buf, deps)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !strings.Contains(buf.String(), "already up to date") {
+		t.Errorf("expected 'already up to date' for downgrade scenario:\n%s", buf.String())
+	}
+
+	// Binary should be untouched.
+	data, _ := os.ReadFile(binaryPath)
+	if string(data) != "old-binary" {
+		t.Error("binary should not have been replaced in downgrade scenario")
+	}
+}
+
+// --- verifyChecksum tests ---
+
+func checksumLine(data []byte, filename string) string {
+	sum := sha256.Sum256(data)
+	return hex.EncodeToString(sum[:]) + "  " + filename + "\n"
+}
+
+func TestVerifyChecksum_Valid(t *testing.T) {
+	data := []byte("binary content")
+	checksumFile := checksumLine(data, "myfile.tar.gz")
+
+	if err := verifyChecksum(data, "myfile.tar.gz", []byte(checksumFile)); err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+}
+
+func TestVerifyChecksum_ValidAmongMultipleEntries(t *testing.T) {
+	data := []byte("binary content")
+	checksumFile := "aaaa  other-file.zip\n" + checksumLine(data, "myfile.tar.gz") + "bbbb  another.tar.gz\n"
+
+	if err := verifyChecksum(data, "myfile.tar.gz", []byte(checksumFile)); err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+}
+
+func TestVerifyChecksum_Mismatch(t *testing.T) {
+	data := []byte("binary content")
+	wrong := strings.Repeat("a", 64) + "  myfile.tar.gz\n"
+
+	err := verifyChecksum(data, "myfile.tar.gz", []byte(wrong))
+	if err == nil {
+		t.Fatal("expected checksum mismatch error")
+	}
+	if !strings.Contains(err.Error(), "checksum mismatch") {
+		t.Errorf("error %q should mention 'checksum mismatch'", err)
+	}
+}
+
+func TestVerifyChecksum_NotFound(t *testing.T) {
+	checksumFile := "abc123  other-file.tar.gz\n"
+
+	err := verifyChecksum([]byte("data"), "myfile.tar.gz", []byte(checksumFile))
+	if err == nil {
+		t.Fatal("expected not-found error")
+	}
+	if !strings.Contains(err.Error(), "no checksum entry") {
+		t.Errorf("error %q should mention 'no checksum entry'", err)
+	}
+}
+
+func TestVerifyChecksum_EmptyFile(t *testing.T) {
+	err := verifyChecksum([]byte("data"), "myfile.tar.gz", []byte(""))
+	if err == nil {
+		t.Fatal("expected error for empty checksums file")
+	}
+}
+
 func TestPerformUpdate_StagedFileCleanedUp(t *testing.T) {
 	_, binaryPath := setupBinary(t)
 	deps := fakeDeps(t, binaryPath)