Merge pull request #12 from piaodazhu/dev

improve test coverage && add benchmark

TODO

@@ -1,7 +1,7 @@
 - [x] Hook API
 - [x] Unit Tests
 - [x] Code Report
-- [ ] Better Coverage
-- [ ] Benchmark
+- [x] Better Coverage
+- [x] Benchmark
 - [x] Register arguments: maxConn, maxServeNum, maxServeTime
 - [x] Active user status

client.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"sort"
 	"sync"
@@ -126,7 +127,12 @@ func register(conn net.Conn, info RegisterInfo, ctrl ControlInfo) error {
 
 // SetLogger Set customized logrus logger for the inner client.
 func (c *ProxyLiteClient) SetLogger(logger *log.Logger) {
-	c.logger = logger
+	if logger == nil {
+		c.logger = log.New()
+		c.logger.SetOutput(io.Discard)
+	} else {
+		c.logger = logger
+	}
 }
 
 // RegisterInnerService Register inner server to proxy server's outer port.

context.go

@@ -28,7 +28,7 @@ func makeContext(tn *tunnel, user *net.Conn, data []byte, kvMap *sync.Map) *Cont
 }
 
 func (ctx *Context) AbortTunnel() error {
-	if ctx.tn.innerConn == nil {
+	if ctx.tn == nil || ctx.tn.innerConn == nil {
 		return errors.New("cannot abort service because inner connection not exists")
 	}
 	return (*ctx.tn.innerConn).Close()

context_test.go

@@ -0,0 +1,37 @@
+package proxylite
+
+import "testing"
+
+func TestContextNil(t *testing.T) {
+	data := []byte{'a'}
+	ctx := makeContext(nil, nil, data, nil)
+	if ctx.AbortTunnel() == nil {
+		t.Error("AbortTunnel")
+	}
+	if ctx.AbortUser() == nil {
+		t.Error("AbortUser")
+	}
+	if ctx.DataBuffer() == nil {
+		t.Error("DataBuffer")
+	}
+
+	blankInfo := ServiceInfo{}
+	if ctx.ServiceInfo() != blankInfo {
+		t.Error("ServiceInfo")
+	}
+	if ctx.UserLocalAddress() != nil {
+		t.Error("UserLocalAddress")
+	}
+	if ctx.UserRemoteAddress() != nil {
+		t.Error("UserRemoteAddress")
+	}
+	if ctx.InnerLocalConn() != nil {
+		t.Error("InnerLocalConn")
+	}
+	if ctx.InnerRemoteConn() != nil {
+		t.Error("InnerRemoteConn")
+	}
+	if _, ok := ctx.GetValue("k"); ok {
+		t.Error("GetValue")
+	}
+}

proxylite_bench_test.go

@@ -0,0 +1,95 @@
+package proxylite
+
+import (
+	"net"
+	"strings"
+	"testing"
+	"time"
+)
+
+var msg = strings.Repeat("hello", 100)
+
+func init() {
+	if err := newTcpEchoServer(":9969", len(msg)); err != nil {
+		panic(err)
+	}
+}
+
+func BenchmarkProxylite(b *testing.B) {
+	proxyServer := NewProxyLiteServer()
+	proxyServer.SetLogger(nil)
+	proxyServer.AddPort(9968, 9968)
+	go func() {
+		if err := proxyServer.Run(":9967"); err != nil {
+			panic(err)
+		}
+	}()
+	time.Sleep(time.Millisecond * 10)
+
+	innerClient := NewProxyLiteClient(":9967")
+	innerClient.SetLogger(nil)
+
+	cancelFunc, done, err := innerClient.RegisterInnerService(
+		RegisterInfo{
+			OuterPort: 9968,
+			InnerAddr: ":9969",
+			Name:      "Echo",
+			Message:   "TCP Echo Server",
+		},
+		ControlInfo{},
+	)
+	if err != nil {
+		panic(err)
+	}
+	defer func() {
+		cancelFunc()
+		<-done
+		proxyServer.Stop()
+	}()
+	time.Sleep(time.Millisecond * 10)
+
+	user, err := net.Dial("tcp", ":9968")
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		err := written(user, []byte(msg), len(msg))
+		if err != nil {
+			b.Error("write 1, ", err)
+		}
+		_, err = readn(user, len(msg))
+		if err != nil {
+			b.Error("read 1, ", err)
+		}
+	}
+	b.StopTimer()
+	user.Close()
+	time.Sleep(time.Millisecond * 100)
+	select {
+	case <-done:
+		b.Error("unexpected quit")
+	default:
+	}
+}
+
+func BenchmarkNoProxy(b *testing.B) {
+	user, err := net.Dial("tcp", ":9969")
+	if err != nil {
+		b.Fail()
+	}
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		err := written(user, []byte(msg), len(msg))
+		if err != nil {
+			b.Error("write 1, ", err)
+		}
+		_, err = readn(user, len(msg))
+		if err != nil {
+			b.Error("read 1, ", err)
+		}
+	}
+	b.StopTimer()
+	user.Close()
+}