github.com/pf-qiu/concourse/v6@v6.7.3-0.20201207032516-1f455d73275f/atc/worker/transport/hijack_streamer.go (about)

     1  package transport
     2  
     3  // WorkerHijackStreamer implements Stream that is using our httpClient,
     4  import (
     5  	"bufio"
     6  	"context"
     7  	"encoding/json"
     8  	"fmt"
     9  	"io"
    10  	"io/ioutil"
    11  	"net"
    12  	"net/http"
    13  	"net/url"
    14  
    15  	"code.cloudfoundry.org/garden"
    16  	"github.com/pf-qiu/concourse/v6/atc/db"
    17  	"github.com/concourse/retryhttp"
    18  	"github.com/tedsuo/rata"
    19  )
    20  
    21  //go:generate counterfeiter . TransportDB
    22  type TransportDB interface {
    23  	GetWorker(name string) (db.Worker, bool, error)
    24  }
    25  
    26  //go:generate counterfeiter . ReadCloser
    27  type ReadCloser interface {
    28  	Read(p []byte) (n int, err error)
    29  	Close() error
    30  }
    31  
    32  //go:generate counterfeiter . RequestGenerator
    33  type RequestGenerator interface {
    34  	CreateRequest(name string, params rata.Params, body io.Reader) (*http.Request, error)
    35  }
    36  
    37  // instead of httpClient defined in default Garden HijackStreamer
    38  type WorkerHijackStreamer struct {
    39  	HttpClient       *http.Client
    40  	HijackableClient retryhttp.HijackableClient
    41  	Req              RequestGenerator
    42  }
    43  
    44  func (h *WorkerHijackStreamer) Stream(handler string, body io.Reader, params rata.Params, query url.Values, contentType string) (io.ReadCloser, error) {
    45  	request, err := h.Req.CreateRequest(handler, params, body)
    46  	if err != nil {
    47  		return nil, err
    48  	}
    49  
    50  	if contentType != "" {
    51  		request.Header.Set("Content-Type", contentType)
    52  	}
    53  
    54  	if query != nil {
    55  		request.URL.RawQuery = query.Encode()
    56  	}
    57  
    58  	httpResp, err := h.HttpClient.Do(request)
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  
    63  	if httpResp.StatusCode < 200 || httpResp.StatusCode > 299 {
    64  		defer httpResp.Body.Close()
    65  
    66  		var result garden.Error
    67  		err := json.NewDecoder(httpResp.Body).Decode(&result)
    68  		if err != nil {
    69  			return nil, fmt.Errorf("bad response: %s", err)
    70  		}
    71  
    72  		return nil, result.Err
    73  	}
    74  
    75  	return httpResp.Body, nil
    76  }
    77  
    78  func (h *WorkerHijackStreamer) Hijack(ctx context.Context, handler string, body io.Reader, params rata.Params, query url.Values, contentType string) (net.Conn, *bufio.Reader, error) {
    79  	request, err := h.Req.CreateRequest(handler, params, body)
    80  	if err != nil {
    81  		return nil, nil, err
    82  	}
    83  
    84  	if contentType != "" {
    85  		request.Header.Set("Content-Type", contentType)
    86  	}
    87  
    88  	if query != nil {
    89  		request.URL.RawQuery = query.Encode()
    90  	}
    91  
    92  	request = request.WithContext(ctx)
    93  
    94  	httpResp, hijackCloser, err := h.HijackableClient.Do(request)
    95  	if err != nil {
    96  		return nil, nil, err
    97  	}
    98  
    99  	if httpResp.StatusCode < 200 || httpResp.StatusCode > 299 {
   100  		defer hijackCloser.Close()
   101  		defer httpResp.Body.Close()
   102  
   103  		errRespBytes, err := ioutil.ReadAll(httpResp.Body)
   104  		if err != nil {
   105  			return nil, nil, fmt.Errorf("Backend error: Exit status: %d, error reading response body: %s", httpResp.StatusCode, err)
   106  		}
   107  		var gerr garden.Error
   108  		if err := gerr.UnmarshalJSON(errRespBytes); err == nil {
   109  			if _, ok := gerr.Err.(garden.ExecutableNotFoundError); ok {
   110  				return nil, nil, gerr.Err
   111  			}
   112  		}
   113  
   114  		return nil, nil, fmt.Errorf("Backend error: Exit status: %d, message: %s", httpResp.StatusCode, errRespBytes)
   115  	}
   116  
   117  	hijackedConn, hijackedResponseReader := hijackCloser.Hijack()
   118  
   119  	return hijackedConn, hijackedResponseReader, nil
   120  }