PixivFE/core/requests.go

// This file implements the core functionality for handling API requests,
// including response caching, retry mechanisms, and request proxying.
package core

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"path"
	"strings"
	"time"

	"codeberg.org/vnpower/pixivfe/v2/config"
	"codeberg.org/vnpower/pixivfe/v2/i18n"
	"codeberg.org/vnpower/pixivfe/v2/server/audit"
	"codeberg.org/vnpower/pixivfe/v2/server/request_context"
	"codeberg.org/vnpower/pixivfe/v2/server/token_manager"
	"codeberg.org/vnpower/pixivfe/v2/server/utils"

	"github.com/hashicorp/go-retryablehttp"
	lru "github.com/hashicorp/golang-lru"
	"github.com/tidwall/gjson"
	"github.com/zeebo/xxh3"
)

// SimpleHTTPResponse represents a simplified HTTP response structure.
type SimpleHTTPResponse struct {
	StatusCode int
	Body       string
}

// CachedItem represents a cached API response along with its expiration time.
type CachedItem struct {
	Response  *SimpleHTTPResponse
	ExpiresAt time.Time
}

var (
	retryClient *retryablehttp.Client
	cacheSeed   uint64
	cache       *lru.Cache

	// shortTTLPaths lists URI paths that require a shorter TTL for their cached responses.
	shortTTLPaths = []string{
		"/ajax/discovery/artworks",
		"/ajax/discovery/novels",
		"/ajax/illust/new",
		"/ranking.php",
	}

	// excludedCachePaths lists API endpoints that should *never* be cached, regardless of any other factors.
	excludedCachePaths = []string{
		"/ranking.php",
	}
)

// InitCache initializes the API response cache based on parameters in GlobalConfig.
//
// It sets up an LRU cache with a specified size and logs the cache parameters.
// If caching is disabled in the configuration, it skips initialization.
func InitCache() {
	if !config.GlobalConfig.CacheEnabled {
		fmt.Println("Cache is disabled, skipping cache initialization")
		return
	}

	var err error

	// Initialize the LRU cache with the configured parameters.
	cache, err = lru.New(config.GlobalConfig.CacheSize)
	if err != nil {
		panic(fmt.Sprintf("Failed to create cache: %v", err))
	}
	fmt.Printf("API response cache size set to: %d items\n", config.GlobalConfig.CacheSize)
	fmt.Printf("API response cache TTL set to: %v\n", config.GlobalConfig.CacheTTL)
	fmt.Printf("API response cache short TTL set to: %v\n", config.GlobalConfig.CacheShortTTL)

	// Create a byte slice to hold the random seed.
	var seedBytes [8]byte

	// Read 8 random bytes from the crypto/rand reader.
	_, err = rand.Read(seedBytes[:])
	if err != nil {
		panic(fmt.Sprintf("Failed to generate cache key seed: %v", err))
	}

	// Convert the byte slice to a uint64 seed using little endian.
	cacheSeed = binary.LittleEndian.Uint64(seedBytes[:])
}

// init sets up a configured retryablehttp client.
func init() {
	retryClient = retryablehttp.NewClient()
	retryClient.RetryMax = config.GlobalConfig.APIMaxRetries
	retryClient.RetryWaitMin = config.GlobalConfig.APIBaseTimeout
	retryClient.RetryWaitMax = config.GlobalConfig.APIMaxBackoffTime
	retryClient.HTTPClient = utils.HttpClient
	retryClient.Logger = nil // Disables the default logger in go-retryablehttp
}

// generateCacheKey creates a unique identifier for caching purposes by hashing
// the combined request URL and the user's token using seeded xxhash.
//
// This ensures that cached responses are specific to both the endpoint and a given authenticated user.
func generateCacheKey(url, userToken string) string {
	combined := url + ":" + userToken

	hash := xxh3.HashStringSeed(combined, cacheSeed)

	return fmt.Sprintf("%x", hash)
}

// isShortTTLPath determines whether a given URL path should use a shorter TTL
// based on predefined paths that have more dynamic content.
func isShortTTLPath(rawURL string) bool {
	parsedURL, err := url.Parse(rawURL)
	if err != nil {
		return false
	}

	urlPath := path.Clean(parsedURL.Path)

	for _, shortPath := range shortTTLPaths {
		if urlPath == shortPath {
			return true
		}
	}

	return false
}

// isExcludedFromCache determines whether a specific API endpoint should *never* be cached, regardless of any other factors.
func isExcludedFromCache(rawURL string) bool {
	parsedURL, err := url.Parse(rawURL)
	if err != nil {
		return false
	}

	urlPath := path.Clean(parsedURL.Path)

	for _, exclPath := range excludedCachePaths {
		if strings.HasPrefix(urlPath, exclPath) {
			return true
		}
	}
	return false
}

// shouldInvalidateCache checks if the incoming request headers indicate that the cache should be invalidated.
func shouldInvalidateCache(incomingHeaders http.Header) bool {
	redirectedHeader := incomingHeaders.Get("X-Handled-Redirected")
	return strings.Contains(strings.ToLower(redirectedHeader), "true")
}

// fetchFromCache retrieves a cached response if it exists and is not expired.
func fetchFromCache(cacheKey string) (*CachedItem, bool) {
	if cachedItem, found := cache.Get(cacheKey); found {
		item := cachedItem.(CachedItem)
		if time.Now().Before(item.ExpiresAt) {
			return &item, true
		}
	}
	return nil, false
}

// shouldCacheResponse determines whether a particular response should be cached based on incoming headers.
func shouldCacheResponse(incomingHeaders http.Header) bool {
	cacheControl := incomingHeaders.Get("Cache-Control")
	return cacheControl == "" || !strings.Contains(strings.ToLower(cacheControl), "no-cache")
}

// determineTTL selects the appropriate TTL based on the URL path.
func determineTTL(url string) time.Duration {
	if isShortTTLPath(url) {
		return config.GlobalConfig.CacheShortTTL
	}
	return config.GlobalConfig.CacheTTL
}

// getCachedResponse attempts to retrieve a cached response if caching is enabled.
//
// It also handles cache invalidation based on redirection headers.
func getCachedResponse(url, userToken string, incomingHeaders http.Header) (*SimpleHTTPResponse, bool) {
	if !config.GlobalConfig.CacheEnabled {
		return nil, false
	}

	if isExcludedFromCache(url) {
		return nil, false
	}

	cacheKey := generateCacheKey(url, userToken)

	if shouldInvalidateCache(incomingHeaders) {
		cache.Remove(cacheKey)
	} else if cachedItem, found := fetchFromCache(cacheKey); found {
		return cachedItem.Response, true
	}

	return nil, false
}

// cacheSuccessfulResponse caches the response if caching is enabled and allowed by headers.
func cacheSuccessfulResponse(url, userToken string, incomingHeaders http.Header, resp *SimpleHTTPResponse) {
	if !config.GlobalConfig.CacheEnabled {
		return
	}

	if isExcludedFromCache(url) {
		return
	}

	if shouldCacheResponse(incomingHeaders) {
		ttl := determineTTL(url)
		cache.Add(generateCacheKey(url, userToken), CachedItem{
			Response:  resp,
			ExpiresAt: time.Now().Add(ttl),
		})
	}
}

// makeRequest is a helper function that encapsulates the common logic for making HTTP requests.
//
// It constructs the request using the provided request function, executes it using the retryable client,
// reads the response, logs the API round trip details for auditing, and returns a simplified response.
func makeRequest(ctx context.Context, reqFunc func(context.Context, string) (*retryablehttp.Request, error), token *token_manager.Token, url string) (*SimpleHTTPResponse, error) {
	// Generate the HTTP request using the provided request function.
	req, err := reqFunc(ctx, token.Value)
	if err != nil {
		return nil, err
	}

	start := time.Now()
	resp, err := retryClient.Do(req)
	end := time.Now()

	if err != nil {
		return nil, i18n.Errorf("failed to make request: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	audit.LogAPIRoundTrip(audit.APIRequestSpan{
		StartTime: start,
		EndTime:   end,
		RequestId: request_context.GetFromContext(ctx).RequestId,
		Response:  resp,
		Error:     err,
		Method:    req.Method,
		Url:       url,
		Token:     token.Value,
		Body:      string(body),
	})

	return &SimpleHTTPResponse{
		StatusCode: resp.StatusCode,
		Body:       string(body),
	}, nil
}

// executeSinglePostRequest handles POST requests without retries.
func executeSinglePostRequest(
	ctx context.Context,
	reqFunc func(context.Context, string) (*retryablehttp.Request, error),
	userToken string,
	url string,
) (*SimpleHTTPResponse, error) {
	token := &token_manager.Token{Value: userToken}
	return makeRequest(ctx, reqFunc, token, url)
}

// retrieveToken obtains a valid token, preferring the userToken if provided by the caller.
func retrieveToken(tokenManager *token_manager.TokenManager, userToken string) (*token_manager.Token, error) {
	if userToken != "" {
		return &token_manager.Token{Value: userToken}, nil
	}

	token := tokenManager.GetToken()
	if token == nil {
		tokenManager.ResetAllTokens()
		return nil,
			i18n.Errorf(
				`All tokens (%d) are timed out, resetting all tokens to their initial good state.
Consider providing additional tokens in PIXIVFE_TOKEN or reviewing API request level backoff configuration.
Please refer the following documentation for additional information:
- https://pixivfe-docs.pages.dev/hosting/obtaining-pixivfe-token/
- https://pixivfe-docs.pages.dev/hosting/environment-variables/#exponential-backoff-configuration`,
				len(config.GlobalConfig.Token),
			)
	}
	return token, nil
}

// executeGetWithRetries performs GET requests with retry logic and caching.
func executeGetWithRetries(
	ctx context.Context,
	reqFunc func(context.Context, string) (*retryablehttp.Request, error),
	userToken, url string,
	incomingHeaders http.Header,
) (*SimpleHTTPResponse, error) {
	var lastErr error
	tokenManager := config.GlobalConfig.TokenManager

	for i := 0; i < config.GlobalConfig.APIMaxRetries; i++ {
		token, err := retrieveToken(tokenManager, userToken)
		if err != nil {
			return nil, err
		}

		resp, err := makeRequest(ctx, reqFunc, token, url)
		if err != nil {
			lastErr = err
			tokenManager.MarkTokenStatus(token, token_manager.TimedOut)
			continue
		}

		if resp.StatusCode == http.StatusOK {
			tokenManager.MarkTokenStatus(token, token_manager.Good)
			cacheSuccessfulResponse(url, userToken, incomingHeaders, resp)
			return resp, nil
		}

		lastErr = i18n.Errorf("HTTP status code: %d", resp.StatusCode)
		tokenManager.MarkTokenStatus(token, token_manager.TimedOut)

		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
			// Continue with the next iteration
		}
	}

	return nil, i18n.Errorf("max retries reached for GET request. Last error: %v", lastErr)
}

// retryRequest handles the logic for executing HTTP requests with automatic retries and token management.
//
// It distinguishes between GET and POST requests, caching only GET requests when appropriate,
// and ensures that tokens are managed correctly for authentication.
func retryRequest(
	ctx context.Context,
	reqFunc func(context.Context, string) (*retryablehttp.Request, error),
	userToken string,
	isPost bool,
	url string,
	incomingHeaders http.Header,
) (*SimpleHTTPResponse, error) {
	if isPost {
		return executeSinglePostRequest(ctx, reqFunc, userToken, url)
	}

	if isExcludedFromCache(url) {
		return executeGetWithRetries(ctx, reqFunc, userToken, url, incomingHeaders)
	}

	if cachedResponse, shouldReturn := getCachedResponse(url, userToken, incomingHeaders); shouldReturn {
		return cachedResponse, nil
	}

	return executeGetWithRetries(ctx, reqFunc, userToken, url, incomingHeaders)
}

// API_GET performs a GET request to the Pixiv API with automatic retries and caching.
func API_GET(
	ctx context.Context,
	url string,
	userToken string,
	incomingHeaders http.Header,
) (*SimpleHTTPResponse, error) {
	return retryRequest(ctx, func(ctx context.Context, token string) (*retryablehttp.Request, error) {
		// Create a new GET request for the specified URL.
		req, err := retryablehttp.NewRequest("GET", url, nil)
		if err != nil {
			return nil, err
		}
		req = req.WithContext(ctx)

		req.Header.Add("User-Agent", config.GetRandomUserAgent())
		req.Header.Add("Accept-Language", config.GlobalConfig.AcceptLanguage)

		req.AddCookie(&http.Cookie{
			Name:  "PHPSESSID",
			Value: token,
		})
		return req, nil
	}, userToken, false, url, incomingHeaders)
}

// API_GET_UnwrapJson performs a GET request using API_GET and processes the JSON response.
//
// It validates the JSON structure, checks for errors within the response,
// and extracts the relevant body content for further use.
func API_GET_UnwrapJson(ctx context.Context, url, userToken string, incomingHeaders http.Header) (string, error) {
	resp, err := API_GET(ctx, url, userToken, incomingHeaders)
	if err != nil {
		return "", err
	}

	if !gjson.Valid(resp.Body) {
		return "", i18n.Errorf("Invalid JSON: %v", resp.Body)
	}

	result := gjson.Parse(resp.Body)

	if result.Get("error").Bool() {
		return "", errors.New(result.Get("message").String())
	}

	body := result.Get("body")

	if !body.Exists() {
		return "", i18n.Error("Incompatible response body")
	}

	return body.String(), nil
}

// createMultipartFormData constructs multipart form data from a map of fields.
//
// It is used to prepare data for POST requests that require multipart encoding.
func createMultipartFormData(fields map[string]string) (*bytes.Buffer, string, error) {
	body := new(bytes.Buffer)
	writer := multipart.NewWriter(body)

	defer writer.Close()

	for key, value := range fields {
		if err := writer.WriteField(key, value); err != nil {
			return nil, "", err
		}
	}

	return body, writer.FormDataContentType(), nil
}

// API_POST performs a POST request to the Pixiv API with support for different payload types.
//
// It handles authentication and constructs the appropriate request body based on the payload.
func API_POST(
	ctx context.Context,
	url string,
	payload interface{},
	userToken, csrf string,
	contentType string,
	incomingHeaders http.Header,
) (*SimpleHTTPResponse, error) {
	if userToken == "" {
		return nil, i18n.Error("userToken is required for POST requests")
	}

	resp, err := retryRequest(ctx, func(ctx context.Context, token string) (*retryablehttp.Request, error) {
		var req *retryablehttp.Request
		var err error

		// Determine the type of payload and construct the request accordingly.
		switch v := payload.(type) {
		case string:
			// If the payload is a string, send it as a raw byte buffer.
			req, err = retryablehttp.NewRequest("POST", url, bytes.NewBuffer([]byte(v)))
		case map[string]string:
			// If the payload is a map, construct multipart form data.
			body, formContentType, err := createMultipartFormData(v)
			if err != nil {
				return nil, err
			}
			req, err = retryablehttp.NewRequest("POST", url, body)
			if err == nil {
				// Update the content type to reflect multipart form data.
				contentType = formContentType
			}
		default:
			// Return an error if the payload type is unsupported.
			return nil, i18n.Error("Unsupported payload type")
		}

		if err != nil {
			return nil, err
		}

		req = req.WithContext(ctx)

		req.Header.Add("User-Agent", config.GetRandomUserAgent())
		req.Header.Add("Accept", "application/json")
		req.Header.Add("x-csrf-token", csrf)

		req.AddCookie(&http.Cookie{
			Name:  "PHPSESSID",
			Value: token,
		})
		req.Header.Add("Content-Type", contentType)

		return req, nil
	}, userToken, true, url, incomingHeaders)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("API request failed with status code: %d, body: %s", resp.StatusCode, resp.Body)
	}

	return resp, nil
}

// ProxyRequest forwards an incoming HTTP request to a target server and writes the response back to the original client.
//
// It handles copying of response headers, status codes, and the response body.
func ProxyRequest(w http.ResponseWriter, req *http.Request) error {
	resp, err := utils.HttpClient.Do(req)
	if err != nil {
		return i18n.Errorf("failed to proxy request: %w", err)
	}
	defer resp.Body.Close()

	header := w.Header()
	for k, v := range resp.Header {
		header[k] = v
	}

	w.WriteHeader(resp.StatusCode)

	_, err = io.Copy(w, resp.Body)
	if err != nil {
		return i18n.Errorf("failed to copy response body: %w", err)
	}

	return nil
}