vendor/github.com/docker/distribution/registry/client/transport/transport.go

   1 package transport
   2
   3 import (
   4         "io"
   5         "net/http"
   6         "sync"
   7 )
   8
   9 // RequestModifier represents an object which will do an inplace
  10 // modification of an HTTP request.
  11 type RequestModifier interface {
  12         ModifyRequest(*http.Request) error
  13 }
  14
  15 type headerModifier http.Header
  16
  17 // NewHeaderRequestModifier returns a new RequestModifier which will
  18 // add the given headers to a request.
  19 func NewHeaderRequestModifier(header http.Header) RequestModifier {
  20         return headerModifier(header)
  21 }
  22
  23 func (h headerModifier) ModifyRequest(req *http.Request) error {
  24         for k, s := range http.Header(h) {
  25                 req.Header[k] = append(req.Header[k], s...)
  26         }
  27
  28         return nil
  29 }
  30
  31 // NewTransport creates a new transport which will apply modifiers to
  32 // the request on a RoundTrip call.
  33 func NewTransport(base http.RoundTripper, modifiers ...RequestModifier) http.RoundTripper {
  34         return &transport{
  35                 Modifiers: modifiers,
  36                 Base:      base,
  37         }
  38 }
  39
  40 // transport is an http.RoundTripper that makes HTTP requests after
  41 // copying and modifying the request
  42 type transport struct {
  43         Modifiers []RequestModifier
  44         Base      http.RoundTripper
  45
  46         mu     sync.Mutex                      // guards modReq
  47         modReq map[*http.Request]*http.Request // original -> modified
  48 }
  49
  50 // RoundTrip authorizes and authenticates the request with an
  51 // access token. If no token exists or token is expired,
  52 // tries to refresh/fetch a new token.
  53 func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) {
  54         req2 := cloneRequest(req)
  55         for _, modifier := range t.Modifiers {
  56                 if err := modifier.ModifyRequest(req2); err != nil {
  57                         return nil, err
  58                 }
  59         }
  60
  61         t.setModReq(req, req2)
  62         res, err := t.base().RoundTrip(req2)
  63         if err != nil {
  64                 t.setModReq(req, nil)
  65                 return nil, err
  66         }
  67         res.Body = &onEOFReader{
  68                 rc: res.Body,
  69                 fn: func() { t.setModReq(req, nil) },
  70         }
  71         return res, nil
  72 }
  73
  74 // CancelRequest cancels an in-flight request by closing its connection.
  75 func (t *transport) CancelRequest(req *http.Request) {
  76         type canceler interface {
  77                 CancelRequest(*http.Request)
  78         }
  79         if cr, ok := t.base().(canceler); ok {
  80                 t.mu.Lock()
  81                 modReq := t.modReq[req]
  82                 delete(t.modReq, req)
  83                 t.mu.Unlock()
  84                 cr.CancelRequest(modReq)
  85         }
  86 }
  87
  88 func (t *transport) base() http.RoundTripper {
  89         if t.Base != nil {
  90                 return t.Base
  91         }
  92         return http.DefaultTransport
  93 }
  94
  95 func (t *transport) setModReq(orig, mod *http.Request) {
  96         t.mu.Lock()
  97         defer t.mu.Unlock()
  98         if t.modReq == nil {
  99                 t.modReq = make(map[*http.Request]*http.Request)
 100         }
 101         if mod == nil {
 102                 delete(t.modReq, orig)
 103         } else {
 104                 t.modReq[orig] = mod
 105         }
 106 }
 107
 108 // cloneRequest returns a clone of the provided *http.Request.
 109 // The clone is a shallow copy of the struct and its Header map.
 110 func cloneRequest(r *http.Request) *http.Request {
 111         // shallow copy of the struct
 112         r2 := new(http.Request)
 113         *r2 = *r
 114         // deep copy of the Header
 115         r2.Header = make(http.Header, len(r.Header))
 116         for k, s := range r.Header {
 117                 r2.Header[k] = append([]string(nil), s...)
 118         }
 119
 120         return r2
 121 }
 122
 123 type onEOFReader struct {
 124         rc io.ReadCloser
 125         fn func()
 126 }
 127
 128 func (r *onEOFReader) Read(p []byte) (n int, err error) {
 129         n, err = r.rc.Read(p)
 130         if err == io.EOF {
 131                 r.runFunc()
 132         }
 133         return
 134 }
 135
 136 func (r *onEOFReader) Close() error {
 137         err := r.rc.Close()
 138         r.runFunc()
 139         return err
 140 }
 141
 142 func (r *onEOFReader) runFunc() {
 143         if fn := r.fn; fn != nil {
 144                 fn()
 145                 r.fn = nil
 146         }
 147 }