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 }