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

     1  package connection
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"encoding/json"
     7  	"fmt"
     8  	"io"
     9  	"io/ioutil"
    10  	"net"
    11  	"net/http"
    12  	"net/http/httputil"
    13  	"net/url"
    14  	"time"
    15  
    16  	"code.cloudfoundry.org/garden"
    17  	"code.cloudfoundry.org/garden/routes"
    18  	"github.com/tedsuo/rata"
    19  )
    20  
    21  // IMPORTANT NOTE: We don't compile this in because we actually use transport.WorkerHijackStreamer
    22  // most of this folder was just copied in from garden/client as a temp workaround to adding ctx to Hijack()
    23  // we needed connection.go in garden/client package and were forced to import the rest of this stuff to make
    24  // connection_suite_test.go to pass
    25  
    26  type DialerFunc func(network, address string) (net.Conn, error)
    27  
    28  type hijackable struct {
    29  	req               *rata.RequestGenerator
    30  	noKeepaliveClient *http.Client
    31  	dialer            DialerFunc
    32  }
    33  
    34  func NewHijackStreamer(network, address string) HijackStreamer {
    35  	return NewHijackStreamerWithDialer(func(string, string) (net.Conn, error) {
    36  		return net.DialTimeout(network, address, 2*time.Second)
    37  	})
    38  }
    39  
    40  func NewHijackStreamerWithDialer(dialFunc DialerFunc) HijackStreamer {
    41  	return &hijackable{
    42  		req:    rata.NewRequestGenerator("http://api", routes.Routes),
    43  		dialer: dialFunc,
    44  		noKeepaliveClient: &http.Client{
    45  			Transport: &http.Transport{
    46  				Dial:              dialFunc,
    47  				DisableKeepAlives: true,
    48  			},
    49  		},
    50  	}
    51  }
    52  
    53  func (h *hijackable) Hijack(ctx context.Context, handler string, body io.Reader, params rata.Params, query url.Values, contentType string) (net.Conn, *bufio.Reader, error) {
    54  	request, err := h.req.CreateRequest(handler, params, body)
    55  	if err != nil {
    56  		return nil, nil, err
    57  	}
    58  
    59  	if contentType != "" {
    60  		request.Header.Set("Content-Type", contentType)
    61  	}
    62  
    63  	if query != nil {
    64  		request.URL.RawQuery = query.Encode()
    65  	}
    66  
    67  	request = request.WithContext(ctx)
    68  
    69  	conn, err := h.dialer("tcp", "api") // net/addr don't matter here
    70  	if err != nil {
    71  		return nil, nil, err
    72  	}
    73  
    74  	client := httputil.NewClientConn(conn, nil)
    75  
    76  	httpResp, err := client.Do(request)
    77  	if err != nil {
    78  		return nil, nil, err
    79  	}
    80  
    81  	if httpResp.StatusCode < 200 || httpResp.StatusCode > 299 {
    82  		defer httpResp.Body.Close()
    83  
    84  		errRespBytes, err := ioutil.ReadAll(httpResp.Body)
    85  		if err != nil {
    86  			return nil, nil, fmt.Errorf("Backend error: Exit status: %d, Body: %s, error reading response body: %s", httpResp.StatusCode, string(errRespBytes), err)
    87  		}
    88  
    89  		var result garden.Error
    90  		err = json.Unmarshal(errRespBytes, &result)
    91  		if err != nil {
    92  			return nil, nil, fmt.Errorf("Backend error: Exit status: %d, Body: %s, error reading response body: %s", httpResp.StatusCode, string(errRespBytes), err)
    93  		}
    94  
    95  		return nil, nil, result.Err
    96  	}
    97  
    98  	hijackedConn, hijackedResponseReader := client.Hijack()
    99  
   100  	return hijackedConn, hijackedResponseReader, nil
   101  }
   102  
   103  func (c *hijackable) Stream(handler string, body io.Reader, params rata.Params, query url.Values, contentType string) (io.ReadCloser, error) {
   104  	request, err := c.req.CreateRequest(handler, params, body)
   105  	if err != nil {
   106  		return nil, err
   107  	}
   108  
   109  	if contentType != "" {
   110  		request.Header.Set("Content-Type", contentType)
   111  	}
   112  
   113  	if query != nil {
   114  		request.URL.RawQuery = query.Encode()
   115  	}
   116  
   117  	httpResp, err := c.noKeepaliveClient.Do(request)
   118  	if err != nil {
   119  		return nil, err
   120  	}
   121  
   122  	if httpResp.StatusCode < 200 || httpResp.StatusCode > 299 {
   123  		defer httpResp.Body.Close()
   124  
   125  		var result garden.Error
   126  		err := json.NewDecoder(httpResp.Body).Decode(&result)
   127  		if err != nil {
   128  			return nil, fmt.Errorf("bad response: %s", err)
   129  		}
   130  
   131  		return nil, result.Err
   132  	}
   133  
   134  	return httpResp.Body, nil
   135  }