test: add tcpforward test

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...'

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()
 }

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)
+}

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 {

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