google.golang.org/grpc@v1.72.2/internal/testutils/proxyserver/proxyserver.go (about)

     1  /*
     2   *
     3   * Copyright 2024 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  // Package proxyserver provides an implementation of a proxy server for testing purposes.
    20  // The server supports only a single incoming connection at a time and is not concurrent.
    21  // It handles only HTTP CONNECT requests; other HTTP methods are not supported.
    22  package proxyserver
    23  
    24  import (
    25  	"bufio"
    26  	"bytes"
    27  	"io"
    28  	"net"
    29  	"net/http"
    30  	"testing"
    31  	"time"
    32  
    33  	"google.golang.org/grpc/internal/testutils"
    34  )
    35  
    36  // ProxyServer represents a test proxy server.
    37  type ProxyServer struct {
    38  	lis       net.Listener
    39  	in        net.Conn            // Connection from the client to the proxy.
    40  	out       net.Conn            // Connection from the proxy to the backend.
    41  	onRequest func(*http.Request) // Function to check the request sent to proxy.
    42  	Addr      string              // Address of the proxy
    43  }
    44  
    45  const defaultTestTimeout = 10 * time.Second
    46  
    47  // Stop closes the ProxyServer and its connections to client and server.
    48  func (p *ProxyServer) stop() {
    49  	p.lis.Close()
    50  	if p.in != nil {
    51  		p.in.Close()
    52  	}
    53  	if p.out != nil {
    54  		p.out.Close()
    55  	}
    56  }
    57  
    58  func (p *ProxyServer) handleRequest(t *testing.T, in net.Conn, waitForServerHello bool) {
    59  	req, err := http.ReadRequest(bufio.NewReader(in))
    60  	if err != nil {
    61  		t.Errorf("failed to read CONNECT req: %v", err)
    62  		return
    63  	}
    64  	if req.Method != http.MethodConnect {
    65  		t.Errorf("unexpected Method %q, want %q", req.Method, http.MethodConnect)
    66  	}
    67  	p.onRequest(req)
    68  
    69  	t.Logf("Dialing to %s", req.URL.Host)
    70  	out, err := net.Dial("tcp", req.URL.Host)
    71  	if err != nil {
    72  		in.Close()
    73  		t.Logf("failed to dial to server: %v", err)
    74  		return
    75  	}
    76  	out.SetDeadline(time.Now().Add(defaultTestTimeout))
    77  	resp := http.Response{StatusCode: http.StatusOK, Proto: "HTTP/1.0"}
    78  	var buf bytes.Buffer
    79  	resp.Write(&buf)
    80  
    81  	if waitForServerHello {
    82  		// Batch the first message from the server with the http connect
    83  		// response. This is done to test the cases in which the grpc client has
    84  		// the response to the connect request and proxied packets from the
    85  		// destination server when it reads the transport.
    86  		b := make([]byte, 50)
    87  		bytesRead, err := out.Read(b)
    88  		if err != nil {
    89  			t.Errorf("Got error while reading server hello: %v", err)
    90  			in.Close()
    91  			out.Close()
    92  			return
    93  		}
    94  		buf.Write(b[0:bytesRead])
    95  	}
    96  	p.in = in
    97  	p.in.Write(buf.Bytes())
    98  	p.out = out
    99  
   100  	go io.Copy(p.in, p.out)
   101  	go io.Copy(p.out, p.in)
   102  }
   103  
   104  // New initializes and starts a proxy server, registers a cleanup to
   105  // stop it, and returns a ProxyServer.
   106  func New(t *testing.T, reqCheck func(*http.Request), waitForServerHello bool) *ProxyServer {
   107  	t.Helper()
   108  	pLis, err := testutils.LocalTCPListener()
   109  	if err != nil {
   110  		t.Fatalf("failed to listen: %v", err)
   111  	}
   112  
   113  	p := &ProxyServer{
   114  		lis:       pLis,
   115  		onRequest: reqCheck,
   116  		Addr:      pLis.Addr().String(),
   117  	}
   118  
   119  	// Start the proxy server.
   120  	go func() {
   121  		for {
   122  			in, err := p.lis.Accept()
   123  			if err != nil {
   124  				return
   125  			}
   126  			// p.handleRequest is not invoked in a goroutine because the test
   127  			// proxy currently supports handling only one connection at a time.
   128  			p.handleRequest(t, in, waitForServerHello)
   129  		}
   130  	}()
   131  	t.Logf("Started proxy at: %q", pLis.Addr().String())
   132  	t.Cleanup(p.stop)
   133  	return p
   134  }