github.com/icyphox/x@v0.0.355-0.20220311094250-029bd783e8b8/fetcher/fetcher.go (about)

     1  package fetcher
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/base64"
     6  	stderrors "errors"
     7  	"io"
     8  	"net/http"
     9  	"os"
    10  	"strings"
    11  
    12  	"github.com/hashicorp/go-retryablehttp"
    13  	"github.com/pkg/errors"
    14  
    15  	"github.com/ory/x/httpx"
    16  	"github.com/ory/x/stringsx"
    17  )
    18  
    19  // Fetcher is able to load file contents from http, https, file, and base64 locations.
    20  type Fetcher struct {
    21  	hc *retryablehttp.Client
    22  }
    23  
    24  type opts struct {
    25  	hc *retryablehttp.Client
    26  }
    27  
    28  var ErrUnknownScheme = stderrors.New("unknown scheme")
    29  
    30  // WithClient sets the http.Client the fetcher uses.
    31  func WithClient(hc *retryablehttp.Client) func(*opts) {
    32  	return func(o *opts) {
    33  		o.hc = hc
    34  	}
    35  }
    36  
    37  func newOpts() *opts {
    38  	return &opts{
    39  		hc: httpx.NewResilientClient(),
    40  	}
    41  }
    42  
    43  // NewFetcher creates a new fetcher instance.
    44  func NewFetcher(opts ...func(*opts)) *Fetcher {
    45  	o := newOpts()
    46  	for _, f := range opts {
    47  		f(o)
    48  	}
    49  	return &Fetcher{hc: o.hc}
    50  }
    51  
    52  // Fetch fetches the file contents from the source.
    53  func (f *Fetcher) Fetch(source string) (*bytes.Buffer, error) {
    54  	switch s := stringsx.SwitchPrefix(source); {
    55  	case s.HasPrefix("http://"), s.HasPrefix("https://"):
    56  		return f.fetchRemote(source)
    57  	case s.HasPrefix("file://"):
    58  		return f.fetchFile(strings.Replace(source, "file://", "", 1))
    59  	case s.HasPrefix("base64://"):
    60  		src, err := base64.StdEncoding.DecodeString(strings.Replace(source, "base64://", "", 1))
    61  		if err != nil {
    62  			return nil, errors.Wrapf(err, "rule: %s", source)
    63  		}
    64  		return bytes.NewBuffer(src), nil
    65  	default:
    66  		return nil, errors.Wrap(ErrUnknownScheme, s.ToUnknownPrefixErr().Error())
    67  	}
    68  }
    69  
    70  func (f *Fetcher) fetchRemote(source string) (*bytes.Buffer, error) {
    71  	res, err := f.hc.Get(source)
    72  	if err != nil {
    73  		return nil, errors.Wrapf(err, "rule: %s", source)
    74  	}
    75  	defer res.Body.Close()
    76  
    77  	if res.StatusCode != http.StatusOK {
    78  		return nil, errors.Errorf("expected http response status code 200 but got %d when fetching: %s", res.StatusCode, source)
    79  	}
    80  
    81  	return f.decode(res.Body)
    82  }
    83  
    84  func (f *Fetcher) fetchFile(source string) (*bytes.Buffer, error) {
    85  	fp, err := os.Open(source)
    86  	if err != nil {
    87  		return nil, errors.Wrapf(err, "unable to fetch from source: %s", source)
    88  	}
    89  	defer fp.Close()
    90  
    91  	return f.decode(fp)
    92  }
    93  
    94  func (f *Fetcher) decode(r io.Reader) (*bytes.Buffer, error) {
    95  	var b bytes.Buffer
    96  	if _, err := io.Copy(&b, r); err != nil {
    97  		return nil, err
    98  	}
    99  	return &b, nil
   100  }