pcgamedb/utils/fetch.go

package utils

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"golang.org/x/net/html/charset"
)

const userAgent string = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36"

type FetchConfig struct {
	Method     string
	Url        string
	Data       interface{}
	RetryTimes int
	Headers    map[string]string
	Cookies    map[string]string
	Timeout    time.Duration
}

type FetchResponse struct {
	StatusCode int
	Data       []byte
	Header     http.Header
	Cookie     []*http.Cookie
}

func Fetch(cfg FetchConfig) (*FetchResponse, error) {
	var req *http.Request
	var resp *http.Response
	var backoff time.Duration = 1
	var reqBody io.Reader = nil
	var err error

	if cfg.RetryTimes == 0 {
		cfg.RetryTimes = 3
	}
	if cfg.Method == "" {
		cfg.Method = "GET"
	}
	if cfg.Timeout == 0 {
		cfg.Timeout = 10 * time.Second
	}

	if cfg.Data != nil && (cfg.Method == "POST" || cfg.Method == "PUT") {
		if cfg.Headers == nil {
			cfg.Headers = map[string]string{}
		}
		newHeaders := make(map[string]string)
		for k, v := range cfg.Headers {
			newHeaders[strings.ToLower(k)] = v
		}
		cfg.Headers = newHeaders
		if _, exist := cfg.Headers["content-type"]; !exist {
			cfg.Headers["content-type"] = "application/json"
		}
		v := cfg.Headers["content-type"]
		if v == "application/x-www-form-urlencoded" {
			switch data := cfg.Data.(type) {
			case map[string]string:
				params := url.Values{}
				for k, v := range data {
					params.Set(k, v)
				}
				reqBody = strings.NewReader(params.Encode())
			case string:
				reqBody = strings.NewReader(data)
			case url.Values:
				reqBody = strings.NewReader(data.Encode())
			default:
				return nil, errors.New("unsupported data type")
			}
		} else if v == "application/json" {
			switch data := cfg.Data.(type) {
			case []byte:
				reqBody = bytes.NewReader(data)
			case string:
				reqBody = strings.NewReader(data)
			case interface{}:
				jsonData, err := json.Marshal(cfg.Data)
				if err != nil {
					return nil, err
				}
				reqBody = bytes.NewReader(jsonData)
			default:
				return nil, errors.New("unsupported data type")
			}
		} else {
			reqBody = strings.NewReader(cfg.Data.(string))
		}
	}

	var bodyBuffer *bytes.Buffer
	if reqBody != nil {
		bodyBuffer = new(bytes.Buffer)
		_, err = io.Copy(bodyBuffer, reqBody)
		if err != nil {
			return nil, err
		}
	}

	for retryTime := 0; retryTime <= cfg.RetryTimes; retryTime++ {
		ctx, cancel := context.WithTimeout(context.Background(), cfg.Timeout)
		defer cancel()

		var currentReqBody io.Reader
		if bodyBuffer != nil {
			currentReqBody = bytes.NewReader(bodyBuffer.Bytes())
		}

		req, err = http.NewRequestWithContext(ctx, cfg.Method, cfg.Url, currentReqBody)
		if err != nil {
			return nil, err
		}
		if v, exist := cfg.Headers["user-agent"]; exist {
			if v != "" {
				req.Header.Set("user-agent", v)
			}
		} else {
			req.Header.Set("user-agent", userAgent)
		}
		for k, v := range cfg.Headers {
			req.Header.Set(k, v)
		}
		for k, v := range cfg.Cookies {
			req.AddCookie(&http.Cookie{Name: k, Value: v})
		}
		resp, err = http.DefaultClient.Do(req)
		if err != nil {
			if isRetryableError(err) {
				err = errors.New("request error: " + err.Error())
				time.Sleep(backoff * time.Second)
				backoff *= 2
				continue
			}
		}

		if resp == nil {
			return nil, errors.New("response is nil")
		}

		if isRetryableStatusCode(resp.StatusCode) {
			err = errors.New("response status code: " + resp.Status)
			time.Sleep(backoff * time.Second)
			backoff *= 2
			continue
		}

		contentType := resp.Header.Get("content-type")
		var reader io.Reader
		if strings.Contains(contentType, "charset=") {
			reader, err = charset.NewReader(resp.Body, contentType)
		} else {
			reader = resp.Body
		}
		if err != nil {
			return nil, err
		}
		dataBytes, err := io.ReadAll(reader)
		if err != nil {
			return nil, err
		}

		res := &FetchResponse{
			StatusCode: resp.StatusCode,
			Header:     resp.Header,
			Cookie:     resp.Cookies(),
			Data:       dataBytes,
		}

		return res, nil
	}
	return nil, err
}

func isRetryableStatusCode(statusCode int) bool {
	switch statusCode {
	case http.StatusInternalServerError,
		http.StatusBadGateway,
		http.StatusServiceUnavailable,
		http.StatusGatewayTimeout,
		http.StatusTooManyRequests:
		return true
	default:
		return false
	}
}

func isRetryableError(err error) bool {
	if err != nil {
		var netErr net.Error
		if errors.As(err, &netErr) && netErr.Timeout() {
			return true
		}
	}
	return false
}

func FetchWithWAFSession(cfg FetchConfig, session *WAFSession) (*FetchResponse, error) {
	if cfg.Cookies == nil {
		cfg.Cookies = map[string]string{}
	}
	for _, cookie := range session.Cookies {
		cfg.Cookies[cookie.Name] = cookie.Value
	}
	if cfg.Headers == nil {
		cfg.Headers = map[string]string{}
	}
	for k, v := range session.Headers {
		cfg.Headers[k] = v
	}
	return Fetch(cfg)
}