Merge pull request #5 from sigman78/sigman/storage-if

Storage interface for FS access

Author: sigman78

internal/wayback/css.go

@@ -2,7 +2,6 @@ package wayback
 
 import (
 	"net/url"
-	"os"
 	"path/filepath"
 	"regexp"
 	"strings"
@@ -92,12 +91,12 @@ func RewriteCSSContent(css, pageURL string, cfg *Config, idx *SnapshotIndex) str
 	return css
 }
 
-// RewriteCSSFile reads a CSS file, rewrites its URLs, and writes it back.
-func RewriteCSSFile(filePath, pageURL string, cfg *Config, idx *SnapshotIndex) error {
-	data, err := os.ReadFile(filePath) //nolint:gosec // G304: path is written by this program
+// RewriteCSSFile reads a CSS file from storage, rewrites its URLs, and writes it back.
+func RewriteCSSFile(store Storage, logicalPath, pageURL string, cfg *Config, idx *SnapshotIndex) error {
+	data, err := store.Get(logicalPath)
 	if err != nil {
 		return err
 	}
 	rewritten := RewriteCSSContent(string(data), pageURL, cfg, idx)
-	return os.WriteFile(filePath, []byte(rewritten), 0600)
+	return store.PutBytes(logicalPath, []byte(rewritten))
 }

internal/wayback/downloader.go

@@ -1,13 +1,12 @@
 package wayback
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"io"
 	"log"
 	"net/http"
-	"os"
-	"path/filepath"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -33,8 +32,9 @@ type Config struct {
 	DownloadExternalAssets bool
 	Debug                  bool
 	StopOnError            bool
-	CDXRatePerMin          int // CDX API requests per minute (default 60)
-	CDXMaxRetries          int // max retry attempts on throttle/5xx (default 5)
+	CDXRatePerMin          int     // CDX API requests per minute (default 60)
+	CDXMaxRetries          int     // max retry attempts on throttle/5xx (default 5)
+	Storage                Storage // if nil, NewLocalStorage(Directory) is used
 }
 
 var downloadHTTPClient = &http.Client{
@@ -69,8 +69,9 @@ func DownloadAll(cfg *Config) error {
 		fmt.Printf("Found %d unique snapshots to download.\n", total)
 	}
 
-	if err := os.MkdirAll(cfg.Directory, 0750); err != nil {
-		return fmt.Errorf("create output dir: %w", err)
+	store := cfg.Storage
+	if store == nil {
+		store = NewLocalStorage(cfg.Directory)
 	}
 
 	pool, err := ants.NewPool(cfg.Threads)
@@ -91,7 +92,7 @@ func DownloadAll(cfg *Config) error {
 			}
 			errCh := make(chan error, 1)
 			if err := pool.Submit(func() {
-				errCh <- downloadOne(ctx, s, cfg, idx, dlProg)
+				errCh <- downloadOne(ctx, s, cfg, store, idx, dlProg)
 			}); err != nil {
 				return fmt.Errorf("submit task: %w", err)
 			}
@@ -119,17 +120,16 @@ func DownloadAll(cfg *Config) error {
 }
 
 // downloadOne downloads a single snapshot and optionally rewrites its links.
-func downloadOne(ctx context.Context, snap Snapshot, cfg *Config, idx *SnapshotIndex, dlProg *Progress) error {
+func downloadOne(ctx context.Context, snap Snapshot, cfg *Config, store Storage, idx *SnapshotIndex, dlProg *Progress) error {
 
 	if ctx.Err() != nil {
 		return ctx.Err()
 	}
 
-	localPath := URLToLocalPath(snap.FileURL, cfg.PrettyPath)
-	localPath = filepath.Join(cfg.Directory, filepath.FromSlash(localPath))
+	logicalPath := URLToLocalPath(snap.FileURL, cfg.PrettyPath)
 
 	// Skip existing files
-	if _, err := os.Stat(localPath); err == nil {
+	if store.Exists(logicalPath) {
 		dlProg.Inc()
 		return nil
 	}
@@ -160,53 +160,27 @@ func downloadOne(ctx context.Context, snap Snapshot, cfg *Config, idx *SnapshotI
 		return fmt.Errorf("HTTP %d for %s", resp.StatusCode, waybackURL)
 	}
 
-	// Ensure parent directory exists
-	if err := os.MkdirAll(filepath.Dir(localPath), 0750); err != nil {
-		return fmt.Errorf("mkdirall: %w", err)
-	}
-
-	// Stream to temp file, then rename atomically
-	tmpFile, err := os.CreateTemp(filepath.Dir(localPath), ".wbdl-*")
-	if err != nil {
-		return fmt.Errorf("create temp: %w", err)
-	}
-	tmpName := tmpFile.Name()
-	defer func() {
-		_ = tmpFile.Close()
-		_ = os.Remove(tmpName) // no-op if renamed
-	}()
-
-	// Read first 512 bytes for content sniffing
+	// Read first 512 bytes for content sniffing, then stream remainder via storage
 	first := make([]byte, 512)
 	n, _ := io.ReadFull(resp.Body, first)
 	first = first[:n]
 
-	if _, err := tmpFile.Write(first); err != nil {
-		return fmt.Errorf("write first bytes: %w", err)
-	}
-	if _, err := io.Copy(tmpFile, resp.Body); err != nil {
-		return fmt.Errorf("write body: %w", err)
-	}
-	if err := tmpFile.Close(); err != nil {
-		return fmt.Errorf("close temp: %w", err)
-	}
-
-	if err := os.Rename(tmpName, localPath); err != nil { //nolint:gosec // G703: localPath is sanitized by URLToLocalPath
-		return fmt.Errorf("rename: %w", err)
+	if err := store.Put(logicalPath, io.MultiReader(bytes.NewReader(first), resp.Body)); err != nil {
+		return fmt.Errorf("store: %w", err)
 	}
 
 	// Post-process HTML / CSS
 	if cfg.RewriteLinks {
 		ct := resp.Header.Get("Content-Type")
 		fileURL := snap.FileURL
 
-		if IsHTMLFile(localPath, ct, first) {
-			if err := ProcessHTML(localPath, fileURL, cfg, idx); err != nil && cfg.Debug {
-				log.Printf("html rewrite %s: %v", localPath, err)
+		if IsHTMLFile(logicalPath, ct, first) {
+			if err := ProcessHTML(store, logicalPath, fileURL, cfg, idx); err != nil && cfg.Debug {
+				log.Printf("html rewrite %s: %v", logicalPath, err)
 			}
-		} else if IsCSSResource(localPath, ct) {
-			if err := RewriteCSSFile(localPath, fileURL, cfg, idx); err != nil && cfg.Debug {
-				log.Printf("css rewrite %s: %v", localPath, err)
+		} else if IsCSSResource(logicalPath, ct) {
+			if err := RewriteCSSFile(store, logicalPath, fileURL, cfg, idx); err != nil && cfg.Debug {
+				log.Printf("css rewrite %s: %v", logicalPath, err)
 			}
 		}
 	}

internal/wayback/html.go

@@ -3,16 +3,15 @@ package wayback
 import (
 	"bytes"
 	"net/url"
-	"os"
 	"path/filepath"
 	"strings"
 
 	"golang.org/x/net/html"
 )
 
 // ProcessHTML rewrites links and canonical tags in an HTML file.
-func ProcessHTML(filePath, pageURL string, cfg *Config, idx *SnapshotIndex) error {
-	data, err := os.ReadFile(filePath) //nolint:gosec // G304: path is written by this program
+func ProcessHTML(store Storage, logicalPath, pageURL string, cfg *Config, idx *SnapshotIndex) error {
+	data, err := store.Get(logicalPath)
 	if err != nil {
 		return err
 	}
@@ -28,7 +27,7 @@ func ProcessHTML(filePath, pageURL string, cfg *Config, idx *SnapshotIndex) erro
 	}
 
 	// Relative directory of the output file (used for RelativeLink)
-	localDir := ToPosix(filepath.ToSlash(filepath.Dir(filePath)))
+	localDir := ToPosix(filepath.ToSlash(filepath.Dir(filepath.Join(cfg.Directory, filepath.FromSlash(logicalPath)))))
 
 	var walk func(*html.Node)
 	walk = func(n *html.Node) {
@@ -75,7 +74,7 @@ func ProcessHTML(filePath, pageURL string, cfg *Config, idx *SnapshotIndex) erro
 	if err := html.Render(&buf, doc); err != nil {
 		return err
 	}
-	return os.WriteFile(filePath, buf.Bytes(), 0600)
+	return store.PutBytes(logicalPath, buf.Bytes())
 }
 
 // attrName returns the relevant URL attribute for a given tag name.

internal/wayback/html_test.go

@@ -1,35 +1,29 @@
 package wayback
 
 import (
-	"os"
-	"path/filepath"
 	"strings"
 	"testing"
 )
 
-// processHTMLInTemp writes htmlContent to a temp file at <dir>/index.html,
-// sets cfg.Directory = dir so that relative-link math is consistent, then
-// runs ProcessHTML and returns the rewritten file contents.
+// processHTMLInTemp writes htmlContent into a LocalStorage backed by a temp
+// directory, runs ProcessHTML, and returns the rewritten file contents.
 func processHTMLInTemp(t *testing.T, htmlContent, pageURL string, cfg *Config) string {
 	t.Helper()
-	dir := t.TempDir()
-	cfg.Directory = dir
-
-	filePath := filepath.Join(dir, "index.html")
-	if err := os.WriteFile(filePath, []byte(htmlContent), 0600); err != nil {
+	store := NewLocalStorage(t.TempDir())
+	if err := store.PutBytes("test.html", []byte(htmlContent)); err != nil {
 		t.Fatalf("write test HTML: %v", err)
 	}
 
 	idx := NewSnapshotIndex()
-	if err := ProcessHTML(filePath, pageURL, cfg, idx); err != nil {
+	if err := ProcessHTML(store, "test.html", pageURL, cfg, idx); err != nil {
 		t.Fatalf("ProcessHTML: %v", err)
 	}
 
-	out, err := os.ReadFile(filePath)
+	got, err := store.Get("test.html")
 	if err != nil {
 		t.Fatalf("read result: %v", err)
 	}
-	return string(out)
+	return string(got)
 }
 
 func testHTMLCfg() *Config {

internal/wayback/storage.go

@@ -0,0 +1,85 @@
+package wayback
+
+import (
+	"io"
+	"os"
+	"path/filepath"
+)
+
+// Storage abstracts reading and writing downloaded snapshot files.
+// Logical paths are forward-slash relative paths as returned by URLToLocalPath
+// (e.g. "example.com/page/index.html"). Implementations map them to wherever
+// files actually live (OS directory, zip archive, memory map, …).
+type Storage interface {
+	// Exists reports whether the logical path already has content.
+	Exists(path string) bool
+	// Put writes the content of r to path. The write is atomic —
+	// no partial file is visible to concurrent readers.
+	Put(path string, r io.Reader) error
+	// Get returns the full content of path.
+	Get(path string) ([]byte, error)
+	// PutBytes writes data to path (convenience wrapper around Put).
+	PutBytes(path string, data []byte) error
+}
+
+// LocalStorage is the default Storage implementation that mirrors the
+// logical layout into a root directory on the OS filesystem.
+type LocalStorage struct {
+	rootDir string
+}
+
+// NewLocalStorage returns a LocalStorage rooted at dir.
+// The root directory is created lazily by Put/PutBytes.
+func NewLocalStorage(dir string) *LocalStorage {
+	return &LocalStorage{rootDir: dir}
+}
+
+// abs converts a logical forward-slash path to an absolute OS path.
+func (s *LocalStorage) abs(path string) string {
+	return filepath.Join(s.rootDir, filepath.FromSlash(path))
+}
+
+// Exists reports whether path already exists in storage.
+func (s *LocalStorage) Exists(path string) bool {
+	_, err := os.Stat(s.abs(path))
+	return err == nil
+}
+
+// Put streams r into path atomically via a temp file + rename.
+func (s *LocalStorage) Put(path string, r io.Reader) error {
+	fullPath := s.abs(path)
+	dir := filepath.Dir(fullPath)
+	if err := os.MkdirAll(dir, 0750); err != nil {
+		return err
+	}
+	tmpFile, err := os.CreateTemp(dir, ".wbdl-*")
+	if err != nil {
+		return err
+	}
+	tmpName := tmpFile.Name()
+	defer func() {
+		_ = tmpFile.Close()
+		_ = os.Remove(tmpName) // no-op if already renamed
+	}()
+	if _, err := io.Copy(tmpFile, r); err != nil {
+		return err
+	}
+	if err := tmpFile.Close(); err != nil {
+		return err
+	}
+	return os.Rename(tmpName, fullPath) //nolint:gosec // G703: fullPath is sanitized by URLToLocalPath
+}
+
+// Get returns the full content of path.
+func (s *LocalStorage) Get(path string) ([]byte, error) {
+	return os.ReadFile(s.abs(path)) //nolint:gosec // G304: path is written by this program
+}
+
+// PutBytes writes data to path, creating parent directories as needed.
+func (s *LocalStorage) PutBytes(path string, data []byte) error {
+	fullPath := s.abs(path)
+	if err := os.MkdirAll(filepath.Dir(fullPath), 0750); err != nil {
+		return err
+	}
+	return os.WriteFile(fullPath, data, 0600)
+}