feat(app): add internal components batch 7

This commit is contained in:
吴璨
2026-09-30 10:20:58 +08:00
parent baa1154cdd
commit aad2388ba4
10 changed files with 1626 additions and 0 deletions
+125
View File
@@ -0,0 +1,125 @@
package app
import (
"encoding/json"
"strconv"
"strings"
"time"
)
type Drama struct {
VIP *bool `json:"vip,omitempty"`
ID string `json:"id"`
Source string `json:"source,omitempty"`
SourceID string `json:"sourceId,omitempty"`
Title string `json:"title"`
Name string `json:"name"`
Desc string `json:"desc"`
Intro string `json:"intro"`
Cover any `json:"cover"`
CoverURL any `json:"coverUrl"`
CoverURLSnake any `json:"cover_url"`
Image any `json:"image"`
ImageURL any `json:"imageUrl"`
ImageURLSnake any `json:"image_url"`
Img any `json:"img"`
Pic any `json:"pic"`
Picture any `json:"picture"`
Poster any `json:"poster"`
Thumb any `json:"thumb"`
Thumbnail any `json:"thumbnail"`
TotalEpisode any `json:"totalEpisode"`
TotalEpisodeSnake any `json:"total_episode"`
ChapterCount any `json:"chapterCount"`
ChapterCountSnake any `json:"chapter_count"`
EpisodeCount any `json:"episodeCount"`
EpisodeCountSnake any `json:"episode_count"`
Total any `json:"total"`
Episodes any `json:"episodes"`
Duration any `json:"duration"`
ChannelName string `json:"channelName"`
Category string `json:"category"`
CategoryName string `json:"categoryName"`
CategoryNameSnake string `json:"category_name"`
TypeName string `json:"typeName"`
TypeNameSnake string `json:"type_name"`
SortName string `json:"sortName"`
SortNameSnake string `json:"sort_name"`
Remark string `json:"remark,omitempty"`
Score string `json:"score,omitempty"`
Views string `json:"views,omitempty"`
Heat string `json:"heat,omitempty"`
OnlineDate string `json:"onlineDate,omitempty"`
Tags []string `json:"tags,omitempty"`
ReleaseStatus string `json:"releaseStatus,omitempty"`
SortMetadata *sortMetadataState `json:"sortMetadata,omitempty"`
}
func (d Drama) DisplayTitle() string {
if strings.TrimSpace(d.Title) != "" {
return d.Title
}
if strings.TrimSpace(d.Name) != "" {
return d.Name
}
return "短剧"
}
type Chapter struct {
VIP bool `json:"vip,omitempty"`
ID string `json:"id"`
Source string `json:"source,omitempty"`
Title string `json:"title"`
VideoURL string `json:"videoUrl"`
CurrentEpisode json.RawMessage `json:"currentEpisode"`
MediaSize int64 `json:"mediaSize"`
PageURL string `json:"pageUrl,omitempty"`
Referer string `json:"referer,omitempty"`
}
func (c Chapter) EpisodeString(fallback int) string {
if len(c.CurrentEpisode) > 0 && string(c.CurrentEpisode) != "null" {
var s string
if err := json.Unmarshal(c.CurrentEpisode, &s); err == nil && strings.TrimSpace(s) != "" {
return strings.TrimSpace(s)
}
var n int
if err := json.Unmarshal(c.CurrentEpisode, &n); err == nil && n > 0 {
return strconv.Itoa(n)
}
var f float64
if err := json.Unmarshal(c.CurrentEpisode, &f); err == nil && f > 0 {
return strconv.Itoa(int(f))
}
}
return strconv.Itoa(fallback)
}
type Task struct {
DramaID string `json:"dramaId"`
DramaTitle string `json:"dramaTitle"`
Chapter Chapter `json:"chapter"`
Index int `json:"index"`
Total int `json:"total"`
OutPath string `json:"outPath"`
ReleaseStatus string `json:"releaseStatus,omitempty"`
DownloadQuality int `json:"downloadQuality,omitempty"`
}
type Result struct {
Task Task `json:"task"`
OK bool `json:"ok"`
Err string `json:"err,omitempty"`
}
type DownloadProgress struct {
Percent int
DownloadedBytes int64
TotalBytes int64
SpeedBytesPerSecond float64
Elapsed time.Duration
Remaining time.Duration
MediaElapsed time.Duration
MediaTotal time.Duration
Phase string
}
+180
View File
@@ -0,0 +1,180 @@
package app
import (
"context"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"sync"
"time"
)
type networkCheckResult struct {
Name string `json:"name"`
Status int `json:"status,omitempty"`
Detail string `json:"detail,omitempty"`
Error string `json:"error,omitempty"`
}
func (a *UIApp) networkCheckResults(ctx context.Context) []networkCheckResult {
targets := []struct{ name, endpoint string }{
{"红果入口", a.downloader.providerBaseURL(sourceHongguo) + "/"},
}
var mediaTask *Task
a.mu.Lock()
concurrency := a.cfg.RequestConcurrency
for _, drama := range a.dramas {
if dramaProvider(drama) == sourceHongguo {
if endpoint, valid := buildImageURL(bestDramaCover(drama)); valid {
targets = append(targets, struct{ name, endpoint string }{"红果封面 CDN", endpoint})
break
}
}
}
for index := len(a.taskOrder) - 1; index >= 0; index-- {
task := a.tasks[a.taskOrder[index]]
if task == nil || isChapterPlaceholderTask(task) {
continue
}
source := sourceFromDramaID(task.DramaID)
if source == sourceHongguo {
copy := task.Source
mediaTask = &copy
if task.Status == uiStatusFailed {
break
}
}
}
a.mu.Unlock()
if concurrency < 1 {
concurrency = 1
}
if concurrency > 2 {
concurrency = 2
}
results := make([]networkCheckResult, len(targets)+1)
jobs := make(chan int)
var workers sync.WaitGroup
for worker := 0; worker < concurrency; worker++ {
workers.Add(1)
go func() {
defer workers.Done()
for index := range jobs {
if index == len(targets) {
results[index] = a.checkHongguoMedia(ctx, mediaTask)
continue
}
target := targets[index]
probeCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
status, err := a.checkNetworkResource(probeCtx, http.MethodGet, target.endpoint, imageReferer(target.endpoint), false)
cancel()
results[index] = networkCheckResult{Name: target.name, Status: status, Error: a.redactError(err)}
}
}()
}
for index := range results {
jobs <- index
}
close(jobs)
workers.Wait()
return results
}
func (a *UIApp) checkNetworkResource(ctx context.Context, method, endpoint, referer string, key bool) (int, error) {
upstream, err := http.NewRequestWithContext(ctx, method, endpoint, nil)
if err != nil {
return 0, err
}
upstream.Header.Set("User-Agent", userAgent)
upstream.Header.Set("Referer", referer)
response, err := a.downloader.client.Do(upstream)
if err != nil {
return 0, publicError(err)
}
defer response.Body.Close()
if response.StatusCode < 200 || response.StatusCode >= 300 {
return response.StatusCode, fmt.Errorf("%s HTTP %d", upstream.URL.Hostname(), response.StatusCode)
}
if key {
body, err := io.ReadAll(io.LimitReader(response.Body, 17))
if err != nil {
return response.StatusCode, err
}
if len(body) != 16 {
return response.StatusCode, fmt.Errorf("密钥响应不是有效的 16 字节 AES-128 密钥")
}
} else {
if _, err := io.Copy(io.Discard, io.LimitReader(response.Body, 32768)); err != nil {
return response.StatusCode, err
}
}
return response.StatusCode, nil
}
func (a *UIApp) checkHongguoMedia(ctx context.Context, task *Task) networkCheckResult {
result := networkCheckResult{Name: "红果实际媒体"}
if task == nil {
result.Detail = "暂无红果分集任务,入口可达不代表媒体可下载"
return result
}
ctx, cancel := context.WithTimeout(ctx, 40*time.Second)
defer cancel()
media, err := a.downloader.resolveProviderMedia(ctx, *task)
if err != nil {
result.Error = "播放地址解析失败: " + a.redactError(err)
return result
}
for depth := 0; depth < 3; depth++ {
variant := selectBestM3U8Variant(media.Playlist, media.URL)
if variant == "" {
break
}
media.Playlist, media.URL, err = a.downloader.fetchMediaPlaylist(ctx, variant, media.Referer)
if err != nil {
result.Error = "播放列表失败: " + a.redactError(err)
return result
}
}
base, _ := url.Parse(media.URL)
segment := media.URL
keyChecked := false
for _, line := range strings.Split(media.Playlist, "\n") {
line = strings.TrimSpace(line)
if strings.HasPrefix(line, "#EXT-X-KEY:") && !strings.Contains(line, "METHOD=NONE") {
match := hlsURIAttribute.FindStringSubmatch(line)
if len(match) == 2 && !keyChecked {
reference, parseErr := url.Parse(match[1])
if parseErr == nil {
result.Status, err = a.checkNetworkResource(ctx, http.MethodGet, base.ResolveReference(reference).String(), media.Referer, true)
} else {
err = parseErr
}
if err != nil {
result.Error = "视频密钥失败: " + a.redactError(err)
return result
}
keyChecked = true
}
} else if line != "" && !strings.HasPrefix(line, "#") {
reference, parseErr := url.Parse(line)
if parseErr != nil {
result.Error = a.redactError(parseErr)
return result
}
segment = base.ResolveReference(reference).String()
break
}
}
result.Status, err = a.checkNetworkResource(ctx, http.MethodHead, segment, media.Referer, false)
if err != nil {
result.Error = "媒体分片检测失败: " + a.redactError(err)
return result
}
result.Detail = "播放地址、首个媒体资源可达(不等于整集下载验证)"
if keyChecked {
result.Detail = "播放列表、16 字节密钥、首个分片均可达(不等于整集下载验证)"
}
return result
}
+150
View File
@@ -0,0 +1,150 @@
package app
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/url"
"sync"
"time"
)
const dnsAlternateSubnet = "8.8.8.0/24"
type dnsCacheEntry struct {
addresses []string
expires time.Time
}
type dnsResolver struct {
client *http.Client
mu sync.Mutex
cache map[string]dnsCacheEntry
inFlight map[string]chan struct{}
endpoints []string
}
func newDNSResolver(transport *http.Transport) *dnsResolver {
lookupTransport := transport.Clone()
lookupTransport.TLSClientConfig = nil
return &dnsResolver{
client: &http.Client{Transport: lookupTransport, Timeout: 4 * time.Second},
cache: map[string]dnsCacheEntry{}, inFlight: map[string]chan struct{}{},
endpoints: []string{"https://dns.alidns.com/resolve", "https://dns.google/resolve"},
}
}
func (resolver *dnsResolver) lookup(ctx context.Context, host string) ([]string, error) {
entry, err := resolver.lookupEntry(ctx, host, "")
return entry.addresses, err
}
func (resolver *dnsResolver) lookupEntry(ctx context.Context, host, subnet string) (dnsCacheEntry, error) {
cacheKey := host + "|" + subnet
for {
resolver.mu.Lock()
if entry, found := resolver.cache[cacheKey]; found && time.Now().Before(entry.expires) {
resolver.mu.Unlock()
return entry, nil
}
if pending := resolver.inFlight[cacheKey]; pending != nil {
resolver.mu.Unlock()
select {
case <-pending:
continue
case <-ctx.Done():
return dnsCacheEntry{}, ctx.Err()
}
}
resolver.inFlight[cacheKey] = make(chan struct{})
resolver.mu.Unlock()
break
}
entry, err := resolver.query(ctx, host, subnet)
resolver.mu.Lock()
if err == nil {
resolver.cache[cacheKey] = entry
}
close(resolver.inFlight[cacheKey])
delete(resolver.inFlight, cacheKey)
resolver.mu.Unlock()
return entry, err
}
func (resolver *dnsResolver) query(ctx context.Context, host, subnet string) (dnsCacheEntry, error) {
lastErr := errors.New("DoH 未返回公网 IPv4 地址")
for _, endpoint := range resolver.endpoints {
if err := ctx.Err(); err != nil {
return dnsCacheEntry{}, err
}
lookupURL, err := url.Parse(endpoint)
if err != nil {
return dnsCacheEntry{}, err
}
query := lookupURL.Query()
query.Set("name", host)
query.Set("type", "A")
if subnet != "" {
query.Set("edns_client_subnet", subnet)
}
lookupURL.RawQuery = query.Encode()
request, err := http.NewRequestWithContext(ctx, http.MethodGet, lookupURL.String(), nil)
if err != nil {
return dnsCacheEntry{}, err
}
request.Header.Set("Accept", "application/dns-json")
response, err := resolver.client.Do(request)
if err != nil {
lastErr = publicError(err)
continue
}
var payload struct {
Status int `json:"Status"`
Answer []struct {
Type int `json:"type"`
TTL int `json:"TTL"`
Data string `json:"data"`
} `json:"Answer"`
}
err = json.NewDecoder(io.LimitReader(response.Body, 65536)).Decode(&payload)
response.Body.Close()
if response.StatusCode != http.StatusOK || err != nil || payload.Status != 0 {
lastErr = errors.New("DoH 服务没有返回有效 DNS 记录")
continue
}
entry := dnsCacheEntry{}
ttl := 300
var addressErr error
for _, answer := range payload.Answer {
if answer.Type != 1 {
continue
}
address := net.ParseIP(answer.Data)
if address == nil || address.To4() == nil {
addressErr = errors.New("DoH 返回无效 IPv4 地址")
break
}
entry.addresses = append(entry.addresses, address.String())
if answer.TTL < ttl {
ttl = answer.TTL
}
}
if addressErr == nil {
addressErr = validateImageAddresses(entry.addresses)
}
if addressErr != nil {
lastErr = fmt.Errorf("DoH 未返回可用公网 IPv4 地址:%w", addressErr)
continue
}
if ttl < 1 {
ttl = 1
}
entry.expires = time.Now().Add(time.Duration(ttl) * time.Second)
return entry, nil
}
return dnsCacheEntry{}, lastErr
}
+117
View File
@@ -0,0 +1,117 @@
package app
import (
"context"
"encoding/json"
"errors"
"io"
"net/http"
"net/http/httptest"
"reflect"
"sync/atomic"
"testing"
"time"
)
func writeDNSFixture(writer http.ResponseWriter, addresses ...string) {
answers := make([]map[string]any, 0, len(addresses))
for _, address := range addresses {
answers = append(answers, map[string]any{"type": 1, "TTL": 37, "data": address})
}
writer.Header().Set("Content-Type", "application/dns-json")
json.NewEncoder(writer).Encode(map[string]any{"Status": 0, "Answer": answers})
}
func TestDNSResolverRetriesUnsafeOrInvalidAnswers(t *testing.T) {
for _, test := range []struct {
name string
addresses []string
body string
status int
}{
{name: "fake-ip", addresses: []string{"198.18.0.42"}},
{name: "private", addresses: []string{"172.18.0.2"}},
{name: "loopback", addresses: []string{"127.0.0.1"}},
{name: "metadata", addresses: []string{"169.254.169.254"}},
{name: "carrier-network", addresses: []string{"100.64.0.1"}},
{name: "reserved", addresses: []string{"203.0.113.1"}},
{name: "mixed-private", addresses: []string{"93.184.216.35", "10.0.0.1"}},
{name: "mixed-fake-ip", addresses: []string{"93.184.216.35", "198.19.0.1"}},
{name: "invalid-address", addresses: []string{"93.184.216.35", "invalid address"}},
{name: "ipv6-in-a-record", addresses: []string{"2606:4700:4700::1111"}},
{name: "empty-answer"},
{name: "nxdomain", body: `{"Status":3}`},
{name: "malformed-json", body: `{"Status":`},
{name: "http-error", addresses: []string{"93.184.216.35"}, status: http.StatusServiceUnavailable},
} {
t.Run(test.name, func(t *testing.T) {
var firstCalls, secondCalls atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
query := request.URL.Query()
if query.Get("name") != "covers.example.org" || query.Get("type") != "A" || query.Get("edns_client_subnet") != dnsAlternateSubnet || request.Header.Get("Accept") != "application/dns-json" {
t.Error("unexpected DNS request", request.URL)
}
if request.URL.Path == "/first" {
firstCalls.Add(1)
if test.status != 0 {
writer.WriteHeader(test.status)
}
if test.body != "" {
io.WriteString(writer, test.body)
} else {
writeDNSFixture(writer, test.addresses...)
}
return
}
secondCalls.Add(1)
writeDNSFixture(writer, "93.184.216.34")
}))
defer server.Close()
standard := http.DefaultTransport.(*http.Transport).Clone()
standard.Proxy = nil
resolver := newDNSResolver(standard)
defer resolver.client.CloseIdleConnections()
resolver.endpoints = []string{server.URL + "/first", server.URL + "/second"}
for range 2 {
entry, err := resolver.lookupEntry(context.Background(), "covers.example.org", dnsAlternateSubnet)
if err != nil || !reflect.DeepEqual(entry.addresses, []string{"93.184.216.34"}) {
t.Fatal("DNS failed to continue to a valid public answer", entry.addresses, err)
}
if remaining := time.Until(entry.expires); remaining <= 0 || remaining > 37*time.Second {
t.Fatal("DNS cache ignored the answer TTL", remaining)
}
}
if firstCalls.Load() != 1 || secondCalls.Load() != 1 {
t.Fatal("DNS fallback did not reuse its successful cache", firstCalls.Load(), secondCalls.Load())
}
})
}
}
func TestDNSResolverDoesNotCacheUnsafeAnswers(t *testing.T) {
var calls atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
calls.Add(1)
writeDNSFixture(writer, "93.184.216.34", "198.18.0.42")
}))
defer server.Close()
standard := http.DefaultTransport.(*http.Transport).Clone()
standard.Proxy = nil
resolver := newDNSResolver(standard)
defer resolver.client.CloseIdleConnections()
resolver.endpoints = []string{server.URL + "/first", server.URL + "/second"}
for range 2 {
addresses, err := resolver.lookup(context.Background(), "covers.example.org")
if err == nil || len(addresses) != 0 {
t.Fatal("unsafe DNS answer was returned", addresses, err)
}
}
if calls.Load() != 4 {
t.Fatal("unsafe DNS answer was cached or prevented trying another service", calls.Load())
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
if _, err := resolver.lookup(ctx, "cancelled.example.org"); !errors.Is(err, context.Canceled) || calls.Load() != 4 {
t.Fatal("cancelled DNS lookup did not stop", calls.Load(), err)
}
}
+264
View File
@@ -0,0 +1,264 @@
package app
import (
"context"
"crypto/tls"
"errors"
"fmt"
"net"
"net/http"
"net/netip"
"net/url"
"strings"
"sync"
"time"
)
type coverRequestKey struct{}
type coverSourceKey struct{}
var nonPublicImageNetworks = []netip.Prefix{
netip.MustParsePrefix("0.0.0.0/8"), netip.MustParsePrefix("100.64.0.0/10"),
netip.MustParsePrefix("192.0.0.0/24"), netip.MustParsePrefix("192.0.2.0/24"),
netip.MustParsePrefix("198.18.0.0/15"), netip.MustParsePrefix("198.51.100.0/24"),
netip.MustParsePrefix("203.0.113.0/24"), netip.MustParsePrefix("240.0.0.0/4"),
netip.MustParsePrefix("64:ff9b::/96"), netip.MustParsePrefix("64:ff9b:1::/48"),
netip.MustParsePrefix("100::/64"), netip.MustParsePrefix("2001::/32"),
netip.MustParsePrefix("2001:db8::/32"), netip.MustParsePrefix("2002::/16"),
}
func publicImageAddress(value string) bool {
address, err := netip.ParseAddr(value)
if err != nil || address.Zone() != "" {
return false
}
address = address.Unmap()
if !address.IsGlobalUnicast() || address.IsPrivate() {
return false
}
for _, network := range nonPublicImageNetworks {
if network.Contains(address) {
return false
}
}
return true
}
func validImageURL(remote *url.URL) bool {
if remote == nil || remote.Scheme != "https" || remote.Opaque != "" || remote.User != nil ||
(remote.Port() != "" && remote.Port() != "443") {
return false
}
host := strings.ToLower(strings.TrimSuffix(remote.Hostname(), "."))
if _, err := netip.ParseAddr(host); err == nil {
return publicImageAddress(host)
}
if len(host) > 253 || !strings.Contains(host, ".") || strings.ContainsAny(host, ":%\\") {
return false
}
for _, suffix := range []string{"localhost", "local", "internal", "lan", "home", "invalid", "test", "example"} {
if host == suffix || strings.HasSuffix(host, "."+suffix) {
return false
}
}
for _, label := range strings.Split(host, ".") {
if label == "" || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' {
return false
}
for _, ch := range label {
if ch != '-' && (ch < 'a' || ch > 'z') && (ch < '0' || ch > '9') {
return false
}
}
}
return true
}
type imageTransport struct {
base http.RoundTripper
standard *http.Transport
resolver *dnsResolver
lookup func(context.Context, string) ([]string, error)
fallback func(context.Context, string) ([]string, error)
mu sync.Mutex
transports map[string]*http.Transport
}
func newImageTransport(base http.RoundTripper, standard *http.Transport) *imageTransport {
resolver := newDNSResolver(standard)
return &imageTransport{
base: base, standard: standard, resolver: resolver,
lookup: net.DefaultResolver.LookupHost, transports: map[string]*http.Transport{},
fallback: func(ctx context.Context, host string) ([]string, error) {
entry, err := resolver.lookupEntry(ctx, host, dnsAlternateSubnet)
return entry.addresses, err
},
}
}
func (transport *imageTransport) RoundTrip(request *http.Request) (*http.Response, error) {
if request.Context().Value(coverRequestKey{}) != true {
return transport.base.RoundTrip(request)
}
if !validImageURL(request.URL) || request.Method != http.MethodGet {
return nil, errors.New("封面地址不受支持")
}
host := strings.TrimSuffix(request.URL.Hostname(), ".")
fallbackUsed := false
fallback := func() ([]string, error) {
fallbackUsed = true
return transport.fallback(request.Context(), host)
}
addresses, err := transport.lookup(request.Context(), host)
if (err != nil || len(addresses) == 0) && transport.fallback != nil && request.Context().Err() == nil {
addresses, err = fallback()
}
if err != nil {
return nil, fmt.Errorf("封面域名解析失败: %w", err)
}
var proxyURL *url.URL
if transport.standard.Proxy != nil {
proxyURL, err = transport.standard.Proxy(request)
if err != nil {
return nil, err
}
}
tried := map[string]bool{}
lastErr := errors.New("封面没有可用的公网地址")
for attempt := 0; attempt < 2; attempt++ {
if addressErr := validateImageAddresses(addresses); addressErr != nil {
source, _ := request.Context().Value(coverSourceKey{}).(string)
if fallbackUsed || transport.fallback == nil || accountSourceGroup(source) == "" || net.ParseIP(host) != nil || request.Context().Err() != nil {
return nil, addressErr
}
addresses, err = fallback()
if err != nil {
return nil, fmt.Errorf("%v;可信 DNS 重解析失败:%w", addressErr, err)
}
if err = validateImageAddresses(addresses); err != nil {
return nil, fmt.Errorf("可信 DNS 未提供可用封面地址:%w", err)
}
}
for _, address := range addresses {
if tried[address] || len(tried) >= 4 {
continue
}
tried[address] = true
pinned := transport.pinned(host, address, proxyURL)
upstream := request.Clone(request.Context())
upstream.Host = request.URL.Host
upstream.URL.Host = net.JoinHostPort(address, "443")
response, roundTripErr := pinned.RoundTrip(upstream)
if response != nil {
response.Request = request
}
if roundTripErr == nil {
return response, nil
}
lastErr = roundTripErr
if request.Context().Err() != nil {
return nil, request.Context().Err()
}
}
if fallbackUsed || transport.fallback == nil || len(tried) >= 4 {
break
}
addresses, err = fallback()
if err != nil {
break
}
}
return nil, fmt.Errorf("封面 CDN 连接失败: %w", lastErr)
}
func validateImageAddresses(addresses []string) error {
if len(addresses) == 0 {
return errors.New("封面没有可用的公网地址")
}
for _, address := range addresses {
if !publicImageAddress(address) {
return fmt.Errorf("封面地址指向本机、内网或保留网络(%s)", address)
}
}
return nil
}
func (transport *imageTransport) pinned(host, address string, proxyURL *url.URL) *http.Transport {
key := host + "|" + address
if proxyURL != nil {
key += "|" + proxyURL.String()
}
transport.mu.Lock()
defer transport.mu.Unlock()
if existing := transport.transports[key]; existing != nil {
return existing
}
pinned := transport.standard.Clone()
if pinned.TLSClientConfig == nil {
pinned.TLSClientConfig = &tls.Config{}
}
pinned.TLSClientConfig.ServerName = host
pinned.Proxy = http.ProxyURL(proxyURL)
if proxyURL != nil && proxyURL.Scheme == "https" {
plain := *proxyURL
plain.Scheme = "http"
if plain.Port() == "" {
plain.Host = net.JoinHostPort(plain.Hostname(), "443")
}
pinned.Proxy = http.ProxyURL(&plain)
pinned.DialContext = transport.tlsProxyDialer(proxyURL.Hostname())
}
if len(transport.transports) >= 32 {
for key, stale := range transport.transports {
stale.CloseIdleConnections()
delete(transport.transports, key)
break
}
}
transport.transports[key] = pinned
return pinned
}
func (transport *imageTransport) tlsProxyDialer(host string) func(context.Context, string, string) (net.Conn, error) {
return func(ctx context.Context, network, address string) (net.Conn, error) {
dial := transport.standard.DialContext
if dial == nil {
dialer := &net.Dialer{Timeout: 10 * time.Second}
dial = dialer.DialContext
}
connection, err := dial(ctx, network, address)
if err != nil {
return nil, err
}
config := &tls.Config{}
if transport.standard.TLSClientConfig != nil {
config = transport.standard.TLSClientConfig.Clone()
}
config.ServerName, config.NextProtos = host, []string{"http/1.1"}
secured := tls.Client(connection, config)
if transport.standard.TLSHandshakeTimeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, transport.standard.TLSHandshakeTimeout)
defer cancel()
}
if err := secured.HandshakeContext(ctx); err != nil {
connection.Close()
return nil, err
}
return secured, nil
}
}
func (transport *imageTransport) CloseIdleConnections() {
transport.resolver.client.CloseIdleConnections()
if closer, ok := transport.base.(interface{ CloseIdleConnections() }); ok {
closer.CloseIdleConnections()
}
transport.mu.Lock()
defer transport.mu.Unlock()
for key, pinned := range transport.transports {
pinned.CloseIdleConnections()
delete(transport.transports, key)
}
}
+179
View File
@@ -0,0 +1,179 @@
package app
import (
"context"
"errors"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
)
func TestDownloaderCoverDNSFallbackUsesConfiguredProxy(t *testing.T) {
var calls atomic.Int32
proxy := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
calls.Add(1)
if request.Method != http.MethodConnect || request.Host != "dns.alidns.com:443" && request.Host != "dns.google:443" {
t.Error("unexpected DNS proxy destination", request.Method, request.Host)
}
http.Error(writer, "local DNS failure fixture", http.StatusBadGateway)
}))
defer proxy.Close()
cfg := defaultConfig()
cfg.dataDir, cfg.OutputDir, cfg.ProxyURL = t.TempDir(), t.TempDir(), proxy.URL
downloader := NewDownloader(cfg)
defer downloader.client.CloseIdleConnections()
transport := downloader.client.Transport.(*imageTransport)
transport.lookup = func(context.Context, string) ([]string, error) {
return []string{"198.18.0.42"}, nil
}
app := &UIApp{downloader: downloader}
ctx := context.WithValue(context.Background(), coverSourceKey{}, sourceHongguo)
_, err := app.loadCoverImage(ctx, "https://covers.example.org/fixture", nil)
if err == nil || !strings.Contains(err.Error(), "可信 DNS 重解析失败") || calls.Load() != 2 {
t.Fatal("registered cover did not use the default DNS fallback through its configured proxy", calls.Load(), err)
}
}
func TestDownloaderCoverDNSFallbackRecoversPublicOrigin(t *testing.T) {
for _, test := range []struct {
name string
addresses []string
err error
dnsCalls int32
}{
{name: "public-dns", addresses: []string{"93.184.216.34"}},
{name: "mihomo-fake-ip", addresses: []string{"198.18.0.42"}, dnsCalls: 1},
{name: "docker-private-ip", addresses: []string{"172.18.0.2"}, dnsCalls: 1},
{name: "private-ipv6", addresses: []string{"fc00::1"}, dnsCalls: 1},
{name: "mapped-private-ipv4", addresses: []string{"::ffff:192.168.1.1"}, dnsCalls: 1},
{name: "mixed-dns", addresses: []string{"93.184.216.35", "198.18.0.42"}, dnsCalls: 1},
{name: "empty-dns", dnsCalls: 1},
{name: "dns-error", err: errors.New("local DNS unavailable"), dnsCalls: 1},
} {
t.Run(test.name, func(t *testing.T) {
var dnsCalls, targetCalls, dials atomic.Int32
target := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
targetCalls.Add(1)
if request.Host != "example.com" || request.TLS.ServerName != "example.com" || request.URL.RawQuery != "signature=%2F%2b%3D" || request.Header.Get("Referer") != "https://source.example.org/" {
t.Error("DNS recovery changed the original Host, SNI, signature or source header")
}
io.WriteString(writer, "local text transport fixture")
}))
defer target.Close()
dns := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
dnsCalls.Add(1)
if request.URL.Query().Get("name") != "example.com" {
t.Error("unexpected DNS hostname", request.URL)
}
writeDNSFixture(writer, "93.184.216.34")
}))
defer dns.Close()
cfg := defaultConfig()
cfg.dataDir, cfg.OutputDir, cfg.ProxyURL = t.TempDir(), t.TempDir(), "direct"
downloader := NewDownloader(cfg)
defer downloader.client.CloseIdleConnections()
transport := downloader.client.Transport.(*imageTransport)
transport.lookup = func(context.Context, string) ([]string, error) { return test.addresses, test.err }
transport.resolver.endpoints = []string{dns.URL}
transport.standard.TLSClientConfig = target.Client().Transport.(*http.Transport).TLSClientConfig.Clone()
transport.standard.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
dials.Add(1)
if address != "93.184.216.34:443" {
t.Error("attempted a connection outside the verified public destination", address)
return nil, errors.New("unexpected connection")
}
return (&net.Dialer{}).DialContext(ctx, network, target.Listener.Addr().String())
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
ctx = context.WithValue(context.WithValue(ctx, coverRequestKey{}, true), coverSourceKey{}, sourceHongguo)
for range 2 {
request, _ := http.NewRequestWithContext(ctx, http.MethodGet, "https://example.com/fixture?signature=%2F%2b%3D", nil)
request.Header.Set("Referer", "https://source.example.org/")
response, err := downloader.client.Do(request)
if err != nil {
t.Fatal("default cover DNS recovery failed", err)
}
body, err := io.ReadAll(response.Body)
response.Body.Close()
if err != nil || string(body) != "local text transport fixture" || response.Request.URL.String() != request.URL.String() {
t.Fatal("DNS recovery changed the response or original URL", err)
}
}
if dnsCalls.Load() != test.dnsCalls || targetCalls.Load() != 2 || dials.Load() != 1 {
t.Fatal("unexpected DNS or connection reuse", dnsCalls.Load(), targetCalls.Load(), dials.Load())
}
})
}
}
func TestDownloaderCoverDNSFallbackRejectsUnsafeDestinations(t *testing.T) {
for _, test := range []struct {
name string
url string
source string
addresses []string
redirect bool
dnsCalls int32
}{
{name: "literal-private", url: "https://10.0.0.1/fixture", source: sourceHongguo},
{name: "literal-fake-ip", url: "https://198.18.0.42/fixture", source: sourceHongguo},
{name: "unregistered-source", url: "https://example.com/fixture"},
{name: "private-answer", url: "https://example.com/fixture", source: sourceHongguo, addresses: []string{"127.0.0.1"}, dnsCalls: 2},
{name: "mixed-answer", url: "https://example.com/fixture", source: sourceHongguo, addresses: []string{"93.184.216.34", "10.0.0.1"}, dnsCalls: 2},
{name: "private-redirect", url: "https://example.com/fixture", source: sourceHongguo, addresses: []string{"169.254.169.254"}, redirect: true, dnsCalls: 3},
} {
t.Run(test.name, func(t *testing.T) {
var dnsCalls, targetCalls atomic.Int32
target := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
targetCalls.Add(1)
http.Redirect(writer, request, "https://redirect.example.com/fixture", http.StatusFound)
}))
defer target.Close()
dns := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
dnsCalls.Add(1)
if test.redirect && request.URL.Query().Get("name") == "example.com" {
writeDNSFixture(writer, "93.184.216.34")
return
}
writeDNSFixture(writer, test.addresses...)
}))
defer dns.Close()
cfg := defaultConfig()
cfg.dataDir, cfg.OutputDir, cfg.ProxyURL = t.TempDir(), t.TempDir(), "direct"
downloader := NewDownloader(cfg)
defer downloader.client.CloseIdleConnections()
transport := downloader.client.Transport.(*imageTransport)
transport.lookup = func(context.Context, string) ([]string, error) { return []string{"198.18.0.42"}, nil }
transport.resolver.endpoints = []string{dns.URL + "/first", dns.URL + "/second"}
transport.standard.TLSClientConfig = target.Client().Transport.(*http.Transport).TLSClientConfig.Clone()
transport.standard.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
if !test.redirect || address != "93.184.216.34:443" {
t.Error("unsafe destination reached the dialer", address)
return nil, errors.New("unexpected connection")
}
return (&net.Dialer{}).DialContext(ctx, network, target.Listener.Addr().String())
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
ctx = context.WithValue(context.WithValue(ctx, coverRequestKey{}, true), coverSourceKey{}, test.source)
request, _ := http.NewRequestWithContext(ctx, http.MethodGet, test.url, nil)
response, err := downloader.client.Do(request)
if response != nil {
response.Body.Close()
}
wantTargets := int32(0)
if test.redirect {
wantTargets = 1
}
if err == nil || dnsCalls.Load() != test.dnsCalls || targetCalls.Load() != wantTargets {
t.Fatal("unsafe cover resolution or redirect was accepted", dnsCalls.Load(), targetCalls.Load(), err)
}
})
}
}
@@ -0,0 +1,76 @@
package app
import (
"bytes"
"context"
"crypto/tls"
"errors"
"net"
"net/http"
"net/http/httptest"
"sync/atomic"
"testing"
"time"
)
func TestRegisteredCoverRecoversFromFakeIPWithoutPrivateConnections(t *testing.T) {
for _, mode := range []string{"public-fallback", "private-fallback", "mixed-fallback", "private-redirect"} {
t.Run(mode, func(t *testing.T) {
data := embySyntheticPoster(t)
var calls, fallbacks atomic.Int32
server := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if request.Host != "p3-novel.byteimg.com" || request.TLS.ServerName != "p3-novel.byteimg.com" || request.Header.Get("Referer") == "" {
t.Error("cover lost its original Host, SNI or source headers")
}
if mode == "private-redirect" {
http.Redirect(writer, request, "https://private.edge.example.org/cover", http.StatusFound)
return
}
writer.Write(data)
}))
defer server.Close()
standard := http.DefaultTransport.(*http.Transport).Clone()
standard.Proxy = nil
standard.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
standard.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
calls.Add(1)
if address != "93.184.216.34:443" {
t.Error("attempted a connection to unverified DNS", address)
return nil, errors.New("unverified connection")
}
return (&net.Dialer{}).DialContext(ctx, network, server.Listener.Addr().String())
}
transport := newImageTransport(standard, standard)
defer transport.CloseIdleConnections()
transport.lookup = func(context.Context, string) ([]string, error) {
return []string{"198.18.0.42"}, nil
}
transport.fallback = func(ctx context.Context, host string) ([]string, error) {
fallbacks.Add(1)
if ctx.Value(coverSourceKey{}) != sourceHongguo {
t.Error("shared image cache dropped the registered source context")
}
if mode == "private-fallback" || host == "private.edge.example.org" {
return []string{"127.0.0.1"}, nil
}
if mode == "mixed-fallback" {
return []string{"93.184.216.34", "10.0.0.1"}, nil
}
return []string{"93.184.216.34"}, nil
}
app := &UIApp{downloader: &Downloader{client: &http.Client{Transport: transport, Timeout: 3 * time.Second}}}
ctx := context.WithValue(context.Background(), coverSourceKey{}, sourceHongguo)
body, err := app.loadCoverImage(ctx, "https://p3-novel.byteimg.com/novel-pic/fixture.image", nil)
if mode == "public-fallback" {
if err != nil || !bytes.Equal(body, data) || calls.Load() != 1 || fallbacks.Load() != 1 {
t.Fatal("registered cover did not recover safely from Fake-IP", calls.Load(), fallbacks.Load(), err)
}
if _, err := app.loadCoverImage(ctx, "https://p3-novel.byteimg.com/novel-pic/fixture.image", nil); err != nil || calls.Load() != 1 {
t.Fatal("successful cover was not cached", err)
}
} else if err == nil || mode != "private-redirect" && calls.Load() != 0 || mode == "private-redirect" && calls.Load() != 1 {
t.Fatal("unsafe fallback or redirect reached the network", calls.Load(), err)
}
})
}
}
+170
View File
@@ -0,0 +1,170 @@
package app
import (
"context"
"crypto/tls"
"errors"
"io"
"net"
"net/http"
"net/http/httptest"
"net/url"
"reflect"
"strings"
"sync"
"sync/atomic"
"testing"
)
func TestImageURLAcceptsPublicCDNsWithoutDomainList(t *testing.T) {
for _, address := range []string{"https://new-covers.example.org/a?token=%2F%2b%3D", "https://another.example.com/cover", "https://93.184.216.34/image", "https://[2606:4700:4700::1111]/image"} {
parsed, err := url.Parse(address)
if err != nil || !validImageURL(parsed) {
t.Fatal("public cover URL required a hard-coded domain", address, err)
}
}
for _, address := range []string{
"http://example.org/a", "file:///etc/passwd", "https://user@example.org/a", "https://example.org:8443/a",
"https://localhost/a", "https://app.localhost/a", "https://router.local/a", "https://private.lan/a",
"https://127.0.0.1/a", "https://10.1.2.3/a", "https://172.16.0.1/a", "https://192.168.1.1/a",
"https://169.254.169.254/a", "https://100.100.100.200/a", "https://0.0.0.0/a", "https://255.255.255.255/a",
"https://[::1]/a", "https://[::ffff:127.0.0.1]/a", "https://[fc00::1]/a", "https://[fe80::1%25en0]/a",
"https://[64:ff9b::a00:1]/a", "https://example.invalid/a", "https://empty../a",
} {
parsed, err := url.Parse(address)
if err == nil && validImageURL(parsed) {
t.Fatal("unsafe cover URL was accepted", address)
}
}
}
func TestImageTransportRejectsPrivateDNSBeforeAnyConnection(t *testing.T) {
for _, addresses := range [][]string{{"127.0.0.1"}, {"10.0.0.1"}, {"100.100.100.200"}, {"::ffff:192.168.1.1"}, {"93.184.216.34", "169.254.169.254"}} {
t.Run(strings.Join(addresses, "/"), func(t *testing.T) {
var calls atomic.Int32
standard := http.DefaultTransport.(*http.Transport).Clone()
standard.DialContext = func(context.Context, string, string) (net.Conn, error) {
calls.Add(1)
return nil, errors.New("must not connect")
}
transport := newImageTransport(standard, standard)
transport.lookup = func(context.Context, string) ([]string, error) { return addresses, nil }
transport.fallback = func(context.Context, string) ([]string, error) {
calls.Add(1)
return []string{"93.184.216.34"}, nil
}
request := httptest.NewRequest(http.MethodGet, "https://rebound.example.org/cover", nil)
request = request.WithContext(context.WithValue(request.Context(), coverRequestKey{}, true))
if _, err := transport.RoundTrip(request); err == nil || calls.Load() != 0 {
t.Fatal("private DNS was fetched or retried as a public destination", calls.Load(), err)
}
})
}
}
func TestImageTransportPinsDNSPreservesSNIAndProxy(t *testing.T) {
for _, mode := range []string{"direct", "http-proxy", "https-proxy"} {
t.Run(mode, func(t *testing.T) {
var targetCalls atomic.Int32
target := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
targetCalls.Add(1)
if r.Host != "new-covers.example.org" || r.TLS.ServerName != "new-covers.example.org" || r.URL.RawQuery != "signature=%2F%2b%3D" || r.Header.Get("Referer") != "https://source.example.org/" {
t.Error("pinned image request changed the hostname, TLS name, query or source headers")
}
io.WriteString(w, "synthetic transport fixture")
}))
defer target.Close()
var mu sync.Mutex
var dials, connects []string
proxyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
mu.Lock()
connects = append(connects, r.Host)
mu.Unlock()
if r.Method != http.MethodConnect || r.Host != "93.184.216.34:443" {
http.Error(w, "unexpected CONNECT destination", http.StatusBadRequest)
return
}
upstream, err := net.Dial("tcp", target.Listener.Addr().String())
if err != nil {
t.Error(err)
return
}
connection, buffer, err := w.(http.Hijacker).Hijack()
if err != nil {
upstream.Close()
t.Error(err)
return
}
buffer.WriteString("HTTP/1.1 200 Connection Established\r\n\r\n")
buffer.Flush()
go func() { io.Copy(upstream, buffer); upstream.Close() }()
io.Copy(connection, upstream)
connection.Close()
})
standard := http.DefaultTransport.(*http.Transport).Clone()
standard.Proxy = nil
standard.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
var proxyAddress string
if mode != "direct" {
var proxy *httptest.Server
if mode == "https-proxy" {
proxy = httptest.NewTLSServer(proxyHandler)
} else {
proxy = httptest.NewServer(proxyHandler)
}
defer proxy.Close()
proxyURL, _ := url.Parse(proxy.URL)
proxyAddress = proxyURL.Host
standard.Proxy = http.ProxyURL(proxyURL)
}
standard.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
mu.Lock()
dials = append(dials, address)
mu.Unlock()
if mode == "direct" && address == "93.184.216.34:443" {
address = target.Listener.Addr().String()
} else if mode == "direct" || address != proxyAddress {
return nil, errors.New("request bypassed its pinned address or configured proxy")
}
return (&net.Dialer{}).DialContext(ctx, network, address)
}
transport := newImageTransport(standard, standard)
defer transport.CloseIdleConnections()
transport.lookup = func(ctx context.Context, host string) ([]string, error) {
if host != "new-covers.example.org" {
t.Error("unexpected hostname", host)
}
return []string{"93.184.216.34"}, nil
}
request, _ := http.NewRequestWithContext(context.WithValue(context.Background(), coverRequestKey{}, true), http.MethodGet, "https://new-covers.example.org/cover?signature=%2F%2b%3D", nil)
request.Header.Set("Referer", "https://source.example.org/")
response, err := (&http.Client{Transport: transport}).Do(request)
if err != nil {
t.Fatal(err)
}
body, err := io.ReadAll(response.Body)
response.Body.Close()
if err != nil || string(body) != "synthetic transport fixture" || response.Request.URL.String() != request.URL.String() {
t.Fatal("pinned response lost the original address", err)
}
mu.Lock()
defer mu.Unlock()
if targetCalls.Load() != 1 || len(dials) != 1 || mode == "direct" && len(connects) != 0 || mode != "direct" && !reflect.DeepEqual(connects, []string{"93.184.216.34:443"}) {
t.Fatal("image request did not follow the expected pinned route", dials, connects, targetCalls.Load())
}
})
}
}
func TestImageSourceRequiresRegisteredCoverForEveryViewer(t *testing.T) {
address := "https://new-covers.example.org/cover?token=fixture"
app := &UIApp{dramas: []Drama{{ID: "hongguo:7000000000000000001", Source: sourceHongguo, CoverURL: address}}}
for _, ctx := range []context.Context{context.Background(), withSourceScope(context.Background(), accountRecord{Admin: true}), withSourceScope(context.Background(), accountRecord{Sources: []string{sourceHongguo}})} {
if source, allowed := app.imageSource(ctx, address); !allowed || source != sourceHongguo {
t.Fatal("registered cover was rejected")
}
if _, allowed := app.imageSource(ctx, address+"-forged"); allowed {
t.Fatal("unregistered cover was accepted")
}
}
}
+287
View File
@@ -0,0 +1,287 @@
package app
import (
"context"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"sync"
"time"
)
type requestGate struct {
busy bool
lastStart time.Time
}
type requestBackoff struct {
host string
until time.Time
status int
reason string
}
type httpStatusError struct {
host string
status int
reason string
}
func (failure *httpStatusError) Error() string {
reason := ""
if failure.reason != "" {
reason = ":" + failure.reason
}
return fmt.Sprintf("%s HTTP %d%s", failure.host, failure.status, reason)
}
func (backoff *requestBackoff) Error() string {
reason := ""
if backoff.reason != "" {
reason = ":" + backoff.reason
}
return fmt.Sprintf("%s HTTP %d%s,已暂停该域名请求,%s 后可重试", backoff.host, backoff.status, reason, backoff.until.Format("15:04:05"))
}
func catalogResponseBlockReason(response *http.Response, body []byte) string {
if len(body) > 64*1024 {
body = body[:64*1024]
}
page := strings.ToLower(string(body))
reason := ""
if strings.Contains(page, "cloudflare") && (strings.Contains(page, "sorry, you have been blocked") || strings.Contains(page, "you are unable to access")) {
reason = "Cloudflare 拒绝了当前请求"
} else if strings.EqualFold(response.Header.Get("Cf-Mitigated"), "challenge") || strings.Contains(page, "_cf_chl_opt") || strings.Contains(page, "<title>just a moment") && strings.Contains(page, "cloudflare") {
reason = "站点要求浏览器验证,当前请求无法通过"
}
return reason
}
func (d *Downloader) catalogResponseError(request *http.Request, response *http.Response, body []byte) error {
reason := catalogResponseBlockReason(response, body)
if d.limiter != nil && reason != "" && response.StatusCode >= 200 && response.StatusCode < 300 {
d.limiter.observeBlock(request, response)
}
host := strings.ToLower(request.URL.Hostname())
if d.limiter != nil && (response.StatusCode == http.StatusForbidden || response.StatusCode == http.StatusTooManyRequests || reason != "") {
d.limiter.mu.Lock()
backoffs := d.limiter.backoffs
if background, _ := request.Context().Value(backgroundCatalogKey{}).(bool); background && response.StatusCode != http.StatusTooManyRequests {
backoffs = d.limiter.backgroundBackoffs
}
if previous := backoffs[host]; previous != nil && previous.status == response.StatusCode {
updated := *previous
if reason != "" {
updated.reason = reason
}
backoffs[host] = &updated
d.limiter.mu.Unlock()
return &updated
}
d.limiter.mu.Unlock()
}
if reason != "" {
return &httpStatusError{host: host, status: response.StatusCode, reason: reason}
}
return &httpStatusError{host: host, status: response.StatusCode}
}
type requestLimiter struct {
concurrency int
active int
foregroundWaiting int
lastForeground time.Time
interval time.Duration
mu sync.Mutex
gates map[string]*requestGate
backoffs map[string]*requestBackoff
backgroundBackoffs map[string]*requestBackoff
changed chan struct{}
}
type backgroundCatalogKey struct{}
func newRequestLimiter(concurrency int, interval time.Duration) *requestLimiter {
if concurrency < 1 {
concurrency = 1
}
return &requestLimiter{concurrency: concurrency, interval: interval, gates: map[string]*requestGate{}, backoffs: map[string]*requestBackoff{}, backgroundBackoffs: map[string]*requestBackoff{}, changed: make(chan struct{})}
}
func (limiter *requestLimiter) configure(concurrency int, interval time.Duration) {
limiter.mu.Lock()
defer limiter.mu.Unlock()
limiter.concurrency = concurrency
limiter.interval = interval
limiter.notifyLocked()
}
func (limiter *requestLimiter) notifyLocked() {
close(limiter.changed)
limiter.changed = make(chan struct{})
}
func (limiter *requestLimiter) acquire(ctx context.Context, request *http.Request) (func(), error) {
background, _ := ctx.Value(backgroundCatalogKey{}).(bool)
host := strings.ToLower(request.URL.Hostname())
source := providerSourceForURL(request.URL.String())
if source == "" {
source = host
if strings.HasSuffix(host, ".cloudfront.net") {
source = "cloudfront"
}
}
limiter.mu.Lock()
if !background {
limiter.foregroundWaiting++
defer func() {
limiter.mu.Lock()
limiter.foregroundWaiting--
limiter.notifyLocked()
limiter.mu.Unlock()
}()
}
gate := limiter.gates[source]
if gate == nil {
gate = &requestGate{}
limiter.gates[source] = gate
}
limiter.mu.Unlock()
for {
if err := ctx.Err(); err != nil {
return nil, err
}
limiter.mu.Lock()
if backoff := limiter.backoffs[host]; backoff != nil && time.Now().Before(backoff.until) {
limiter.mu.Unlock()
return nil, backoff
}
if backoff := limiter.backgroundBackoffs[host]; background && backoff != nil && time.Now().Before(backoff.until) {
limiter.mu.Unlock()
return nil, backoff
}
delay := time.Until(gate.lastStart.Add(limiter.interval))
if background {
if idleDelay := time.Until(limiter.lastForeground.Add(limiter.interval)); idleDelay > delay {
delay = idleDelay
}
}
lowPriorityReady := !background || limiter.foregroundWaiting == 0 && limiter.active == 0
if lowPriorityReady && !gate.busy && limiter.active < limiter.concurrency && delay <= 0 {
gate.busy = true
gate.lastStart = time.Now()
limiter.active++
limiter.mu.Unlock()
break
}
changed := limiter.changed
limiter.mu.Unlock()
var timer *time.Timer
var deadline <-chan time.Time
if delay > 0 {
timer = time.NewTimer(delay)
deadline = timer.C
}
select {
case <-changed:
case <-deadline:
case <-ctx.Done():
}
if timer != nil {
timer.Stop()
}
}
var once sync.Once
return func() {
once.Do(func() {
limiter.mu.Lock()
limiter.active--
gate.busy = false
if !background {
limiter.lastForeground = time.Now()
}
limiter.notifyLocked()
limiter.mu.Unlock()
})
}, nil
}
func (limiter *requestLimiter) observe(request *http.Request, response *http.Response) {
if response.StatusCode != http.StatusForbidden && response.StatusCode != http.StatusTooManyRequests && !strings.EqualFold(response.Header.Get("Cf-Mitigated"), "challenge") {
return
}
limiter.observeBlock(request, response)
}
func (limiter *requestLimiter) observeBlock(request *http.Request, response *http.Response) {
delay := time.Minute
if seconds, err := strconv.Atoi(response.Header.Get("Retry-After")); err == nil && seconds > 0 {
delay = time.Duration(seconds) * time.Second
} else if until, err := http.ParseTime(response.Header.Get("Retry-After")); err == nil && time.Until(until) > 0 {
delay = time.Until(until)
}
if delay > 30*time.Minute {
delay = 30 * time.Minute
}
host := strings.ToLower(request.URL.Hostname())
limiter.mu.Lock()
backoffs := limiter.backoffs
if background, _ := request.Context().Value(backgroundCatalogKey{}).(bool); background && response.StatusCode != http.StatusTooManyRequests {
backoffs = limiter.backgroundBackoffs
}
backoffs[host] = &requestBackoff{host: host, until: time.Now().Add(delay), status: response.StatusCode}
limiter.notifyLocked()
limiter.mu.Unlock()
}
type limitedResponseBody struct {
io.ReadCloser
release func()
}
func (body *limitedResponseBody) Close() error {
err := body.ReadCloser.Close()
body.release()
return err
}
func (d *Downloader) doCatalogRequest(request *http.Request) (*http.Response, error) {
return d.doCatalogRequestWithTimeout(request, 0)
}
func (d *Downloader) doCatalogRequestWithTimeout(request *http.Request, timeout time.Duration) (*http.Response, error) {
return d.doPreparedCatalogRequest(request, timeout, nil)
}
func (d *Downloader) doPreparedCatalogRequest(request *http.Request, timeout time.Duration, prepare func(*http.Request)) (*http.Response, error) {
if background, _ := request.Context().Value(backgroundCatalogKey{}).(bool); background && (timeout <= 0 || timeout > 8*time.Second) {
timeout = 8 * time.Second
}
release, err := d.limiter.acquire(request.Context(), request)
if err != nil {
return nil, err
}
cancel := func() {}
if timeout > 0 {
var ctx context.Context
ctx, cancel = context.WithTimeout(request.Context(), timeout)
request = request.WithContext(ctx)
}
if prepare != nil {
prepare(request)
}
response, err := d.client.Do(request)
if err != nil {
cancel()
release()
return nil, err
}
d.limiter.observe(request, response)
response.Body = &limitedResponseBody{ReadCloser: response.Body, release: func() { cancel(); release() }}
return response, nil
}
+78
View File
@@ -0,0 +1,78 @@
package app
import (
"context"
"errors"
"net/http"
"strings"
"sync/atomic"
"testing"
)
func TestCatalogBlockReasonSurvivesCooldown(t *testing.T) {
var calls atomic.Int32
d := rankingTestDownloader(t, func(request *http.Request) (*http.Response, error) {
calls.Add(1)
return rankingHTTPResponse(request, 403, `<title>Attention Required! | Cloudflare</title>Sorry, you have been blocked. sensitive-ray-id`), nil
})
for i := 0; i < 2; i++ {
_, err := d.fetchProviderText(context.Background(), "https://blocked.example/text", "")
var backoff *requestBackoff
if !errors.As(err, &backoff) || !strings.Contains(err.Error(), "Cloudflare 拒绝了当前请求") || strings.Contains(err.Error(), "sensitive-ray-id") {
t.Fatal("block reason was lost or raw page was exposed", err)
}
}
if calls.Load() != 1 {
t.Fatal("cooldown repeated an upstream request")
}
}
func TestCatalogBackgroundBackoffStaysSeparate(t *testing.T) {
for _, status := range []int{403, 429} {
t.Run(http.StatusText(status), func(t *testing.T) {
var calls atomic.Int32
d := rankingTestDownloader(t, func(request *http.Request) (*http.Response, error) {
calls.Add(1)
if optional, _ := request.Context().Value(backgroundCatalogKey{}).(bool); optional {
return rankingHTTPResponse(request, status, "optional unavailable"), nil
}
return rankingHTTPResponse(request, 200, "foreground ready"), nil
})
optional := context.WithValue(context.Background(), backgroundCatalogKey{}, true)
if _, err := d.fetchProviderText(optional, "https://catalog.example/detail", ""); err == nil {
t.Fatal("missing optional failure")
}
body, err := d.fetchProviderText(context.Background(), "https://catalog.example/list", "")
if status == 403 && (err != nil || body != "foreground ready" || calls.Load() != 2) {
t.Fatal("optional 403 blocked normal catalog", err)
}
if status == 429 && (err == nil || calls.Load() != 1) {
t.Fatal("site-wide rate limit was ignored")
}
})
}
}
func TestCatalogChallenge200StopsWithoutPublishingSuccess(t *testing.T) {
var calls atomic.Int32
d := rankingTestDownloader(t, func(request *http.Request) (*http.Response, error) {
calls.Add(1)
response := rankingHTTPResponse(request, 200, "<title>Just a moment...</title>Cloudflare fixture")
response.Header.Set("Cf-Mitigated", "challenge")
return response, nil
})
for index := 0; index < 2; index++ {
body, err := d.fetchProviderText(context.Background(), "https://challenge.example/text", "")
var backoff *requestBackoff
if body != "" || !errors.As(err, &backoff) || !strings.Contains(err.Error(), "浏览器验证") {
t.Fatal("challenge response treated as successful content", err)
}
}
if calls.Load() != 1 {
t.Fatal("challenge repeated requests during cooldown")
}
response := &http.Response{StatusCode: 200, Header: make(http.Header)}
if catalogResponseBlockReason(response, []byte("<title>Local fixture</title>Just a moment, Cloudflare text in a description")) != "" {
t.Fatal("ordinary description was mistaken for a challenge")
}
}