diff --git a/config.go b/config.go index 7ae196b..c8ec41b 100644 --- a/config.go +++ b/config.go @@ -11,8 +11,9 @@ type UpstreamMatch struct { } type Upstream struct { - Server string `yaml:"server"` - Match UpstreamMatch `yaml:"match"` + Server string `yaml:"server"` + Match UpstreamMatch `yaml:"match"` + AllowedRedirect *string `yaml:"allowed-redirect"` } func (upstream Upstream) GetPath(orig string) (string, bool, error) { diff --git a/config.yaml b/config.yaml index 7d1627a..5c0d42d 100644 --- a/config.yaml +++ b/config.yaml @@ -83,6 +83,7 @@ upstream: match: match: /chaotic-aur/(.*) replace: /$1 + allowed-redirect: "^https?://cf-builds.garudalinux.org/.*/chaotic-aur/.*$" misc: first-chunk-bytes: 31457280 # 1024*1024*30, all upstreams with 200 response will wait for the first chunks to select a upstream by bandwidth, instead of by latency. defaults to 50M diff --git a/server.go b/server.go index f9723f1..18a24cc 100644 --- a/server.go +++ b/server.go @@ -17,13 +17,27 @@ import ( "time" ) +type reqCtxKey int + +const ( + reqCtxAllowedRedirect reqCtxKey = iota +) + var zeroTime time.Time var ( httpClient = http.Client{ - - // No redirect! + // check allowed redirect CheckRedirect: func(req *http.Request, via []*http.Request) error { + lastRequest := via[len(via)-1] + if allowedRedirect, ok := lastRequest.Context().Value(reqCtxAllowedRedirect).(string); ok { + if matched, err := regexp.MatchString(allowedRedirect, req.URL.String()); err != nil { + return err + } else if !matched { + return http.ErrUseLastResponse + } + return nil + } return http.ErrUseLastResponse }, } @@ -492,6 +506,11 @@ func (server *Server) tryUpstream(ctx context.Context, upstreamIdx int, r *http. if lastModified != zeroTime { method = http.MethodGet } + + if upstream.AllowedRedirect != nil { + ctx = context.WithValue(ctx, reqCtxAllowedRedirect, *upstream.AllowedRedirect) + } + request, err := http.NewRequestWithContext(ctx, method, newurl, nil) if err != nil { return nil, nil, err