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
}

// CachePolicy defines the caching behavior for a request.
type CachePolicy struct {
	ShouldUseCached bool                // whether to attempt fetching from the cache
	ShouldStore     bool                // whether to store the API response in the cache after retrieval
	CachedResponse  *SimpleHTTPResponse // the cached response if available and valid
}

// CachingResult encapsulates the response and whether it was retrieved from the cache.
type CachingResult struct {
	Response  *SimpleHTTPResponse
	FromCache bool
}

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

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

	// excludedCachePaths lists API endpoints for which responses 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)
}

// determineCachePolicy determines the caching policy for a given request.
// It returns a CachePolicy struct indicating whether to fetch from cache,
// whether to store the response in cache, and the cached response if available.
func determineCachePolicy(rawURL, userToken string, headers http.Header) CachePolicy {
	if !config.GlobalConfig.CacheEnabled {
		return CachePolicy{}
	}

	parsedURL, err := url.Parse(rawURL)
	if err != nil {
		return CachePolicy{}
	}

	urlPath := path.Clean(parsedURL.Path)

	// Check if the path is excluded from caching
	for _, exclPath := range excludedCachePaths {
		if strings.HasPrefix(urlPath, exclPath) {
			return CachePolicy{}
		}
	}

	// Retrieve the Cache-Control header from the downstream request
	// and check for "no-cache"
	cacheControl := headers.Get("Cache-Control")
	lowerCacheControl := strings.ToLower(cacheControl)
	if strings.Contains(lowerCacheControl, "no-cache") {
		return CachePolicy{}
	}

	cacheKey := generateCacheKey(rawURL, userToken)

	// Check if cache invalidation is requested
	redirectedHeader := headers.Get("X-Handled-Redirected")
	if strings.Contains(strings.ToLower(redirectedHeader), "true") {
		cache.Remove(cacheKey)
		return CachePolicy{}
	}

	// Attempt to fetch from cache
	if cachedItem, found := cache.Get(cacheKey); found {
		item := cachedItem.(CachedItem)
		if time.Now().Before(item.ExpiresAt) {
			return CachePolicy{
				ShouldUseCached: true,
				CachedResponse:  item.Response,
			}
		}
		// Cache expired
		cache.Remove(cacheKey)
	}

	// Determine if the API response should be cached based on headers
	shouldStore := cacheControl == "" || (!strings.Contains(lowerCacheControl, "no-store") && !strings.Contains(lowerCacheControl, "no-cache"))

	return CachePolicy{
		ShouldUseCached: shouldStore, // Use cached version if storage is allowed
		ShouldStore:     shouldStore,
	}
}

// manageCaching manages the caching logic during and after a request.
// It either returns a cached response or allows the request to proceed and optionally caches the new response.
func manageCaching(rawURL, userToken string, headers http.Header, response *SimpleHTTPResponse) CachingResult {
	policy := determineCachePolicy(rawURL, userToken, headers)
	if policy.ShouldUseCached && policy.CachedResponse != nil {
		return CachingResult{
			Response:  policy.CachedResponse,
			FromCache: true,
		}
	}

	if policy.ShouldStore && response != nil {
		ttl := config.GlobalConfig.CacheTTL
		parsedURL, err := url.Parse(rawURL)
		if err != nil {
			// Handle parsing error by returning the response without caching
			return CachingResult{
				Response:  response,
				FromCache: false,
			}
		}
		urlPath := path.Clean(parsedURL.Path)
		for _, shortPath := range shortTTLPaths {
			if strings.HasPrefix(urlPath, shortPath) {
				ttl = config.GlobalConfig.CacheShortTTL
				break
			}
		}

		cache.Add(generateCacheKey(rawURL, userToken), CachedItem{
			Response:  response,
			ExpiresAt: time.Now().Add(ttl),
		})
	}

	return CachingResult{
		Response:  response,
		FromCache: false,
	}
}

// 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, i18n.Errorf("failed to create API request with token: %w", url, err)
	}

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

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

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

	audit.LogAPIRoundTrip(audit.APIRequestSpan{
		StartTime:  start,
		EndTime:    end,
		RequestId:  request_context.GetFromContext(ctx).RequestId,
		Response:   resp,
		ErrorField: 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)

			// Handle caching logic after we receive an HTTP 200 response
			cachingResult := manageCaching(url, userToken, incomingHeaders, resp)
			if cachingResult.FromCache {
				return cachingResult.Response, nil
			}

			return cachingResult.Response, 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.
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)
	}

	// Determine caching policy before we send a request
	policy := determineCachePolicy(url, userToken, incomingHeaders)
	if policy.ShouldUseCached && policy.CachedResponse != nil {
		return policy.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
}