proxy testutils refactor

clientconn.go

@@ -228,10 +228,8 @@ func DialContext(ctx context.Context, target string, opts ...DialOption) (conn *
 	//
 	// WithTargetResolutionEnabled in `grpc.Dial` ensures that it preserves
 	// behavior: when default scheme passthrough is used, skip hostname
-	// resolution, when any other scheme like "dns" is used for resolution,
-	// perform resolution on the client as expected.
-	opts = append([]DialOption{withDefaultScheme("passthrough"), WithTargetResolutionEnabled()}, opts...)
-
+	// resolution, when "dns" is used for resolution,
+	// perform resolution on the client.
 	opts = append([]DialOption{withDefaultScheme("passthrough"), WithTargetResolutionEnabled()}, opts...)
 	cc, err := NewClient(target, opts...)
 	if err != nil {

dialoptions.go

@@ -94,8 +94,8 @@ type dialOptions struct {
 	idleTimeout                 time.Duration
 	defaultScheme               string
 	maxCallAttempts             int
-	targetResolutionEnabled     bool // Specifies if client should perform target resolution when proxy is enabled.
-	useProxy                    bool // Specifies if a proxy should be used.
+	targetResolutionEnabled     bool // Specifies if target hostnames should be resolved when proxying is enabled.
+	useProxy                    bool // Specifies if a server should be connected via proxy.
 }
 
 // DialOption configures how we set up the connection.
@@ -384,7 +384,8 @@ func WithNoProxy() DialOption {
 }
 
 // WithTargetResolutionEnabled returns a DialOption which enables target
-// resolution on client. This is ignored if WithNoProxy is used.
+// resolution on client even when "dns" scheme is used. This is ignored if
+// WithNoProxy is used.
 //
 // # Experimental
 //

internal/transport/http2_client.go

@@ -157,8 +157,8 @@ type http2Client struct {
 func dial(ctx context.Context, fn func(context.Context, string) (net.Conn, error), addr resolver.Address, grpcUA string) (net.Conn, error) {
 	address := addr.Addr
 
-	if _, present := proxyattributes.Get(addr); present {
-		return proxyDial(ctx, addr, grpcUA)
+	if opts, present := proxyattributes.Get(addr); present {
+		return proxyDial(ctx, addr, grpcUA, opts)
 	}
 	networkType, ok := networktype.Get(addr)
 	if fn != nil {

internal/transport/proxy.go

@@ -54,14 +54,13 @@ func basicAuth(username, password string) string {
 	return base64.StdEncoding.EncodeToString([]byte(auth))
 }
 
-func doHTTPConnectHandshake(ctx context.Context, conn net.Conn, addr resolver.Address, grpcUA string) (_ net.Conn, err error) {
+func doHTTPConnectHandshake(ctx context.Context, conn net.Conn, addr resolver.Address, grpcUA string, opts proxyattributes.Options) (_ net.Conn, err error) {
 	defer func() {
 		if err != nil {
 			conn.Close()
 		}
 	}()
 
-	opts, _ := proxyattributes.Get(addr)
 	req := &http.Request{
 		Method: http.MethodConnect,
 		URL:    &url.URL{Host: opts.ConnectAddr},
@@ -104,12 +103,12 @@ func doHTTPConnectHandshake(ctx context.Context, conn net.Conn, addr resolver.Ad
 }
 
 // proxyDial establishes a TCP connection to the specified address and performs an HTTP CONNECT handshake.
-func proxyDial(ctx context.Context, addr resolver.Address, grpcUA string) (net.Conn, error) {
+func proxyDial(ctx context.Context, addr resolver.Address, grpcUA string, opts proxyattributes.Options) (net.Conn, error) {
 	conn, err := internal.NetDialerWithTCPKeepalive().DialContext(ctx, "tcp", addr.Addr)
 	if err != nil {
 		return nil, err
 	}
-	return doHTTPConnectHandshake(ctx, conn, addr, grpcUA)
+	return doHTTPConnectHandshake(ctx, conn, addr, grpcUA, opts)
 }
 
 func sendHTTPRequest(ctx context.Context, req *http.Request, conn net.Conn) error {

test/proxy_test.go

@@ -19,21 +19,24 @@
 package test
 
 import (
+	"bufio"
+	"bytes"
 	"context"
 	"encoding/base64"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"net/url"
 	"os"
 	"testing"
+	"time"
 
 	"golang.org/x/net/http/httpproxy"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/internal/resolver/delegatingresolver"
 	"google.golang.org/grpc/internal/stubserver"
-	"google.golang.org/grpc/internal/testutils"
 	testgrpc "google.golang.org/grpc/interop/grpc_testing"
 	"google.golang.org/grpc/resolver"
 	"google.golang.org/grpc/resolver/manual"
@@ -70,6 +73,86 @@ func setupDNS(t *testing.T) *manual.Resolver {
 	return mr
 }
 
+// proxyServer represents a test proxy server.
+type proxyServer struct {
+	lis          net.Listener
+	in           net.Conn                  // Connection from the client to the proxy.
+	out          net.Conn                  // Connection from the proxy to the backend.
+	requestCheck func(*http.Request) error // Function to check the request sent to proxy.
+}
+
+// Stop closes the ProxyServer and its connections to client and server.
+func (p *proxyServer) stop() {
+	p.lis.Close()
+	if p.in != nil {
+		p.in.Close()
+	}
+	if p.out != nil {
+		p.out.Close()
+	}
+}
+
+// Creates and starts a proxy server.
+func newProxyServer(lis net.Listener, reqCheck func(*http.Request) error, errCh chan error, doneCh chan struct{}, backendAddr string, resOnClient bool, proxyStarted func()) *proxyServer {
+	p := &proxyServer{
+		lis:          lis,
+		requestCheck: reqCheck,
+	}
+
+	// Start the proxy server.
+	go func() {
+		in, err := p.lis.Accept()
+		if err != nil {
+			return
+		}
+		p.in = in
+		// This will be used in tests to check if the proxy server is started.
+		if proxyStarted != nil {
+			proxyStarted()
+		}
+		req, err := http.ReadRequest(bufio.NewReader(in))
+		if err != nil {
+			errCh <- fmt.Errorf("failed to read CONNECT req: %v", err)
+			return
+		}
+		if err := p.requestCheck(req); err != nil {
+			resp := http.Response{StatusCode: http.StatusMethodNotAllowed}
+			resp.Write(p.in)
+			p.in.Close()
+			errCh <- fmt.Errorf("get wrong CONNECT req: %+v, error: %v", req, err)
+			return
+		}
+		var out net.Conn
+		// If resolution is done on client,connect to address received in
+		// CONNECT request or else connect to backend address directly. This is
+		// to mimick the name resolution on proxy server.
+		if resOnClient {
+			out, err = net.Dial("tcp", req.URL.Host)
+		} else {
+			out, err = net.Dial("tcp", backendAddr)
+		}
+		if err != nil {
+			errCh <- fmt.Errorf("failed to dial to server: %v", err)
+			return
+		}
+		out.SetDeadline(time.Now().Add(defaultTestTimeout))
+
+		// Response OK to client
+		resp := http.Response{StatusCode: http.StatusOK, Proto: "HTTP/1.0"}
+		var buf bytes.Buffer
+		resp.Write(&buf)
+		p.in.Write(buf.Bytes())
+		p.out = out
+
+		// Perform the proxy function, i.e pass the data from client to server
+		// and server to client.
+		go io.Copy(p.in, p.out)
+		go io.Copy(p.out, p.in)
+		close(doneCh)
+	}()
+	return p
+}
+
 // setupProxy initializes and starts a proxy server, registers a cleanup to
 // stop it, and returns the proxy's listener and helper channels.
 func setupProxy(t *testing.T, backendAddr string, resolutionOnClient bool, reqCheck func(*http.Request) error) (net.Listener, chan error, chan struct{}, chan struct{}) {
@@ -86,8 +169,8 @@ func setupProxy(t *testing.T, backendAddr string, resolutionOnClient bool, reqCh
 		close(errCh)
 	})
 
-	proxyServer := testutils.NewProxyServer(pLis, reqCheck, errCh, doneCh, backendAddr, resolutionOnClient, func() { close(proxyStartedCh) })
-	t.Cleanup(proxyServer.Stop)
+	proxyServer := newProxyServer(pLis, reqCheck, errCh, doneCh, backendAddr, resolutionOnClient, func() { close(proxyStartedCh) })
+	t.Cleanup(proxyServer.stop)
 
 	return pLis, errCh, doneCh, proxyStartedCh
 }