From dc6c94a9aa085d400a01a6c0374bd3e58c32e97b Mon Sep 17 00:00:00 2001 From: ICKelin Date: Thu, 6 May 2021 20:32:09 +0800 Subject: [PATCH] test: add tcpforward test --- build_exec.sh | 2 +- opennotrd/core/tcpforward.go | 23 ++--- opennotrd/core/tcpforward_test.go | 143 ++++++++++++++++++++++++++++++ opennotrd/core/udpforward.go | 1 + opennotrd/run.go | 13 ++- 5 files changed, 167 insertions(+), 15 deletions(-) create mode 100644 opennotrd/core/tcpforward_test.go diff --git a/build_exec.sh b/build_exec.sh index f68786b..3c742a1 100755 --- a/build_exec.sh +++ b/build_exec.sh @@ -8,7 +8,7 @@ cd $WORKSPACE/opennotr echo 'building client...' GOOS=darwin go build -o $BIN/$EXEC_PREFIX_darwin_amd64 GOOS=linux go build -o $BIN/$EXEC_PREFIX_linux_amd64 -GOARCH=arm GOOS=linux go build -o $BIN/$EXEC_PREFIX-_arm +GOARCH=arm GOOS=linux go build -o $BIN/$EXEC_PREFIX_arm GOARCH=arm64 GOOS=linux go build -o $BIN/$EXEC_PREFIX_arm64 echo 'building server...' diff --git a/opennotrd/core/tcpforward.go b/opennotrd/core/tcpforward.go index 6a7223a..0432e5c 100644 --- a/opennotrd/core/tcpforward.go +++ b/opennotrd/core/tcpforward.go @@ -20,25 +20,28 @@ func NewTCPForward() *TCPForward { } } -func (f *TCPForward) ListenAndServe(listenAddr string) error { +func (f *TCPForward) Listen(listenAddr string) (net.Listener, error) { listener, err := net.Listen("tcp", listenAddr) if err != nil { - return err + return nil, err } - defer listener.Close() // set socket with ip transparent option file, err := listener.(*net.TCPListener).File() if err != nil { - return err + return nil, err } defer file.Close() err = syscall.SetsockoptInt(int(file.Fd()), syscall.SOL_IP, syscall.IP_TRANSPARENT, 1) if err != nil { - return err + return nil, err } + return listener, nil +} + +func (f *TCPForward) Serve(listener net.Listener) error { for { conn, err := listener.Accept() if err != nil { @@ -53,22 +56,23 @@ func (f *TCPForward) ListenAndServe(listenAddr string) error { } func (f *TCPForward) forwardTCP(conn net.Conn) { + defer conn.Close() + dip, dport, _ := net.SplitHostPort(conn.LocalAddr().String()) sip, sport, _ := net.SplitHostPort(conn.RemoteAddr().String()) sess := f.sessMgr.GetSession(dip) if sess == nil { logs.Error("no route to host: %s", dip) - conn.Close() return } stream, err := sess.conn.OpenStream() if err != nil { logs.Error("open stream fail: %v", err) - conn.Close() return } + defer stream.Close() // todo rewrite to client configuration targetIP := "127.0.0.1" @@ -78,14 +82,13 @@ func (f *TCPForward) forwardTCP(conn net.Conn) { stream.SetWriteDeadline(time.Time{}) if err != nil { logs.Error("stream write fail: %v", err) - conn.Close() - stream.Close() return } wg := &sync.WaitGroup{} wg.Add(1) defer wg.Wait() + go func() { defer wg.Done() defer stream.Close() @@ -99,6 +102,4 @@ func (f *TCPForward) forwardTCP(conn net.Conn) { // and two goroutine 4KB mem used buf := make([]byte, 4096) io.CopyBuffer(conn, stream, buf) - stream.Close() - conn.Close() } diff --git a/opennotrd/core/tcpforward_test.go b/opennotrd/core/tcpforward_test.go new file mode 100644 index 0000000..610170d --- /dev/null +++ b/opennotrd/core/tcpforward_test.go @@ -0,0 +1,143 @@ +package core + +import ( + "fmt" + "io" + "net" + "os" + "testing" + "time" + + "github.com/hashicorp/yamux" +) + +// client -----> tproxy | opennotr server <------ opennotr client + +var backendAddr = "127.0.0.1:8522" +var serverAddr = "127.0.0.1:8521" +var tproxyAddr = "127.0.0.1:8520" +var vip = "100.64.240.10" + +type mockConn struct { + net.Conn + addr mockAddr +} + +type mockAddr struct{} + +func (addr mockAddr) Network() string { + return "tcp" +} + +func (addr mockAddr) String() string { + return "100.64.240.10:8522" +} + +func (c *mockConn) LocalAddr() net.Addr { + return c.addr +} + +func (c *mockConn) Write(buf []byte) (int, error) { + fmt.Printf("receive %d bytes\n", len(buf)) + return len(buf), nil +} + +func runBackend(t *testing.T) { + conn, err := net.Dial("tcp", serverAddr) + if err != nil { + t.Error(err) + return + } + defer conn.Close() + sess, err := yamux.Client(conn, nil) + if err != nil { + t.Error(err) + t.FailNow() + } + defer sess.Close() + + for { + stream, err := sess.AcceptStream() + if err != nil { + t.Error(err) + t.FailNow() + } + + go func() { + defer stream.Close() + }() + } +} + +func runserver(t *testing.T, listener net.Listener) { + for { + conn, err := listener.Accept() + if err != nil { + break + } + + go func() { + sess, err := yamux.Server(conn, nil) + if err != nil { + t.Error(err) + t.FailNow() + } + + sessMgr := GetSessionManager() + sessMgr.AddSession(vip, &Session{conn: sess}) + }() + } +} + +func runtproxy(t *testing.T, tcpfw *TCPForward, listener net.Listener) { + for { + conn, err := listener.Accept() + if err != nil { + break + } + + go func() { + // forward test + mConn := &mockConn{} + mConn.Conn = conn + tcpfw.forwardTCP(mConn) + }() + } +} + +func TestTCPForward(t *testing.T) { + // listen tproxy + tcpfw := NewTCPForward() + listener, err := tcpfw.Listen(tproxyAddr) + if err != nil { + t.Error(err) + return + } + defer listener.Close() + + srvlistener, err := net.Listen("tcp", serverAddr) + if err != nil { + t.Error(err) + return + } + defer srvlistener.Close() + + go runBackend(t) + go runserver(t, srvlistener) + go runtproxy(t, tcpfw, listener) + + conn, err := net.Dial("tcp", tproxyAddr) + if err != nil { + t.FailNow() + } + defer conn.Close() + + go func() { + for i := 0; i < 100; i++ { + time.Sleep(time.Second * 1) + conn.Write([]byte("ping\n")) + } + }() + + io.Copy(os.Stdout, conn) +} diff --git a/opennotrd/core/udpforward.go b/opennotrd/core/udpforward.go index e41a7b8..e42f6f7 100644 --- a/opennotrd/core/udpforward.go +++ b/opennotrd/core/udpforward.go @@ -139,6 +139,7 @@ func (f *UDPForward) ListenAndServe(listenAddr string) error { return nil } +// forwardUDP reads from stream and write to tofd via rawsocket func (f *UDPForward) forwardUDP(stream *yamux.Stream, tofd int, fromaddr, toaddr *net.UDPAddr) { hdr := make([]byte, 2) for { diff --git a/opennotrd/run.go b/opennotrd/run.go index 7443179..83a7f47 100644 --- a/opennotrd/run.go +++ b/opennotrd/run.go @@ -3,7 +3,6 @@ package opennotrd import ( "flag" "fmt" - "log" "github.com/ICKelin/opennotr/opennotrd/core" "github.com/ICKelin/opennotr/opennotrd/plugin" @@ -45,14 +44,22 @@ func Run() { if len(cfg.ResolverConfig.EtcdEndpoints) > 0 { resolver, err = core.NewResolve(cfg.ResolverConfig.EtcdEndpoints) if err != nil { - log.Println(err) + logs.Error("new resolve fail: %v", err) return } } // up local tcp,udp service // we use tproxy to route traffic to the tcp port and udp port here. - go core.NewTCPForward().ListenAndServe(cfg.ServerConfig.TCPForwardListen) + tcpfw := core.NewTCPForward() + listener, err := tcpfw.Listen(cfg.ServerConfig.TCPForwardListen) + if err != nil { + logs.Error("listen tproxy tcp fail: %v", err) + return + } + + go tcpfw.Serve(listener) + go core.NewUDPForward().ListenAndServe(cfg.ServerConfig.UDPForwardListen) // server provides tcp server for opennotr client