github.com/zhuohuang-hust/src-cbuild@v0.0.0-20230105071821-c7aab3e7c840/mergeCode/libnetwork/cmd/proxy/network_proxy_test.go (about)

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"flag"
     6  	"fmt"
     7  	"io"
     8  	"net"
     9  	"strings"
    10  	"testing"
    11  	"time"
    12  )
    13  
    14  var _ = flag.Bool("incontainer", false, "Indicates if the test is running in a container")
    15  
    16  var testBuf = []byte("Buffalo buffalo Buffalo buffalo buffalo buffalo Buffalo buffalo")
    17  var testBufSize = len(testBuf)
    18  
    19  type EchoServer interface {
    20  	Run()
    21  	Close()
    22  	LocalAddr() net.Addr
    23  }
    24  
    25  type TCPEchoServer struct {
    26  	listener net.Listener
    27  	testCtx  *testing.T
    28  }
    29  
    30  type UDPEchoServer struct {
    31  	conn    net.PacketConn
    32  	testCtx *testing.T
    33  }
    34  
    35  func NewEchoServer(t *testing.T, proto, address string) EchoServer {
    36  	var server EchoServer
    37  	if strings.HasPrefix(proto, "tcp") {
    38  		listener, err := net.Listen(proto, address)
    39  		if err != nil {
    40  			t.Fatal(err)
    41  		}
    42  		server = &TCPEchoServer{listener: listener, testCtx: t}
    43  	} else {
    44  		socket, err := net.ListenPacket(proto, address)
    45  		if err != nil {
    46  			t.Fatal(err)
    47  		}
    48  		server = &UDPEchoServer{conn: socket, testCtx: t}
    49  	}
    50  	return server
    51  }
    52  
    53  func (server *TCPEchoServer) Run() {
    54  	go func() {
    55  		for {
    56  			client, err := server.listener.Accept()
    57  			if err != nil {
    58  				return
    59  			}
    60  			go func(client net.Conn) {
    61  				if _, err := io.Copy(client, client); err != nil {
    62  					server.testCtx.Logf("can't echo to the client: %v\n", err.Error())
    63  				}
    64  				client.Close()
    65  			}(client)
    66  		}
    67  	}()
    68  }
    69  
    70  func (server *TCPEchoServer) LocalAddr() net.Addr { return server.listener.Addr() }
    71  func (server *TCPEchoServer) Close()              { server.listener.Close() }
    72  
    73  func (server *UDPEchoServer) Run() {
    74  	go func() {
    75  		readBuf := make([]byte, 1024)
    76  		for {
    77  			read, from, err := server.conn.ReadFrom(readBuf)
    78  			if err != nil {
    79  				return
    80  			}
    81  			for i := 0; i != read; {
    82  				written, err := server.conn.WriteTo(readBuf[i:read], from)
    83  				if err != nil {
    84  					break
    85  				}
    86  				i += written
    87  			}
    88  		}
    89  	}()
    90  }
    91  
    92  func (server *UDPEchoServer) LocalAddr() net.Addr { return server.conn.LocalAddr() }
    93  func (server *UDPEchoServer) Close()              { server.conn.Close() }
    94  
    95  func testProxyAt(t *testing.T, proto string, proxy Proxy, addr string) {
    96  	defer proxy.Close()
    97  	go proxy.Run()
    98  	client, err := net.Dial(proto, addr)
    99  	if err != nil {
   100  		t.Fatalf("Can't connect to the proxy: %v", err)
   101  	}
   102  	defer client.Close()
   103  	client.SetDeadline(time.Now().Add(10 * time.Second))
   104  	if _, err = client.Write(testBuf); err != nil {
   105  		t.Fatal(err)
   106  	}
   107  	recvBuf := make([]byte, testBufSize)
   108  	if _, err = client.Read(recvBuf); err != nil {
   109  		t.Fatal(err)
   110  	}
   111  	if !bytes.Equal(testBuf, recvBuf) {
   112  		t.Fatal(fmt.Errorf("Expected [%v] but got [%v]", testBuf, recvBuf))
   113  	}
   114  }
   115  
   116  func testProxy(t *testing.T, proto string, proxy Proxy) {
   117  	testProxyAt(t, proto, proxy, proxy.FrontendAddr().String())
   118  }
   119  
   120  func TestTCP4Proxy(t *testing.T) {
   121  	backend := NewEchoServer(t, "tcp", "127.0.0.1:0")
   122  	defer backend.Close()
   123  	backend.Run()
   124  	frontendAddr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
   125  	proxy, err := NewProxy(frontendAddr, backend.LocalAddr())
   126  	if err != nil {
   127  		t.Fatal(err)
   128  	}
   129  	testProxy(t, "tcp", proxy)
   130  }
   131  
   132  func TestTCP6Proxy(t *testing.T) {
   133  	backend := NewEchoServer(t, "tcp", "[::1]:0")
   134  	defer backend.Close()
   135  	backend.Run()
   136  	frontendAddr := &net.TCPAddr{IP: net.IPv6loopback, Port: 0}
   137  	proxy, err := NewProxy(frontendAddr, backend.LocalAddr())
   138  	if err != nil {
   139  		t.Fatal(err)
   140  	}
   141  	testProxy(t, "tcp", proxy)
   142  }
   143  
   144  func TestTCPDualStackProxy(t *testing.T) {
   145  	// If I understand `godoc -src net favoriteAddrFamily` (used by the
   146  	// net.Listen* functions) correctly this should work, but it doesn't.
   147  	t.Skip("No support for dual stack yet")
   148  	backend := NewEchoServer(t, "tcp", "[::1]:0")
   149  	defer backend.Close()
   150  	backend.Run()
   151  	frontendAddr := &net.TCPAddr{IP: net.IPv6loopback, Port: 0}
   152  	proxy, err := NewProxy(frontendAddr, backend.LocalAddr())
   153  	if err != nil {
   154  		t.Fatal(err)
   155  	}
   156  	ipv4ProxyAddr := &net.TCPAddr{
   157  		IP:   net.IPv4(127, 0, 0, 1),
   158  		Port: proxy.FrontendAddr().(*net.TCPAddr).Port,
   159  	}
   160  	testProxyAt(t, "tcp", proxy, ipv4ProxyAddr.String())
   161  }
   162  
   163  func TestUDP4Proxy(t *testing.T) {
   164  	backend := NewEchoServer(t, "udp", "127.0.0.1:0")
   165  	defer backend.Close()
   166  	backend.Run()
   167  	frontendAddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
   168  	proxy, err := NewProxy(frontendAddr, backend.LocalAddr())
   169  	if err != nil {
   170  		t.Fatal(err)
   171  	}
   172  	testProxy(t, "udp", proxy)
   173  }
   174  
   175  func TestUDP6Proxy(t *testing.T) {
   176  	backend := NewEchoServer(t, "udp", "[::1]:0")
   177  	defer backend.Close()
   178  	backend.Run()
   179  	frontendAddr := &net.UDPAddr{IP: net.IPv6loopback, Port: 0}
   180  	proxy, err := NewProxy(frontendAddr, backend.LocalAddr())
   181  	if err != nil {
   182  		t.Fatal(err)
   183  	}
   184  	testProxy(t, "udp", proxy)
   185  }
   186  
   187  func TestUDPWriteError(t *testing.T) {
   188  	frontendAddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
   189  	// Hopefully, this port will be free: */
   190  	backendAddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 25587}
   191  	proxy, err := NewProxy(frontendAddr, backendAddr)
   192  	if err != nil {
   193  		t.Fatal(err)
   194  	}
   195  	defer proxy.Close()
   196  	go proxy.Run()
   197  	client, err := net.Dial("udp", "127.0.0.1:25587")
   198  	if err != nil {
   199  		t.Fatalf("Can't connect to the proxy: %v", err)
   200  	}
   201  	defer client.Close()
   202  	// Make sure the proxy doesn't stop when there is no actual backend:
   203  	client.Write(testBuf)
   204  	client.Write(testBuf)
   205  	backend := NewEchoServer(t, "udp", "127.0.0.1:25587")
   206  	defer backend.Close()
   207  	backend.Run()
   208  	client.SetDeadline(time.Now().Add(10 * time.Second))
   209  	if _, err = client.Write(testBuf); err != nil {
   210  		t.Fatal(err)
   211  	}
   212  	recvBuf := make([]byte, testBufSize)
   213  	if _, err = client.Read(recvBuf); err != nil {
   214  		t.Fatal(err)
   215  	}
   216  	if !bytes.Equal(testBuf, recvBuf) {
   217  		t.Fatal(fmt.Errorf("Expected [%v] but got [%v]", testBuf, recvBuf))
   218  	}
   219  }