| package main |
| |
| import ( |
| "bytes" |
| "fmt" |
| "io" |
| "net" |
| "runtime" |
| "strings" |
| "testing" |
| "time" |
| |
| "github.com/ishidawataru/sctp" |
| "gotest.tools/v3/skip" |
| ) |
| |
| var testBuf = []byte("Buffalo buffalo Buffalo buffalo buffalo buffalo Buffalo buffalo") |
| var testBufSize = len(testBuf) |
| |
| type EchoServer interface { |
| Run() |
| Close() |
| LocalAddr() net.Addr |
| } |
| |
| type EchoServerOptions struct { |
| TCPHalfClose bool |
| } |
| |
| type StreamEchoServer struct { |
| listener net.Listener |
| testCtx *testing.T |
| opts EchoServerOptions |
| } |
| |
| type UDPEchoServer struct { |
| conn net.PacketConn |
| testCtx *testing.T |
| } |
| |
| func NewEchoServer(t *testing.T, proto, address string, opts EchoServerOptions) EchoServer { |
| var server EchoServer |
| if !strings.HasPrefix(proto, "tcp") && opts.TCPHalfClose { |
| t.Fatalf("TCPHalfClose is not supported for %s", proto) |
| } |
| |
| switch { |
| case strings.HasPrefix(proto, "tcp"): |
| listener, err := net.Listen(proto, address) |
| if err != nil { |
| t.Fatal(err) |
| } |
| server = &StreamEchoServer{listener: listener, testCtx: t, opts: opts} |
| case strings.HasPrefix(proto, "udp"): |
| socket, err := net.ListenPacket(proto, address) |
| if err != nil { |
| t.Fatal(err) |
| } |
| server = &UDPEchoServer{conn: socket, testCtx: t} |
| case strings.HasPrefix(proto, "sctp"): |
| addr, err := sctp.ResolveSCTPAddr(proto, address) |
| if err != nil { |
| t.Fatal(err) |
| } |
| listener, err := sctp.ListenSCTP(proto, addr) |
| if err != nil { |
| t.Fatal(err) |
| } |
| server = &StreamEchoServer{listener: listener, testCtx: t} |
| default: |
| t.Fatalf("unknown protocol: %s", proto) |
| } |
| return server |
| } |
| |
| func (server *StreamEchoServer) Run() { |
| go func() { |
| for { |
| client, err := server.listener.Accept() |
| if err != nil { |
| return |
| } |
| go func(client net.Conn) { |
| if server.opts.TCPHalfClose { |
| data, err := io.ReadAll(client) |
| if err != nil { |
| server.testCtx.Logf("io.ReadAll() failed for the client: %v\n", err.Error()) |
| } |
| if _, err := client.Write(data); err != nil { |
| server.testCtx.Logf("can't echo to the client: %v\n", err.Error()) |
| } |
| client.(*net.TCPConn).CloseWrite() |
| } else { |
| if _, err := io.Copy(client, client); err != nil { |
| server.testCtx.Logf("can't echo to the client: %v\n", err.Error()) |
| } |
| client.Close() |
| } |
| }(client) |
| } |
| }() |
| } |
| |
| func (server *StreamEchoServer) LocalAddr() net.Addr { return server.listener.Addr() } |
| func (server *StreamEchoServer) Close() { server.listener.Close() } |
| |
| func (server *UDPEchoServer) Run() { |
| go func() { |
| readBuf := make([]byte, 1024) |
| for { |
| read, from, err := server.conn.ReadFrom(readBuf) |
| if err != nil { |
| return |
| } |
| for i := 0; i != read; { |
| written, err := server.conn.WriteTo(readBuf[i:read], from) |
| if err != nil { |
| break |
| } |
| i += written |
| } |
| } |
| }() |
| } |
| |
| func (server *UDPEchoServer) LocalAddr() net.Addr { return server.conn.LocalAddr() } |
| func (server *UDPEchoServer) Close() { server.conn.Close() } |
| |
| func testProxyAt(t *testing.T, proto string, proxy Proxy, addr string, halfClose bool) { |
| defer proxy.Close() |
| go proxy.Run() |
| var client net.Conn |
| var err error |
| if strings.HasPrefix(proto, "sctp") { |
| var a *sctp.SCTPAddr |
| a, err = sctp.ResolveSCTPAddr(proto, addr) |
| if err != nil { |
| t.Fatal(err) |
| } |
| client, err = sctp.DialSCTP(proto, nil, a) |
| } else { |
| client, err = net.Dial(proto, addr) |
| } |
| |
| if err != nil { |
| t.Fatalf("Can't connect to the proxy: %v", err) |
| } |
| defer client.Close() |
| client.SetDeadline(time.Now().Add(10 * time.Second)) |
| if _, err = client.Write(testBuf); err != nil { |
| t.Fatal(err) |
| } |
| if halfClose { |
| if proto != "tcp" { |
| t.Fatalf("halfClose is not supported for %s", proto) |
| } |
| client.(*net.TCPConn).CloseWrite() |
| } |
| recvBuf := make([]byte, testBufSize) |
| if _, err = client.Read(recvBuf); err != nil { |
| t.Fatal(err) |
| } |
| if !bytes.Equal(testBuf, recvBuf) { |
| t.Fatal(fmt.Errorf("Expected [%v] but got [%v]", testBuf, recvBuf)) |
| } |
| } |
| |
| func testProxy(t *testing.T, proto string, proxy Proxy, halfClose bool) { |
| testProxyAt(t, proto, proxy, proxy.FrontendAddr().String(), halfClose) |
| } |
| |
| func testTCP4Proxy(t *testing.T, halfClose bool) { |
| backend := NewEchoServer(t, "tcp", "127.0.0.1:0", EchoServerOptions{TCPHalfClose: halfClose}) |
| defer backend.Close() |
| backend.Run() |
| frontendAddr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0} |
| proxy, err := NewProxy(frontendAddr, backend.LocalAddr()) |
| if err != nil { |
| t.Fatal(err) |
| } |
| testProxy(t, "tcp", proxy, halfClose) |
| } |
| |
| func TestTCP4Proxy(t *testing.T) { |
| testTCP4Proxy(t, false) |
| } |
| |
| func TestTCP4ProxyHalfClose(t *testing.T) { |
| testTCP4Proxy(t, true) |
| } |
| |
| func TestTCP6Proxy(t *testing.T) { |
| t.Skip("Need to start CI docker with --ipv6") |
| backend := NewEchoServer(t, "tcp", "[::1]:0", EchoServerOptions{}) |
| defer backend.Close() |
| backend.Run() |
| frontendAddr := &net.TCPAddr{IP: net.IPv6loopback, Port: 0} |
| proxy, err := NewProxy(frontendAddr, backend.LocalAddr()) |
| if err != nil { |
| t.Fatal(err) |
| } |
| testProxy(t, "tcp", proxy, false) |
| } |
| |
| func TestTCPDualStackProxy(t *testing.T) { |
| // If I understand `godoc -src net favoriteAddrFamily` (used by the |
| // net.Listen* functions) correctly this should work, but it doesn't. |
| t.Skip("No support for dual stack yet") |
| backend := NewEchoServer(t, "tcp", "[::1]:0", EchoServerOptions{}) |
| defer backend.Close() |
| backend.Run() |
| frontendAddr := &net.TCPAddr{IP: net.IPv6loopback, Port: 0} |
| proxy, err := NewProxy(frontendAddr, backend.LocalAddr()) |
| if err != nil { |
| t.Fatal(err) |
| } |
| ipv4ProxyAddr := &net.TCPAddr{ |
| IP: net.IPv4(127, 0, 0, 1), |
| Port: proxy.FrontendAddr().(*net.TCPAddr).Port, |
| } |
| testProxyAt(t, "tcp", proxy, ipv4ProxyAddr.String(), false) |
| } |
| |
| func TestUDP4Proxy(t *testing.T) { |
| backend := NewEchoServer(t, "udp", "127.0.0.1:0", EchoServerOptions{}) |
| defer backend.Close() |
| backend.Run() |
| frontendAddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0} |
| proxy, err := NewProxy(frontendAddr, backend.LocalAddr()) |
| if err != nil { |
| t.Fatal(err) |
| } |
| testProxy(t, "udp", proxy, false) |
| } |
| |
| func TestUDP6Proxy(t *testing.T) { |
| t.Skip("Need to start CI docker with --ipv6") |
| backend := NewEchoServer(t, "udp", "[::1]:0", EchoServerOptions{}) |
| defer backend.Close() |
| backend.Run() |
| frontendAddr := &net.UDPAddr{IP: net.IPv6loopback, Port: 0} |
| proxy, err := NewProxy(frontendAddr, backend.LocalAddr()) |
| if err != nil { |
| t.Fatal(err) |
| } |
| testProxy(t, "udp", proxy, false) |
| } |
| |
| func TestUDPWriteError(t *testing.T) { |
| frontendAddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0} |
| // Hopefully, this port will be free: */ |
| backendAddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 25587} |
| proxy, err := NewProxy(frontendAddr, backendAddr) |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer proxy.Close() |
| go proxy.Run() |
| client, err := net.Dial("udp", "127.0.0.1:25587") |
| if err != nil { |
| t.Fatalf("Can't connect to the proxy: %v", err) |
| } |
| defer client.Close() |
| // Make sure the proxy doesn't stop when there is no actual backend: |
| client.Write(testBuf) |
| client.Write(testBuf) |
| backend := NewEchoServer(t, "udp", "127.0.0.1:25587", EchoServerOptions{}) |
| defer backend.Close() |
| backend.Run() |
| client.SetDeadline(time.Now().Add(10 * time.Second)) |
| if _, err = client.Write(testBuf); err != nil { |
| t.Fatal(err) |
| } |
| recvBuf := make([]byte, testBufSize) |
| if _, err = client.Read(recvBuf); err != nil { |
| t.Fatal(err) |
| } |
| if !bytes.Equal(testBuf, recvBuf) { |
| t.Fatal(fmt.Errorf("Expected [%v] but got [%v]", testBuf, recvBuf)) |
| } |
| } |
| |
| func TestSCTP4Proxy(t *testing.T) { |
| skip.If(t, runtime.GOOS == "windows", "sctp is not supported on windows") |
| |
| backend := NewEchoServer(t, "sctp", "127.0.0.1:0", EchoServerOptions{}) |
| defer backend.Close() |
| backend.Run() |
| frontendAddr := &sctp.SCTPAddr{IPAddrs: []net.IPAddr{{IP: net.IPv4(127, 0, 0, 1)}}, Port: 0} |
| proxy, err := NewProxy(frontendAddr, backend.LocalAddr()) |
| if err != nil { |
| t.Fatal(err) |
| } |
| testProxy(t, "sctp", proxy, false) |
| } |
| |
| func TestSCTP6Proxy(t *testing.T) { |
| t.Skip("Need to start CI docker with --ipv6") |
| skip.If(t, runtime.GOOS == "windows", "sctp is not supported on windows") |
| |
| backend := NewEchoServer(t, "sctp", "[::1]:0", EchoServerOptions{}) |
| defer backend.Close() |
| backend.Run() |
| frontendAddr := &sctp.SCTPAddr{IPAddrs: []net.IPAddr{{IP: net.IPv6loopback}}, Port: 0} |
| proxy, err := NewProxy(frontendAddr, backend.LocalAddr()) |
| if err != nil { |
| t.Fatal(err) |
| } |
| testProxy(t, "sctp", proxy, false) |
| } |
