Configurable Timeouts for HTTP Requests

cmd/all.go

@@ -11,6 +11,7 @@ import (
 	"github.com/Snider/Borg/pkg/compress"
 	"github.com/Snider/Borg/pkg/datanode"
 	"github.com/Snider/Borg/pkg/github"
+	"github.com/Snider/Borg/pkg/httpclient"
 	"github.com/Snider/Borg/pkg/tim"
 	"github.com/Snider/Borg/pkg/trix"
 	"github.com/Snider/Borg/pkg/ui"
@@ -42,7 +43,15 @@ func NewAllCmd() *cobra.Command {
 				return err
 			}
 
-			repos, err := GithubClient.GetPublicRepos(cmd.Context(), owner)
+			totalTimeout, _ := cmd.Flags().GetDuration("timeout")
+			connectTimeout, _ := cmd.Flags().GetDuration("connect-timeout")
+			tlsTimeout, _ := cmd.Flags().GetDuration("tls-timeout")
+			headerTimeout, _ := cmd.Flags().GetDuration("header-timeout")
+
+			httpClient := httpclient.NewClient(totalTimeout, connectTimeout, tlsTimeout, headerTimeout)
+			githubClient := github.NewGithubClient(httpClient)
+
+			repos, err := githubClient.GetPublicRepos(cmd.Context(), owner)
 			if err != nil {
 				return err
 			}

cmd/all_test.go

@@ -23,7 +23,7 @@ func TestAllCmd_Good(t *testing.T) {
 		},
 	})
 	oldNewAuthenticatedClient := github.NewAuthenticatedClient
-	github.NewAuthenticatedClient = func(ctx context.Context) *http.Client {
+	github.NewAuthenticatedClient = func(ctx context.Context, baseClient *http.Client) *http.Client {
 		return mockGithubClient
 	}
 	defer func() {
@@ -67,7 +67,7 @@ func TestAllCmd_Bad(t *testing.T) {
 		},
 	})
 	oldNewAuthenticatedClient := github.NewAuthenticatedClient
-	github.NewAuthenticatedClient = func(ctx context.Context) *http.Client {
+	github.NewAuthenticatedClient = func(ctx context.Context, baseClient *http.Client) *http.Client {
 		return mockGithubClient
 	}
 	defer func() {
@@ -96,7 +96,7 @@ func TestAllCmd_Ugly(t *testing.T) {
 			},
 		})
 		oldNewAuthenticatedClient := github.NewAuthenticatedClient
-		github.NewAuthenticatedClient = func(ctx context.Context) *http.Client {
+		github.NewAuthenticatedClient = func(ctx context.Context, baseClient *http.Client) *http.Client {
 			return mockGithubClient
 		}
 		defer func() {

cmd/collect.go

@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"time"
+
 	"github.com/spf13/cobra"
 )
 
@@ -11,11 +13,18 @@ func init() {
 	RootCmd.AddCommand(GetCollectCmd())
 }
 func NewCollectCmd() *cobra.Command {
-	return &cobra.Command{
+	cmd := &cobra.Command{
 		Use:   "collect",
 		Short: "Collect a resource from a URI.",
 		Long:  `Collect a resource from a URI and store it in a DataNode.`,
 	}
+
+	cmd.PersistentFlags().Duration("timeout", 0, "Total request timeout (e.g., 60s). 0 means no timeout.")
+	cmd.PersistentFlags().Duration("connect-timeout", 10*time.Second, "TCP connection establishment timeout (e.g., 10s)")
+	cmd.PersistentFlags().Duration("tls-timeout", 10*time.Second, "TLS handshake timeout (e.g., 10s)")
+	cmd.PersistentFlags().Duration("header-timeout", 30*time.Second, "Time to receive response headers timeout (e.g., 30s)")
+
+	return cmd
 }
 
 func GetCollectCmd() *cobra.Command {

cmd/collect_github_repo.go

@@ -6,6 +6,7 @@ import (
 	"os"
 
 	"github.com/Snider/Borg/pkg/compress"
+	"github.com/Snider/Borg/pkg/httpclient"
 	"github.com/Snider/Borg/pkg/tim"
 	"github.com/Snider/Borg/pkg/trix"
 	"github.com/Snider/Borg/pkg/ui"
@@ -44,6 +45,13 @@ func NewCollectGithubRepoCmd() *cobra.Command {
 				return fmt.Errorf("invalid compression: %s (must be 'none', 'gz', or 'xz')", compression)
 			}
 
+			totalTimeout, _ := cmd.Flags().GetDuration("timeout")
+			connectTimeout, _ := cmd.Flags().GetDuration("connect-timeout")
+			tlsTimeout, _ := cmd.Flags().GetDuration("tls-timeout")
+			headerTimeout, _ := cmd.Flags().GetDuration("header-timeout")
+
+			_ = httpclient.NewClient(totalTimeout, connectTimeout, tlsTimeout, headerTimeout)
+
 			prompter := ui.NewNonInteractivePrompter(ui.GetVCSQuote)
 			prompter.Start()
 			defer prompter.Stop()

cmd/collect_github_repos.go

@@ -4,20 +4,24 @@ import (
 	"fmt"
 
 	"github.com/Snider/Borg/pkg/github"
+	"github.com/Snider/Borg/pkg/httpclient"
 	"github.com/spf13/cobra"
 )
 
-var (
-	// GithubClient is the github client used by the command. It can be replaced for testing.
-	GithubClient = github.NewGithubClient()
-)
-
 var collectGithubReposCmd = &cobra.Command{
 	Use:   "repos [user-or-org]",
 	Short: "Collects all public repositories for a user or organization",
 	Args:  cobra.ExactArgs(1),
 	RunE: func(cmd *cobra.Command, args []string) error {
-		repos, err := GithubClient.GetPublicRepos(cmd.Context(), args[0])
+		totalTimeout, _ := cmd.Flags().GetDuration("timeout")
+		connectTimeout, _ := cmd.Flags().GetDuration("connect-timeout")
+		tlsTimeout, _ := cmd.Flags().GetDuration("tls-timeout")
+		headerTimeout, _ := cmd.Flags().GetDuration("header-timeout")
+
+		httpClient := httpclient.NewClient(totalTimeout, connectTimeout, tlsTimeout, headerTimeout)
+		githubClient := github.NewGithubClient(httpClient)
+
+		repos, err := githubClient.GetPublicRepos(cmd.Context(), args[0])
 		if err != nil {
 			return err
 		}

cmd/collect_pwa.go

@@ -5,6 +5,7 @@ import (
 	"os"
 
 	"github.com/Snider/Borg/pkg/compress"
+	"github.com/Snider/Borg/pkg/httpclient"
 	"github.com/Snider/Borg/pkg/pwa"
 	"github.com/Snider/Borg/pkg/tim"
 	"github.com/Snider/Borg/pkg/trix"
@@ -13,17 +14,9 @@ import (
 	"github.com/spf13/cobra"
 )
 
-type CollectPWACmd struct {
-	cobra.Command
-	PWAClient pwa.PWAClient
-}
-
 // NewCollectPWACmd creates a new collect pwa command
-func NewCollectPWACmd() *CollectPWACmd {
-	c := &CollectPWACmd{
-		PWAClient: pwa.NewPWAClient(),
-	}
-	c.Command = cobra.Command{
+func NewCollectPWACmd() *cobra.Command {
+	cmd := &cobra.Command{
 		Use:   "pwa [url]",
 		Short: "Collect a single PWA using a URI",
 		Long: `Collect a single PWA and store it in a DataNode.
@@ -44,24 +37,32 @@ Examples:
 			compression, _ := cmd.Flags().GetString("compression")
 			password, _ := cmd.Flags().GetString("password")
 
-			finalPath, err := CollectPWA(c.PWAClient, pwaURL, outputFile, format, compression, password)
+			totalTimeout, _ := cmd.Flags().GetDuration("timeout")
+			connectTimeout, _ := cmd.Flags().GetDuration("connect-timeout")
+			tlsTimeout, _ := cmd.Flags().GetDuration("tls-timeout")
+			headerTimeout, _ := cmd.Flags().GetDuration("header-timeout")
+
+			httpClient := httpclient.NewClient(totalTimeout, connectTimeout, tlsTimeout, headerTimeout)
+			pwaClient := pwa.NewPWAClient(httpClient)
+
+			finalPath, err := CollectPWA(pwaClient, pwaURL, outputFile, format, compression, password)
 			if err != nil {
 				return err
 			}
 			fmt.Fprintln(cmd.OutOrStdout(), "PWA saved to", finalPath)
 			return nil
 		},
 	}
-	c.Flags().String("uri", "", "The URI of the PWA to collect (can also be passed as positional arg)")
-	c.Flags().String("output", "", "Output file for the DataNode")
-	c.Flags().String("format", "datanode", "Output format (datanode, tim, trix, or stim)")
-	c.Flags().String("compression", "none", "Compression format (none, gz, or xz)")
-	c.Flags().String("password", "", "Password for encryption (required for stim format)")
-	return c
+	cmd.Flags().String("uri", "", "The URI of the PWA to collect (can also be passed as positional arg)")
+	cmd.Flags().String("output", "", "Output file for the DataNode")
+	cmd.Flags().String("format", "datanode", "Output format (datanode, tim, trix, or stim)")
+	cmd.Flags().String("compression", "none", "Compression format (none, gz, or xz)")
+	cmd.Flags().String("password", "", "Password for encryption (required for stim format)")
+	return cmd
 }
 
 func init() {
-	collectCmd.AddCommand(&NewCollectPWACmd().Command)
+	collectCmd.AddCommand(NewCollectPWACmd())
 }
 func CollectPWA(client pwa.PWAClient, pwaURL string, outputFile string, format string, compression string, password string) (string, error) {
 	if pwaURL == "" {

cmd/collect_website.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/schollz/progressbar/v3"
 	"github.com/Snider/Borg/pkg/compress"
+	"github.com/Snider/Borg/pkg/httpclient"
 	"github.com/Snider/Borg/pkg/tim"
 	"github.com/Snider/Borg/pkg/trix"
 	"github.com/Snider/Borg/pkg/ui"
@@ -51,7 +52,14 @@ func NewCollectWebsiteCmd() *cobra.Command {
 				bar = ui.NewProgressBar(-1, "Crawling website")
 			}
 
-			dn, err := website.DownloadAndPackageWebsite(websiteURL, depth, bar)
+			totalTimeout, _ := cmd.Flags().GetDuration("timeout")
+			connectTimeout, _ := cmd.Flags().GetDuration("connect-timeout")
+			tlsTimeout, _ := cmd.Flags().GetDuration("tls-timeout")
+			headerTimeout, _ := cmd.Flags().GetDuration("header-timeout")
+
+			client := httpclient.NewClient(totalTimeout, connectTimeout, tlsTimeout, headerTimeout)
+
+			dn, err := website.DownloadAndPackageWebsite(websiteURL, depth, bar, client)
 			if err != nil {
 				return fmt.Errorf("error downloading and packaging website: %w", err)
 			}

cmd/collect_website_test.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+	"net/http"
 	"path/filepath"
 	"strings"
 	"testing"
@@ -14,7 +15,7 @@ import (
 func TestCollectWebsiteCmd_Good(t *testing.T) {
 	// Mock the website downloader
 	oldDownloadAndPackageWebsite := website.DownloadAndPackageWebsite
-	website.DownloadAndPackageWebsite = func(startURL string, maxDepth int, bar *progressbar.ProgressBar) (*datanode.DataNode, error) {
+	website.DownloadAndPackageWebsite = func(startURL string, maxDepth int, bar *progressbar.ProgressBar, client *http.Client) (*datanode.DataNode, error) {
 		return datanode.New(), nil
 	}
 	defer func() {
@@ -35,7 +36,7 @@ func TestCollectWebsiteCmd_Good(t *testing.T) {
 func TestCollectWebsiteCmd_Bad(t *testing.T) {
 	// Mock the website downloader to return an error
 	oldDownloadAndPackageWebsite := website.DownloadAndPackageWebsite
-	website.DownloadAndPackageWebsite = func(startURL string, maxDepth int, bar *progressbar.ProgressBar) (*datanode.DataNode, error) {
+	website.DownloadAndPackageWebsite = func(startURL string, maxDepth int, bar *progressbar.ProgressBar, client *http.Client) (*datanode.DataNode, error) {
 		return nil, fmt.Errorf("website error")
 	}
 	defer func() {

examples/all/main.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"log"
+	"net/http"
 	"os"
 
 	"github.com/Snider/Borg/pkg/github"
@@ -13,7 +14,7 @@ import (
 func main() {
 	log.Println("Collecting all repositories for a user...")
 
-	repos, err := github.NewGithubClient().GetPublicRepos(context.Background(), "Snider")
+	repos, err := github.NewGithubClient(http.DefaultClient).GetPublicRepos(context.Background(), "Snider")
 	if err != nil {
 		log.Fatalf("Failed to get public repos: %v", err)
 	}

examples/collect_pwa/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"log"
+	"net/http"
 	"os"
 
 	"github.com/Snider/Borg/pkg/pwa"
@@ -10,7 +11,7 @@ import (
 func main() {
 	log.Println("Collecting PWA...")
 
-	client := pwa.NewPWAClient()
+	client := pwa.NewPWAClient(http.DefaultClient)
 	pwaURL := "https://squoosh.app"
 
 	manifestURL, err := client.FindManifest(pwaURL)

examples/collect_website/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"log"
+	"net/http"
 	"os"
 
 	"github.com/Snider/Borg/pkg/website"
@@ -11,7 +12,7 @@ func main() {
 	log.Println("Collecting website...")
 
 	// Download and package the website.
-	dn, err := website.DownloadAndPackageWebsite("https://example.com", 2, nil)
+	dn, err := website.DownloadAndPackageWebsite("https://example.com", 2, nil, http.DefaultClient)
 	if err != nil {
 		log.Fatalf("Failed to collect website: %v", err)
 	}

pkg/github/github.go

@@ -21,30 +21,37 @@ type GithubClient interface {
 }
 
 // NewGithubClient creates a new GithubClient.
-func NewGithubClient() GithubClient {
-	return &githubClient{}
+func NewGithubClient(client *http.Client) GithubClient {
+	return &githubClient{client: client}
 }
 
-type githubClient struct{}
+type githubClient struct {
+	client *http.Client
+}
 
 // NewAuthenticatedClient creates a new authenticated http client.
-var NewAuthenticatedClient = func(ctx context.Context) *http.Client {
+var NewAuthenticatedClient = func(ctx context.Context, baseClient *http.Client) *http.Client {
 	token := os.Getenv("GITHUB_TOKEN")
 	if token == "" {
-		return http.DefaultClient
+		return baseClient
 	}
 	ts := oauth2.StaticTokenSource(
 		&oauth2.Token{AccessToken: token},
 	)
-	return oauth2.NewClient(ctx, ts)
+	authedClient := oauth2.NewClient(ctx, ts)
+	authedClient.Transport = &oauth2.Transport{
+		Base:   baseClient.Transport,
+		Source: ts,
+	}
+	return authedClient
 }
 
 func (g *githubClient) GetPublicRepos(ctx context.Context, userOrOrg string) ([]string, error) {
 	return g.getPublicReposWithAPIURL(ctx, "https://api.github.com", userOrOrg)
 }
 
 func (g *githubClient) getPublicReposWithAPIURL(ctx context.Context, apiURL, userOrOrg string) ([]string, error) {
-	client := NewAuthenticatedClient(ctx)
+	client := NewAuthenticatedClient(ctx, g.client)
 	var allCloneURLs []string
 	url := fmt.Sprintf("%s/users/%s/repos", apiURL, userOrOrg)
 	isFirstRequest := true

pkg/github/github_test.go

@@ -154,7 +154,7 @@ func TestFindNextURL_Ugly(t *testing.T) {
 
 func TestNewAuthenticatedClient_Good(t *testing.T) {
 	t.Setenv("GITHUB_TOKEN", "test-token")
-	client := NewAuthenticatedClient(context.Background())
+	client := NewAuthenticatedClient(context.Background(), http.DefaultClient)
 	if client == http.DefaultClient {
 		t.Error("expected an authenticated client, but got http.DefaultClient")
 	}
@@ -163,17 +163,17 @@ func TestNewAuthenticatedClient_Good(t *testing.T) {
 func TestNewAuthenticatedClient_Bad(t *testing.T) {
 	// Unset the variable to ensure it's not present
 	t.Setenv("GITHUB_TOKEN", "")
-	client := NewAuthenticatedClient(context.Background())
+	client := NewAuthenticatedClient(context.Background(), http.DefaultClient)
 	if client != http.DefaultClient {
 		t.Error("expected http.DefaultClient when no token is set, but got something else")
 	}
 }
 
 // setupMockClient is a helper function to inject a mock http.Client.
 func setupMockClient(t *testing.T, mock *http.Client) *githubClient {
-	client := &githubClient{}
+	client := &githubClient{client: mock}
 	originalNewAuthenticatedClient := NewAuthenticatedClient
-	NewAuthenticatedClient = func(ctx context.Context) *http.Client {
+	NewAuthenticatedClient = func(ctx context.Context, baseClient *http.Client) *http.Client {
 		return mock
 	}
 	// Restore the original function after the test