github.com/yandex/pandora@v0.5.32/components/guns/http/connect_test.go (about)

     1  package phttp
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"net/url"
     9  	"strings"
    10  	"testing"
    11  
    12  	"github.com/stretchr/testify/require"
    13  	"github.com/yandex/pandora/core/aggregator/netsample"
    14  	"go.uber.org/zap"
    15  )
    16  
    17  var tunnelHandler = func(t *testing.T, originURL string, compareURI bool) http.Handler {
    18  	u, err := url.Parse(originURL)
    19  	require.NoError(t, err)
    20  	originHost := u.Host
    21  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    22  		if compareURI {
    23  			require.Equal(t, originHost, r.RequestURI)
    24  		}
    25  
    26  		toOrigin, err := net.Dial("tcp", originHost)
    27  		require.NoError(t, err)
    28  
    29  		conn, bufReader, err := w.(http.Hijacker).Hijack()
    30  		require.NoError(t, err)
    31  		require.Equal(t, bufReader.Reader.Buffered(), 0, "Current implementation should not send requested data before got response.")
    32  
    33  		_, err = io.WriteString(conn, "HTTP/1.1 200 Connection established\r\n\r\n")
    34  		require.NoError(t, err)
    35  
    36  		go func() { _, _ = io.Copy(toOrigin, conn) }()
    37  		go func() { _, _ = io.Copy(conn, toOrigin) }()
    38  	})
    39  }
    40  
    41  func TestDo(t *testing.T) {
    42  	tests := []struct {
    43  		name      string
    44  		tunnelSSL bool
    45  	}{
    46  		{
    47  			name:      "HTTP client",
    48  			tunnelSSL: false,
    49  		},
    50  		{
    51  			name:      "HTTPS client",
    52  			tunnelSSL: true,
    53  		},
    54  	}
    55  	for _, tt := range tests {
    56  		t.Run(tt.name, func(t *testing.T) {
    57  			tunnelSSL := tt.tunnelSSL
    58  			origin := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
    59  				rw.WriteHeader(http.StatusOK)
    60  			}))
    61  			defer origin.Close()
    62  
    63  			var proxy *httptest.Server
    64  			if tunnelSSL {
    65  				proxy = httptest.NewTLSServer(tunnelHandler(t, origin.URL, true))
    66  			} else {
    67  				proxy = httptest.NewServer(tunnelHandler(t, origin.URL, true))
    68  			}
    69  			defer proxy.Close()
    70  
    71  			req, err := http.NewRequest("GET", origin.URL, nil)
    72  			require.NoError(t, err)
    73  
    74  			conf := DefaultConnectGunConfig()
    75  			conf.Client.ConnectSSL = tunnelSSL
    76  			scheme := "http://"
    77  			if tunnelSSL {
    78  				scheme = "https://"
    79  			}
    80  			conf.Target = strings.TrimPrefix(proxy.URL, scheme)
    81  
    82  			client := newConnectClient(conf.Client, conf.Target)
    83  
    84  			res, err := client.Do(req)
    85  			require.NoError(t, err)
    86  			require.Equal(t, res.StatusCode, http.StatusOK)
    87  		})
    88  	}
    89  }
    90  
    91  func TestNewConnectGun(t *testing.T) {
    92  	origin := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
    93  		rw.WriteHeader(http.StatusOK)
    94  	}))
    95  	defer origin.Close()
    96  	proxy := httptest.NewServer(tunnelHandler(t, origin.URL, false))
    97  	defer proxy.Close()
    98  
    99  	log := zap.NewNop()
   100  	conf := DefaultConnectGunConfig()
   101  	conf.Target = proxy.Listener.Addr().String()
   102  	connectGun := NewConnectGun(conf, log)
   103  
   104  	results := &netsample.TestAggregator{}
   105  	_ = connectGun.Bind(results, testDeps())
   106  
   107  	connectGun.Shoot(newAmmoURL(t, origin.URL))
   108  
   109  	require.Equal(t, len(results.Samples), 1)
   110  	require.NoError(t, results.Samples[0].Err())
   111  	require.Equal(t, results.Samples[0].ProtoCode(), http.StatusOK)
   112  }