diff options
| -rw-r--r-- | internal/client/crawl.go | 39 |
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) } |