goyomi/internal/httpclient/client.go

package httpclient

import (
	"context"
	"io"
	"net/http"
	"net/http/cookiejar"
	"strconv"
	"sync"
	"time"

	"golang.org/x/time/rate"
)

const defaultUserAgent = "Mozilla/5.0 (Linux; Android 10; K) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Mobile Safari/537.36"

type Client struct {
	http      *http.Client
	rateLimit float64
	burst     int
	referer   string

	mu       sync.Mutex
	limiters map[string]*rate.Limiter
}

type Option func(*Client)

func WithRateLimit(rps float64, burst int) Option {
	return func(c *Client) {
		c.rateLimit = rps
		c.burst = burst
	}
}

func WithTimeout(d time.Duration) Option {
	return func(c *Client) { c.http.Timeout = d }
}

func WithReferer(referer string) Option {
	return func(c *Client) { c.referer = referer }
}

func NewClient(opts ...Option) *Client {
	jar, _ := cookiejar.New(nil)
	c := &Client{
		http: &http.Client{
			Timeout: 30 * time.Second,
			Jar:     jar,
		},
		rateLimit: 1,
		burst:     1,
		limiters:  map[string]*rate.Limiter{},
	}
	for _, o := range opts {
		o(c)
	}
	return c
}

func (c *Client) limiter(host string) *rate.Limiter {
	c.mu.Lock()
	defer c.mu.Unlock()
	l, ok := c.limiters[host]
	if !ok {
		l = rate.NewLimiter(rate.Limit(c.rateLimit), c.burst)
		c.limiters[host] = l
	}
	return l
}

func (c *Client) Do(req *http.Request) (*http.Response, error) {
	if err := c.limiter(req.URL.Host).Wait(req.Context()); err != nil {
		return nil, err
	}
	if c.referer != "" && req.Header.Get("Referer") == "" {
		req.Header.Set("Referer", c.referer)
	}
	if req.Header.Get("User-Agent") == "" {
		req.Header.Set("User-Agent", defaultUserAgent)
	}

	const maxRetries = 3
	for attempt := 0; attempt <= maxRetries; attempt++ {
		resp, err := c.http.Do(req)
		if err != nil {
			return nil, err
		}
		if resp.StatusCode != http.StatusTooManyRequests {
			return resp, nil
		}
		resp.Body.Close()
		if attempt == maxRetries {
			return resp, nil
		}
		sleep := retryAfter(resp)
		select {
		case <-req.Context().Done():
			return nil, req.Context().Err()
		case <-time.After(sleep):
		}
	}
	panic("unreachable")
}

func retryAfter(resp *http.Response) time.Duration {
	ra := resp.Header.Get("Retry-After")
	if ra == "" {
		return 5 * time.Second
	}
	if secs, err := strconv.ParseFloat(ra, 64); err == nil {
		return time.Duration(secs * float64(time.Second))
	}
	if t, err := http.ParseTime(ra); err == nil {
		d := time.Until(t)
		if d > 0 {
			return d
		}
	}
	return 5 * time.Second
}

func (c *Client) Get(ctx context.Context, url string, headers map[string]string) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	for k, v := range headers {
		req.Header.Set(k, v)
	}
	return c.Do(req)
}

func (c *Client) Post(ctx context.Context, url string, body io.Reader, contentType string, headers map[string]string) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", contentType)
	for k, v := range headers {
		req.Header.Set(k, v)
	}
	return c.Do(req)
}

// HTTPClient returns the underlying *http.Client (for passing to graphql helper etc.)
func (c *Client) HTTPClient() *http.Client { return c.http }