PixivFE/server/proxy_checker/proxy_checker.go

// Package proxy_checker handles the verification of image proxy servers ("proxies").
//
// It tests each proxy to ensure they can correctly serve a specific image
// by comparing the content CRC32 of the response.
package proxy_checker

import (
	"context"
	"fmt"
	"hash/crc32"
	"io"
	"net/http"
	"strings"
	"sync"

	"codeberg.org/vnpower/pixivfe/v2/audit"
	"codeberg.org/vnpower/pixivfe/v2/config"
	"codeberg.org/vnpower/pixivfe/v2/server/utils"
)

const (
	testImagePath  = "/img-original/img/2024/01/21/20/50/51/115365120_p0.jpg"
	testImageCRC32 = "66c66408" // Using CRC32 for simplicity
)

var (
	workingProxies      []string
	workingProxiesMutex sync.RWMutex
)

// CheckProxies tests each proxy in the built-in proxy list for functionality.
//
// It uses goroutines to test proxies concurrently then updates the list of working proxies.
func CheckProxies(ctx context.Context, auditor *audit.Auditor) {
	auditor.SugaredLogger.Info("Starting image proxy check...")

	var wg sync.WaitGroup
	var mutex sync.Mutex
	var newWorkingProxies []string

	auditor.SugaredLogger.Infof("Total image proxies to check: %d", len(config.BuiltinProxyList))

	// Iterate over each proxy URL in the configured list.
	for _, proxy := range config.BuiltinProxyList {
		wg.Add(1)

		go func(proxyURL string) {
			defer wg.Done()

			// Test the proxy and receive the result and the HTTP response.
			isWorking, resp := testProxy(ctx, auditor, proxyURL)

			// Capture the HTTP status code for logging purposes.
			status := ""
			if resp != nil {
				status = resp.Status
			}

			if isWorking {
				// If the proxy works, add it to the list of new working proxies.
				mutex.Lock()
				newWorkingProxies = append(newWorkingProxies, proxyURL)
				mutex.Unlock()

				// Log success for the proxy.
				auditor.SugaredLogger.Infof("[OK]  %s %s", proxyURL, status)
			} else {
				// Log failure for the proxy.
				auditor.SugaredLogger.Warnf("[ERR] %s %s", proxyURL, status)
			}
		}(proxy) // Pass the current proxy URL to the goroutine.
	}

	wg.Wait()

	updateWorkingProxies(auditor, newWorkingProxies)
}

// testProxy checks if a given proxy is able to serve the test image correctly.
//
// It validates the response by comparing its CRC32 checksum with the expected value.
func testProxy(ctx context.Context, auditor *audit.Auditor, proxyBaseURL string) (bool, *http.Response) {
	fullURL := fmt.Sprintf("%s%s", strings.TrimRight(proxyBaseURL, "/"), testImagePath)
	auditor.SugaredLogger.Infof("Testing image proxy %s with full URL: %s", proxyBaseURL, fullURL)

	req, err := http.NewRequestWithContext(ctx, "GET", fullURL, nil)
	if err != nil {
		auditor.SugaredLogger.Errorf("Error creating request for image proxy %s: %v", proxyBaseURL, err)
		return false, nil
	}

	resp, err := utils.HttpClient.Do(req)
	if err != nil {
		auditor.SugaredLogger.Errorf("Error testing image proxy %s: %v", proxyBaseURL, err)
		return false, nil
	}

	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return false, resp
	}

	// Calculate CRC32 checksum of response body
	hash := crc32.NewIEEE()
	if _, err := io.Copy(hash, resp.Body); err != nil {
		auditor.SugaredLogger.Errorf("Error reading response body from image proxy %s: %v", proxyBaseURL, err)
		return false, resp
	}

	checksumHex := fmt.Sprintf("%08x", hash.Sum32())

	if checksumHex != testImageCRC32 {
		auditor.SugaredLogger.Warnf("CRC32 mismatch for image proxy %s. Expected: %s, Got: %s", proxyBaseURL, testImageCRC32, checksumHex)
	}

	return checksumHex == testImageCRC32, resp
}

// updateWorkingProxies safely updates the global list of working proxies.
func updateWorkingProxies(auditor *audit.Auditor, newProxies []string) {
	workingProxiesMutex.Lock()
	defer workingProxiesMutex.Unlock()

	workingProxies = newProxies
	auditor.SugaredLogger.Infof("Updated working image proxies. Count: %d", len(workingProxies))
}

// GetWorkingProxies returns a copy of the current list of working proxies.
//
// A read lock is used to allow concurrent reads while preventing data races.
func GetWorkingProxies() []string {
	workingProxiesMutex.RLock()
	defer workingProxiesMutex.RUnlock()

	return append([]string{}, workingProxies...)
}