Files
NianAIGC/backend/internal/assets/remote.go

109 lines
3.0 KiB
Go

package assets
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"net/http"
"net/url"
)
var errRemotePolicyRequired = errors.New("remote fetcher requires a bounded client and destination policy")
type HTTPRemoteFetcher struct {
client *http.Client
maxBytes int64
policy DestinationPolicy
}
// DestinationPolicy is the injectable outbound-network policy seam.
type DestinationPolicy interface {
Validate(context.Context, *url.URL) error
}
// NewPolicyHTTPRemoteFetcher preserves the existing constructor while offering
// production composition a fail-closed destination-policy seam.
func NewPolicyHTTPRemoteFetcher(client *http.Client, maxBytes int64, policy DestinationPolicy) (*HTTPRemoteFetcher, error) {
if client == nil || client.Timeout <= 0 || policy == nil {
return nil, errRemotePolicyRequired
}
base := *client
previousRedirect := base.CheckRedirect
base.CheckRedirect = func(req *http.Request, via []*http.Request) error {
if len(via) >= 10 {
return errors.New("stopped after 10 redirects")
}
if err := policy.Validate(req.Context(), req.URL); err != nil {
return err
}
if previousRedirect != nil {
return previousRedirect(req, via)
}
return nil
}
return &HTTPRemoteFetcher{client: &base, maxBytes: maxBytes, policy: policy}, nil
}
func NewHTTPRemoteFetcher(client *http.Client, maxBytes int64) *HTTPRemoteFetcher {
if client == nil {
client = http.DefaultClient
}
base := *client
previousRedirect := base.CheckRedirect
base.CheckRedirect = func(req *http.Request, via []*http.Request) error {
if req.URL.Scheme != "http" && req.URL.Scheme != "https" {
return ErrRemoteProtocol
}
if previousRedirect != nil {
return previousRedirect(req, via)
}
if len(via) >= 10 {
return errors.New("stopped after 10 redirects")
}
return nil
}
return &HTTPRemoteFetcher{client: &base, maxBytes: maxBytes}
}
func (f *HTTPRemoteFetcher) Fetch(ctx context.Context, rawURL string) (Blob, error) {
u, err := url.Parse(rawURL)
if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" {
return Blob{}, ErrRemoteProtocol
}
if u.User != nil {
return Blob{}, ErrRemoteProtocol
}
if f.policy != nil {
if err := f.policy.Validate(ctx, u); err != nil {
return Blob{}, err
}
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
if err != nil {
return Blob{}, err
}
resp, err := f.client.Do(req)
if err != nil {
return Blob{}, err
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return Blob{}, fmt.Errorf("remote asset returned HTTP %d", resp.StatusCode)
}
if f.maxBytes <= 0 {
return Blob{}, errors.New("remote size limit must be positive")
}
if resp.ContentLength > f.maxBytes {
return Blob{}, ErrRemoteTooLarge
}
b, err := io.ReadAll(io.LimitReader(resp.Body, f.maxBytes+1))
if err != nil {
return Blob{}, err
}
if int64(len(b)) > f.maxBytes {
return Blob{}, ErrRemoteTooLarge
}
return Blob{Body: io.NopCloser(bytes.NewReader(b)), ContentType: resp.Header.Get("Content-Type"), Size: int64(len(b))}, nil
}