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 }