feat: add idle timeout handling for streaming crawl connections

1 files changed, 37 insertions, 2 deletions

diff --git a/internal/client/crawl.go b/internal/client/crawl.go
index 330bb88..74217f8 100644
--- a/internal/client/crawl.go
+++ b/internal/client/crawl.go
@@ -2,12 +2,27 @@ package client
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/url"
+	"time"
 )
 
+type idleTimeoutConn struct {
+	net.Conn
+	idleTimeout time.Duration
+}
+
+func (c *idleTimeoutConn) Read(b []byte) (int, error) {
+	if err := c.Conn.SetReadDeadline(time.Now().Add(c.idleTimeout)); err != nil {
+		return 0, err
+	}
+	return c.Conn.Read(b)
+}
+
 func (c *Client) CrawlDomain(domain string, subdomains bool, callback func(line string) error) error {
 	path := fmt.Sprintf("/crawl/%s", url.PathEscape(domain))
 	fullURL := c.config.BaseURL + path
@@ -18,14 +33,34 @@ func (c *Client) CrawlDomain(domain string, subdomains bool, callback func(line
 		fullURL += "?" + params.Encode()
 	}
 
-	req, err := http.NewRequest(http.MethodGet, fullURL, nil)
+	ctx := context.Background()
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil)
 	if err != nil {
 		return fmt.Errorf("failed to create request: %w", err)
 	}
 
 	req.Header.Set("User-Agent", c.config.UserAgent)
 
-	resp, err := c.httpClient.Do(req)
+	idleTimeout := 60 * time.Second
+	streamClient := &http.Client{
+		Transport: &http.Transport{
+			DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+				conn, err := (&net.Dialer{
+					Timeout:   30 * time.Second,
+					KeepAlive: 30 * time.Second,
+				}).DialContext(ctx, network, addr)
+				if err != nil {
+					return nil, err
+				}
+				return &idleTimeoutConn{
+					Conn:        conn,
+					idleTimeout: idleTimeout,
+				}, nil
+			},
+		},
+	}
+
+	resp, err := streamClient.Do(req)
 	if err != nil {
 		return fmt.Errorf("failed to execute request: %w", err)
 	}