Merge branch 'efixler:main' into main

Author: efixler

cmd/scrape-server/main.go

@@ -19,10 +19,12 @@ import (
 	"github.com/efixler/envflags"
 	"github.com/efixler/scrape/fetch"
 	"github.com/efixler/scrape/fetch/trafilatura"
+	"github.com/efixler/scrape/internal"
 	"github.com/efixler/scrape/internal/auth"
 	"github.com/efixler/scrape/internal/cmd"
 	"github.com/efixler/scrape/internal/headless"
 	"github.com/efixler/scrape/internal/server"
+	"github.com/efixler/scrape/internal/storage"
 	"github.com/efixler/scrape/resource"
 	"github.com/efixler/scrape/ua"
 	"github.com/efixler/webutil/graceful"
@@ -50,23 +52,26 @@ func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	dbh := dbFlags.MustDatabase()
 	dbFlags = nil
-	normalClient := fetch.MustClient(fetch.WithUserAgent(userAgent.Get().String()))
+	directClient := fetch.MustClient(fetch.WithUserAgent(userAgent.Get().String()))
 	var headlessFetcher fetch.URLFetcher = nil
 	if headlessEnabled.Get() {
 		headlessClient := headless.MustChromeClient(ctx, userAgent.Get().String(), 6)
 		headlessFetcher = trafilatura.MustNew(headlessClient)
 	}
 
-	// TODO: Implement options pattern for NewScrapeServer
-	ss, _ := server.NewScrapeServer(
+	ss := server.MustScrapeServer(
 		ctx,
-		dbh,
-		trafilatura.MustNew(normalClient),
-		headlessFetcher,
+		server.WithURLFetcher(
+			internal.NewStorageBackedFetcher(
+				trafilatura.MustNew(directClient),
+				storage.NewURLDataStore(dbh),
+			),
+		),
+		server.WithHeadlessIf(headlessFetcher),
+		server.WithAuthorizationIf(*signingKey.Get()),
 	)
 
-	if sk := *signingKey.Get(); len(sk) > 0 {
-		ss.SigningKey = sk
+	if ss.AuthEnabled() {
 		slog.Info("scrape-server authorization via JWT is enabled")
 	} else {
 		slog.Info("scrape-server authorization is disabled, running in open access mode")

cmd/scrape/main.go

@@ -185,13 +185,11 @@ func initFetcher(dbh *database.DBHandle) (*internal.StorageBackedFetcher, error)
 			fetch.WithUserAgent(userAgent.Get().String()),
 		)
 	}
-	fetcher, err := internal.NewStorageBackedFetcher(
+	fetcher := internal.NewStorageBackedFetcher(
 		trafilatura.MustNew(client),
 		storage.NewURLDataStore(dbh),
 	)
-	if err != nil {
-		return nil, fmt.Errorf("error creating storage backed fetcher: %s", err)
-	}
+
 	err = fetcher.Open(dbh.Ctx)
 	if err != nil {
 		return nil, fmt.Errorf("error opening storage backed fetcher: %s", err)

internal/benchmark_test.go

@@ -60,7 +60,7 @@ func makeFetcher(dbPath string, ctx context.Context) (*StorageBackedFetcher, err
 	topts := *trafilatura.DefaultOptions
 	topts.Transport = t
 
-	fetcher, err := NewStorageBackedFetcher(
+	fetcher := NewStorageBackedFetcher(
 		trafilatura.Factory(topts),
 		sqlite.Factory(dbPath),
 	)

internal/scrape.go

@@ -28,12 +28,12 @@ type StorageBackedFetcher struct {
 func NewStorageBackedFetcher(
 	fetcher fetch.URLFetcher,
 	storage store.URLDataStore,
-) (*StorageBackedFetcher, error) {
+) *StorageBackedFetcher {
 	return &StorageBackedFetcher{
 		Fetcher: fetcher,
 		Storage: storage,
 		saving:  new(sync.WaitGroup),
-	}, nil
+	}
 }
 
 // The context passed to Open() will be passed on to child components

internal/scrape_test.go

@@ -43,10 +43,8 @@ func TestFetchStoresAndRetrieves(t *testing.T) {
 	dbh := database.New(sqlite.MustNew(sqlite.InMemoryDB()))
 	storage := storage.NewURLDataStore(dbh)
 
-	fetcher, err := NewStorageBackedFetcher(tf, storage)
-	if err != nil {
-		t.Fatal(err)
-	}
+	fetcher := NewStorageBackedFetcher(tf, storage)
+
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	err = fetcher.Open(ctx)
@@ -168,10 +166,8 @@ func TestFetchUnstored(t *testing.T) {
 	}
 	dbh := database.New(sqlite.MustNew(sqlite.InMemoryDB()))
 	storage := storage.NewURLDataStore(dbh)
-	fetcher, err := NewStorageBackedFetcher(tf, storage)
-	if err != nil {
-		t.Fatal(err)
-	}
+	fetcher := NewStorageBackedFetcher(tf, storage)
+
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	err = fetcher.Open(ctx)

internal/server/home.go

@@ -0,0 +1,61 @@
+package server
+
+import (
+	"bytes"
+	"embed"
+	"html/template"
+	"log/slog"
+	"net/http"
+	"time"
+
+	"github.com/efixler/scrape/internal/auth"
+)
+
+//go:embed templates/index.html
+var home embed.FS
+
+func mustHomeTemplate(ss *scrapeServer) *template.Template {
+	tmpl := template.New("home")
+	var authTokenF = func() string { return "" }
+	var authEnabledF = func() bool { return ss.AuthEnabled() }
+	if authEnabledF() {
+		authTokenF = func() string {
+			c, err := auth.NewClaims(
+				auth.WithSubject("home"),
+				auth.ExpiresIn(60*time.Minute),
+			)
+			if err != nil {
+				slog.Error("Error creating claims for home view", "error", err)
+				return ""
+			}
+			s, err := c.Sign(ss.SigningKey())
+			if err != nil {
+				slog.Error("Error signing claims for home view", "error", err)
+				return ""
+			}
+			return s
+		}
+	}
+	funcMap := template.FuncMap{
+		"AuthToken":   authTokenF,
+		"AuthEnabled": authEnabledF,
+	}
+	tmpl = tmpl.Funcs(funcMap)
+	homeSource, _ := home.ReadFile("templates/index.html")
+	tmpl = template.Must(tmpl.Parse(string(homeSource)))
+	return tmpl
+}
+
+func homeHandler(ss *scrapeServer) http.HandlerFunc {
+	tmpl := mustHomeTemplate(ss)
+	return func(w http.ResponseWriter, r *http.Request) {
+		var buf bytes.Buffer
+		if err := tmpl.Execute(&buf, nil); err != nil {
+			http.Error(w, "Error rendering home page", http.StatusInternalServerError)
+			return
+		}
+		w.Header().Set("Content-Type", "text/html; charset=utf-8")
+		w.WriteHeader(http.StatusOK)
+		w.Write(buf.Bytes())
+	}
+}

internal/server/home_test.go

@@ -0,0 +1,65 @@
+package server
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	"github.com/efixler/scrape/internal/auth"
+)
+
+func TestMustTemplate(t *testing.T) {
+	tests := []struct {
+		name        string
+		key         auth.HMACBase64Key
+		expectToken bool
+	}{
+		{
+			name:        "with key",
+			key:         auth.MustNewHS256SigningKey(),
+			expectToken: true,
+		},
+		{
+			name:        "no key",
+			key:         nil,
+			expectToken: false,
+		},
+		{
+			name:        "empty key",
+			key:         auth.HMACBase64Key([]byte{}),
+			expectToken: false,
+		},
+	}
+	for _, test := range tests {
+		ss := MustScrapeServer(
+			context.Background(),
+			WithURLFetcher(&mockUrlFetcher{}),
+			WithAuthorizationIf(test.key),
+		)
+		tmpl := mustHomeTemplate(ss)
+		tmpl, err := tmpl.Parse("{{AuthToken}}")
+		if err != nil {
+			t.Fatalf("[%s] Error parsing template: %s", test.name, err)
+		}
+		var buf bytes.Buffer
+		err = tmpl.Execute(&buf, nil)
+		if err != nil {
+			t.Fatalf("[%s] Error executing template: %s", test.name, err)
+		}
+		output := buf.String()
+		if !test.expectToken && output != "" {
+			t.Fatalf("[%s] Expected empty output, got %s", test.name, output)
+		}
+		if test.expectToken {
+			switch output {
+			case "":
+				t.Fatalf("[%s] Expected non-empty token, got empty", test.name)
+			default:
+				_, err := auth.VerifyToken(test.key, output)
+				if err != nil {
+					t.Fatalf("[%s] Error verifying token: %s", test.name, err)
+				}
+			}
+		}
+	}
+}