add test case

tests/echo_server.go

@@ -1,7 +1,6 @@
 package tests
 
 import (
-	"bufio"
 	"fmt"
 	"io"
 	"net"
@@ -11,8 +10,8 @@ import (
 	frpNet "github.com/fatedier/frp/utils/net"
 )
 
-func StartEchoServer() {
-	l, err := frpNet.ListenTcp("127.0.0.1", 10701)
+func StartTcpEchoServer() {
+	l, err := frpNet.ListenTcp("127.0.0.1", TEST_TCP_ECHO_PORT)
 	if err != nil {
 		fmt.Printf("echo server listen error: %v\n", err)
 		return
@@ -30,7 +29,7 @@ func StartEchoServer() {
 }
 
 func StartUdpEchoServer() {
-	l, err := frpNet.ListenUDP("127.0.0.1", 10703)
+	l, err := frpNet.ListenUDP("127.0.0.1", TEST_UDP_ECHO_PORT)
 	if err != nil {
 		fmt.Printf("udp echo server listen error: %v\n", err)
 		return
@@ -48,7 +47,7 @@ func StartUdpEchoServer() {
 }
 
 func StartUnixDomainServer() {
-	unixPath := "/tmp/frp_echo_server.sock"
+	unixPath := TEST_UNIX_DOMAIN_ADDR
 	os.Remove(unixPath)
 	syscall.Umask(0)
 	l, err := net.Listen("unix", unixPath)
@@ -69,17 +68,20 @@ func StartUnixDomainServer() {
 }
 
 func echoWorker(c net.Conn) {
-	br := bufio.NewReader(c)
+	buf := make([]byte, 2048)
+
 	for {
-		buf, err := br.ReadString('\n')
-		if err == io.EOF {
-			break
-		}
+		n, err := c.Read(buf)
 		if err != nil {
-			fmt.Printf("echo server read error: %v\n", err)
-			return
+			if err == io.EOF {
+				c.Close()
+				break
+			} else {
+				fmt.Printf("echo server read error: %v\n", err)
+				return
+			}
 		}
 
-		c.Write([]byte(buf + "\n"))
+		c.Write(buf[:n])
 	}
 }

tests/func_test.go

@@ -12,43 +12,67 @@ import (
 	"time"
 
 	frpNet "github.com/fatedier/frp/utils/net"
+	"github.com/stretchr/testify/assert"
 )
 
 var (
-	ECHO_PORT     int64  = 10711
-	UDP_ECHO_PORT int64  = 10712
-	HTTP_PORT     int64  = 10710
-	ECHO_TEST_STR string = "Hello World\n"
-	HTTP_RES_STR  string = "Hello World"
+	TEST_STR                    = "frp is a fast reverse proxy to help you expose a local server behind a NAT or firewall to the internet."
+	TEST_TCP_PORT        int64  = 10701
+	TEST_TCP_FRP_PORT    int64  = 10801
+	TEST_TCP_EC_FRP_PORT int64  = 10901
+	TEST_TCP_ECHO_STR    string = "tcp type:" + TEST_STR
+
+	TEST_UDP_PORT     int64  = 10702
+	TEST_UDP_FRP_PORT int64  = 10802
+	TEST_UDP_ECHO_STR string = "udp type:" + TEST_STR
+
+	TEST_UNIX_DOMAIN_ADDR     string = "/tmp/frp_echo_server.sock"
+	TEST_UNIX_DOMAIN_FRP_PORT int64  = 10803
+	TEST_UNIX_DOMAIN_STR      string = "unix domain type:" + TEST_STR
+
+	TEST_HTTP_PORT      int64  = 10704
+	TEST_HTTP_FRP_PORT  int64  = 10804
+	TEST_HTTP_WEB01_STR string = "http web01:" + TEST_STR
 )
 
 func init() {
-	go StartEchoServer()
+	go StartTcpEchoServer()
 	go StartUdpEchoServer()
-	go StartHttpServer()
 	go StartUnixDomainServer()
+	go StartHttpServer()
 	time.Sleep(500 * time.Millisecond)
 }
 
-func TestEchoServer(t *testing.T) {
-	c, err := frpNet.ConnectTcpServer(fmt.Sprintf("127.0.0.1:%d", ECHO_PORT))
-	if err != nil {
-		t.Fatalf("connect to echo server error: %v", err)
-	}
-	timer := time.Now().Add(time.Duration(5) * time.Second)
-	c.SetDeadline(timer)
+func TestTcpServer(t *testing.T) {
+	assert := assert.New(t)
+	// Normal
+	addr := fmt.Sprintf("127.0.0.1:%d", TEST_TCP_FRP_PORT)
+	res, err := sendTcpMsg(addr, TEST_TCP_ECHO_STR)
+	assert.NoError(err)
+	assert.Equal(TEST_TCP_ECHO_STR, res)
 
-	c.Write([]byte(ECHO_TEST_STR + "\n"))
+	// Encrytion and compression
+	addr = fmt.Sprintf("127.0.0.1:%d", TEST_TCP_EC_FRP_PORT)
+	res, err = sendTcpMsg(addr, TEST_TCP_ECHO_STR)
+	assert.NoError(err)
+	assert.Equal(TEST_TCP_ECHO_STR, res)
+}
 
-	br := bufio.NewReader(c)
-	buf, err := br.ReadString('\n')
-	if err != nil {
-		t.Fatalf("read from echo server error: %v", err)
-	}
+func TestUdpEchoServer(t *testing.T) {
+	assert := assert.New(t)
+	// Normal
+	addr := fmt.Sprintf("127.0.0.1:%d", TEST_UDP_FRP_PORT)
+	res, err := sendUdpMsg(addr, TEST_UDP_ECHO_STR)
+	assert.NoError(err)
+	assert.Equal(TEST_UDP_ECHO_STR, res)
 
-	if ECHO_TEST_STR != buf {
-		t.Fatalf("content error, send [%s], get [%s]", strings.Trim(ECHO_TEST_STR, "\n"), strings.Trim(buf, "\n"))
-	}
+func TestUnixDomainServer(t *testing.T) {
+	assert := assert.New(t)
+	// Normal
+	addr := fmt.Sprintf("127.0.0.1:%d", TEST_UNIX_DOMAIN_FRP_PORT)
+	res, err := sendTcpMsg(addr, TEST_UNIX_DOMAIN_STR)
+	assert.NoError(err)
+	assert.Equal(TEST_UNIX_DOMAIN_STR, res)
 }
 
 func TestHttpServer(t *testing.T) {
@@ -71,49 +95,3 @@ func TestHttpServer(t *testing.T) {
 		t.Fatalf("http code from http server error [%d]", res.StatusCode)
 	}
 }
-
-func TestUdpEchoServer(t *testing.T) {
-	addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:10712")
-	if err != nil {
-		t.Fatalf("do udp request error: %v", err)
-	}
-	conn, err := net.DialUDP("udp", nil, addr)
-	if err != nil {
-		t.Fatalf("dial udp server error: %v", err)
-	}
-	defer conn.Close()
-	_, err = conn.Write([]byte("hello frp\n"))
-	if err != nil {
-		t.Fatalf("write to udp server error: %v", err)
-	}
-	data := make([]byte, 20)
-	n, err := conn.Read(data)
-	if err != nil {
-		t.Fatalf("read from udp server error: %v", err)
-	}
-
-	if string(bytes.TrimSpace(data[:n])) != "hello frp" {
-		t.Fatalf("message got from udp server error, get %s", string(data[:n-1]))
-	}
-}
-
-func TestUnixDomainServer(t *testing.T) {
-	c, err := frpNet.ConnectTcpServer(fmt.Sprintf("127.0.0.1:%d", 10704))
-	if err != nil {
-		t.Fatalf("connect to echo server error: %v", err)
-	}
-	timer := time.Now().Add(time.Duration(5) * time.Second)
-	c.SetDeadline(timer)
-
-	c.Write([]byte(ECHO_TEST_STR + "\n"))
-
-	br := bufio.NewReader(c)
-	buf, err := br.ReadString('\n')
-	if err != nil {
-		t.Fatalf("read from echo server error: %v", err)
-	}
-
-	if ECHO_TEST_STR != buf {
-		t.Fatalf("content error, send [%s], get [%s]", strings.Trim(ECHO_TEST_STR, "\n"), strings.Trim(buf, "\n"))
-	}
-}

tests/util.go

@@ -0,0 +1,57 @@
+package test
+
+import (
+	"fmt"
+	"net"
+	"time"
+
+	frpNet "github.com/fatedier/frp/utils/net"
+)
+
+func sendTcpMsg(addr string, msg string) (res string, err error) {
+	c, err := frpNet.ConnectTcpServer(addr)
+	defer c.Close()
+	if err != nil {
+		err = fmt.Errorf("connect to tcp server error: %v", err)
+		return
+	}
+
+	timer := time.Now().Add(5 * time.Second)
+	c.SetDeadline(timer)
+	c.Write([]byte(msg))
+
+	buf := make([]byte, 2048)
+	n, errRet := c.Read(buf)
+	if errRet != nil {
+		err = fmt.Errorf("read from tcp server error: %v", errRet)
+		return
+	}
+	return string(buf[:n]), nil
+}
+
+func sendUdpMsg(addr string, msg string) (res string, err error) {
+	udpAddr, errRet := net.ResolveUDPAddr("udp", addr)
+	if errRet != nil {
+		err = fmt.Errorf("resolve udp addr error: %v", err)
+		return
+	}
+	conn, errRet := net.DialUDP("udp", nil, udpAddr)
+	if errRet != nil {
+		err = fmt.Errorf("dial udp server error: %v", err)
+		return
+	}
+	defer conn.Close()
+	_, err = conn.Write([]byte(msg))
+	if err != nil {
+		err = fmt.Errorf("write to udp server error: %v", err)
+		return
+	}
+
+	buf := make([]byte, 2048)
+	n, errRet := conn.Read(buf)
+	if errRet != nil {
+		err = fmt.Errorf("read from udp server error: %v", err)
+		return
+	}
+	return string(buf[:n]), nil
+}