goyomi/internal/httpclient/flare/flare.go

package flare

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"
	"sync"
	"time"

	"goyomi/internal/httpclient"
	"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 {
	fsClient   *httpclient.FlareSolverrClient
	rateLimit  float64
	burst      int
	referer    string

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

type Response struct {
	*http.Response
	Body io.ReadCloser
}

func (r *Response) Close() error {
	return r.Body.Close()
}

type Option func(*Client)

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

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

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

func NewClient(opts ...Option) *Client {
	c := &Client{
		limiters: make(map[string]*rate.Limiter),
	}
	for _, opt := range opts {
		opt(c)
	}
	if c.rateLimit == 0 {
		c.rateLimit = 1
	}
	if c.burst == 0 {
		c.burst = 1
	}
	fsClient, err := httpclient.NewFlareSolverrClient()
	if err == nil {
		c.fsClient = fsClient
	}
	return c
}

func (c *Client) SetFlareSolverrClient(fs *httpclient.FlareSolverrClient) {
	c.fsClient = fs
}

func (c *Client) doRequest(ctx context.Context, method string, rawURL string, body string) (*Response, error) {
	if c.fsClient == nil {
		return nil, fmt.Errorf("FlareSolverr client not configured")
	}
	c.mu.Lock()
	limiter, ok := c.limiters[rawURL]
	if !ok {
		limiter = rate.NewLimiter(rate.Limit(c.rateLimit), c.burst)
		c.limiters[rawURL] = limiter
	}
	c.mu.Unlock()
	if err := limiter.Wait(ctx); err != nil {
		return nil, err
	}

	var html string
	var cookies []*http.Cookie
	var err error

	if method == http.MethodGet {
		html, cookies, err = c.fsClient.Get(ctx, rawURL)
	} else {
		html, cookies, err = c.fsClient.Post(ctx, rawURL, body)
	}
	if err != nil {
		return nil, err
	}

	jar, _ := cookiejar.New(nil)
	if len(cookies) > 0 {
		httpURL, _ := http.NewRequest(http.MethodGet, rawURL, nil)
		jar.SetCookies(httpURL.URL, cookies)
	}

	fakeResp := &http.Response{
		StatusCode: http.StatusOK,
		Header:     http.Header{},
		Body:       io.NopCloser(strings.NewReader(html)),
		Request:    &http.Request{URL: &url.URL{Path: rawURL}},
	}
	return &Response{Response: fakeResp, Body: fakeResp.Body}, nil
}

func (c *Client) Get(ctx context.Context, rawURL string) (*Response, error) {
	return c.doRequest(ctx, http.MethodGet, rawURL, "")
}

func (c *Client) Post(ctx context.Context, rawURL string, bodyType string, body io.Reader) (*Response, error) {
	bodyStr, _ := io.ReadAll(body)
	return c.doRequest(ctx, http.MethodPost, rawURL, string(bodyStr))
}

func (c *Client) GetHTML(ctx context.Context, url string) (io.ReadCloser, error) {
	resp, err := c.Get(ctx, url)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return nil, fmt.Errorf("HTTP %d for %s", resp.StatusCode, url)
	}
	return resp.Body, nil
}

func (c *Client) PostHTML(ctx context.Context, url string, body string) (io.ReadCloser, error) {
	resp, err := c.Post(ctx, url, "application/x-www-form-urlencoded", strings.NewReader(body))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return nil, fmt.Errorf("HTTP %d for %s", resp.StatusCode, url)
	}
	return resp.Body, nil
}

func (c *Client) GetBytes(ctx context.Context, url string) ([]byte, error) {
	resp, err := c.Get(ctx, url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("HTTP %d for %s", resp.StatusCode, url)
	}
	return io.ReadAll(resp.Body)
}

func (c *Client) SetProxy(proxyURL string) error {
	return nil
}

func (c *Client) Do(req *http.Request) (*http.Response, error) {
	url := req.URL.String()
	var body string
	if req.Body != nil {
		b, _ := io.ReadAll(req.Body)
		body = string(b)
	}
	resp, err := c.doRequest(req.Context(), req.Method, url, body)
	if err != nil {
		return nil, err
	}
	return resp.Response, nil
}