github.com/rck/u-root@v0.0.0-20180106144920-7eb602e381bb/pkg/pxe/schemes.go (about)

     1  package pxe
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"io"
     7  	"net/http"
     8  	"net/url"
     9  	"os"
    10  	"path/filepath"
    11  
    12  	"pack.ag/tftp"
    13  )
    14  
    15  var (
    16  	// ErrNoSuchScheme is returned by Schemes.GetFile and
    17  	// Schemes.LazyGetFile if there is no registered FileScheme
    18  	// implementation for the given URI scheme.
    19  	ErrNoSuchScheme = errors.New("no such scheme")
    20  )
    21  
    22  // FileScheme represents the implementation of a URI scheme and gives access to
    23  // downloading files of that scheme.
    24  //
    25  // For example, an http FileScheme implementation would download files using
    26  // the HTTP protocol.
    27  type FileScheme interface {
    28  	// GetFile returns a reader that gives the contents of `u`.
    29  	//
    30  	// It may do so by downloading `u` and placing it in a buffer, or by
    31  	// returning an io.Reader that downloads the file.
    32  	GetFile(u *url.URL) (io.Reader, error)
    33  }
    34  
    35  var (
    36  	// DefaultHTTPClient is the default HTTP FileScheme.
    37  	//
    38  	// It is not recommended to use this for HTTPS. We recommend creating an
    39  	// http.Client that accepts only a private pool of certificates.
    40  	DefaultHTTPClient = NewHTTPClient(http.DefaultClient)
    41  
    42  	// DefaultTFTPClient is the default TFTP FileScheme.
    43  	DefaultTFTPClient FileScheme
    44  
    45  	// DefaultSchemes are the schemes supported by PXE by default.
    46  	DefaultSchemes = Schemes{
    47  		"tftp": NewCachedFileScheme(DefaultTFTPClient),
    48  		"http": NewCachedFileScheme(DefaultHTTPClient),
    49  		"file": NewCachedFileScheme(&LocalFileClient{}),
    50  	}
    51  )
    52  
    53  func init() {
    54  	c, err := tftp.NewClient()
    55  	if err != nil {
    56  		panic(fmt.Sprintf("tftp.NewClient failed: %v", err))
    57  	}
    58  	DefaultTFTPClient = NewTFTPClient(c)
    59  }
    60  
    61  // Schemes is a map of URI scheme identifier -> implementation that can
    62  // download a file for that scheme.
    63  type Schemes map[string]FileScheme
    64  
    65  // RegisterScheme calls DefaultSchemes.Register.
    66  func RegisterScheme(scheme string, fs FileScheme) {
    67  	DefaultSchemes.Register(scheme, fs)
    68  }
    69  
    70  // Register registers a scheme identified by `scheme` to be `fs`.
    71  func (s Schemes) Register(scheme string, fs FileScheme) {
    72  	s[scheme] = fs
    73  }
    74  
    75  func parseURI(uri string, wd *url.URL) (*url.URL, error) {
    76  	u, err := url.Parse(uri)
    77  	if err != nil {
    78  		return nil, err
    79  	}
    80  
    81  	if len(u.Scheme) == 0 {
    82  		u.Scheme = wd.Scheme
    83  
    84  		if len(u.Host) == 0 {
    85  			// If this is not there, it was likely just a path.
    86  			u.Host = wd.Host
    87  			u.Path = filepath.Join(wd.Path, filepath.Clean(u.Path))
    88  		}
    89  	}
    90  	return u, nil
    91  }
    92  
    93  // GetFile downloads a file via DefaultSchemes. See Schemes.GetFile for
    94  // details.
    95  func GetFile(uri string, wd *url.URL) (io.Reader, error) {
    96  	return DefaultSchemes.GetFile(uri, wd)
    97  }
    98  
    99  // GetFile downloads the file with the given `uri`. `uri.Scheme` is used to
   100  // select the FileScheme via `s`.
   101  //
   102  // If `s` does not contain a FileScheme for `uri.Scheme`, ErrNoSuchScheme is
   103  // returned.
   104  //
   105  // If `uri` is just a relative path and not a full URI, `wd` is used as the
   106  // "working directory" of that relative path; the resulting URI is roughly
   107  // `path.Join(wd.String(), uri)`.
   108  func (s Schemes) GetFile(uri string, wd *url.URL) (io.Reader, error) {
   109  	u, err := parseURI(uri, wd)
   110  	if err != nil {
   111  		return nil, err
   112  	}
   113  
   114  	fg, ok := s[u.Scheme]
   115  	if !ok {
   116  		return nil, ErrNoSuchScheme
   117  	}
   118  	return fg.GetFile(u)
   119  }
   120  
   121  // LazyGetFile calls LazyGetFile on DefaultSchemes. See Schemes.LazyGetFile.
   122  func LazyGetFile(uri string, wd *url.URL) (io.Reader, error) {
   123  	return DefaultSchemes.LazyGetFile(uri, wd)
   124  }
   125  
   126  // LazyGetFile returns a reader that will download the file given by `uri` when
   127  // Read is called, based on `uri`s scheme. See Schemes.GetFile for more
   128  // details.
   129  func (s Schemes) LazyGetFile(uri string, wd *url.URL) (io.Reader, error) {
   130  	u, err := parseURI(uri, wd)
   131  	if err != nil {
   132  		return nil, err
   133  	}
   134  
   135  	fg, ok := s[u.Scheme]
   136  	if !ok {
   137  		return nil, fmt.Errorf("could not get file based on scheme %q: no such scheme registered", u.Scheme)
   138  	}
   139  
   140  	return NewLazyOpener(func() (io.Reader, error) {
   141  		return fg.GetFile(u)
   142  	}), nil
   143  }
   144  
   145  // TFTPClient implements FileScheme for TFTP files.
   146  type TFTPClient struct {
   147  	c *tftp.Client
   148  }
   149  
   150  // NewTFTPClient returns a new TFTP client based on the given tftp.Client.
   151  func NewTFTPClient(c *tftp.Client) FileScheme {
   152  	return &TFTPClient{
   153  		c: c,
   154  	}
   155  }
   156  
   157  // GetFile implements FileScheme.GetFile.
   158  func (t *TFTPClient) GetFile(u *url.URL) (io.Reader, error) {
   159  	r, err := t.c.Get(u.String())
   160  	if err != nil {
   161  		return nil, fmt.Errorf("failed to get %q: %v", u, err)
   162  	}
   163  	return r, nil
   164  }
   165  
   166  // HTTPClient implements FileScheme for HTTP files.
   167  type HTTPClient struct {
   168  	c *http.Client
   169  }
   170  
   171  // NewHTTPClient returns a new HTTP FileScheme based on the given http.Client.
   172  func NewHTTPClient(c *http.Client) *HTTPClient {
   173  	return &HTTPClient{
   174  		c: c,
   175  	}
   176  }
   177  
   178  // GetFile implements FileScheme.GetFile.
   179  func (h HTTPClient) GetFile(u *url.URL) (io.Reader, error) {
   180  	resp, err := h.c.Get(u.String())
   181  	if err != nil {
   182  		return nil, fmt.Errorf("Could not download file %s: %v", u, err)
   183  	}
   184  	if resp.StatusCode != 200 {
   185  		return nil, fmt.Errorf("could not download file %s: response %v", u, resp)
   186  	}
   187  	return resp.Body, nil
   188  }
   189  
   190  // LocalFileClient implements FileScheme for files on disk.
   191  type LocalFileClient struct{}
   192  
   193  // GetFile implements FileScheme.GetFile.
   194  func (lfs LocalFileClient) GetFile(u *url.URL) (io.Reader, error) {
   195  	return os.Open(filepath.Clean(u.Path))
   196  }
   197  
   198  type cachedFile struct {
   199  	cr  *CachingReader
   200  	err error
   201  }
   202  
   203  // CachedFileScheme implements FileScheme and caches files downloaded from a
   204  // FileScheme.
   205  type CachedFileScheme struct {
   206  	fs FileScheme
   207  
   208  	// cache is a map of URI string -> cached file or error object.
   209  	cache map[string]cachedFile
   210  }
   211  
   212  // NewCachedFileScheme returns a caching wrapper for the given FileScheme `fs`.
   213  func NewCachedFileScheme(fs FileScheme) FileScheme {
   214  	return &CachedFileScheme{
   215  		fs:    fs,
   216  		cache: make(map[string]cachedFile),
   217  	}
   218  }
   219  
   220  // GetFile implements FileScheme.GetFile.
   221  func (cc *CachedFileScheme) GetFile(u *url.URL) (io.Reader, error) {
   222  	uri := u.String()
   223  	if cf, ok := cc.cache[uri]; ok {
   224  		if cf.err != nil {
   225  			return nil, cf.err
   226  		}
   227  		return cf.cr.NewReader(), nil
   228  	}
   229  
   230  	r, err := cc.fs.GetFile(u)
   231  	if err != nil {
   232  		cc.cache[uri] = cachedFile{err: err}
   233  		return nil, err
   234  	}
   235  	cr := NewCachingReader(r)
   236  	cc.cache[uri] = cachedFile{cr: cr}
   237  	return cr.NewReader(), nil
   238  }