github.com/zhouyu0/docker-note@v0.0.0-20190722021225-b8d3825084db/builder/builder-next/reqbodyhandler.go (about)

     1  package buildkit
     2  
     3  import (
     4  	"io"
     5  	"net/http"
     6  	"strings"
     7  	"sync"
     8  
     9  	"github.com/moby/buildkit/identity"
    10  	"github.com/pkg/errors"
    11  )
    12  
    13  const urlPrefix = "build-context-"
    14  
    15  type reqBodyHandler struct {
    16  	mu sync.Mutex
    17  	rt http.RoundTripper
    18  
    19  	requests map[string]io.ReadCloser
    20  }
    21  
    22  func newReqBodyHandler(rt http.RoundTripper) *reqBodyHandler {
    23  	return &reqBodyHandler{
    24  		rt:       rt,
    25  		requests: map[string]io.ReadCloser{},
    26  	}
    27  }
    28  
    29  func (h *reqBodyHandler) newRequest(rc io.ReadCloser) (string, func()) {
    30  	id := identity.NewID()
    31  	h.mu.Lock()
    32  	h.requests[id] = rc
    33  	h.mu.Unlock()
    34  	return "http://" + urlPrefix + id, func() {
    35  		h.mu.Lock()
    36  		delete(h.requests, id)
    37  		h.mu.Unlock()
    38  	}
    39  }
    40  
    41  func (h *reqBodyHandler) RoundTrip(req *http.Request) (*http.Response, error) {
    42  	host := req.URL.Host
    43  	if strings.HasPrefix(host, urlPrefix) {
    44  		if req.Method != "GET" {
    45  			return nil, errors.Errorf("invalid request")
    46  		}
    47  		id := strings.TrimPrefix(host, urlPrefix)
    48  		h.mu.Lock()
    49  		rc, ok := h.requests[id]
    50  		delete(h.requests, id)
    51  		h.mu.Unlock()
    52  
    53  		if !ok {
    54  			return nil, errors.Errorf("context not found")
    55  		}
    56  
    57  		resp := &http.Response{
    58  			Status:        "200 OK",
    59  			StatusCode:    200,
    60  			Body:          rc,
    61  			ContentLength: -1,
    62  		}
    63  
    64  		return resp, nil
    65  	}
    66  	return h.rt.RoundTrip(req)
    67  }