support hookfunc

README.md

@@ -1,3 +1,7 @@
+[![Go Reference](https://pkg.go.dev/badge/github.com/piaodazhu/proxylite.svg)](https://pkg.go.dev/github.com/piaodazhu/proxylite)
+[![Go Report Card](https://goreportcard.com/badge/github.com/piaodazhu/proxylite)](https://goreportcard.com/report/github.com/piaodazhu/proxylite)
+<!-- [![codecov](https://codecov.io/gh/piaodazhu/proxylite/branch/master/graph/badge.svg?token=KMOWEKDPN5)](https://codecov.io/gh/piaodazhu/proxylite) -->
+
 [English](./README.md)|[中文](./README_ZH.md)
 ## proxylite
 `proxylite` is a dynamic TCP reverse proxy Golang package for NAT or firewall traversal. It is god damn easy to integrate network traversal feature into your project with proxylite. 

README_ZH.md

@@ -1,3 +1,7 @@
+[![Go Reference](https://pkg.go.dev/badge/github.com/piaodazhu/proxylite.svg)](https://pkg.go.dev/github.com/piaodazhu/proxylite)
+[![Go Report Card](https://goreportcard.com/badge/github.com/piaodazhu/proxylite)](https://goreportcard.com/report/github.com/piaodazhu/proxylite)
+<!-- [![codecov](https://codecov.io/gh/piaodazhu/proxylite/branch/master/graph/badge.svg?token=KMOWEKDPN5)](https://codecov.io/gh/piaodazhu/proxylite) -->
+
 [English](./README.md)|[中文](./README_ZH.md)
 
 ---- 以下中文文档由ChatGPT生成 ----

TODO

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

context.go

@@ -0,0 +1,83 @@
+package proxylite
+
+import (
+	"errors"
+	"net"
+	"sync"
+)
+
+type Context struct {
+	tn   *tunnel
+	user net.Conn
+	data []byte
+	kvs *sync.Map
+}
+
+func makeContext(tn *tunnel, user *net.Conn, data []byte, kvMap *sync.Map) *Context {
+	if tn == nil {
+		return nil
+	}
+	ctx := &Context{
+		tn:   tn,
+		data: data,
+		kvs: kvMap,
+	}
+	if user == nil {
+		ctx.user = nil
+	} else {
+		ctx.user = *user
+	}
+	return ctx
+}
+
+func (ctx *Context) AbortTunnel() error {
+	if ctx.tn.innerConn == nil {
+		return errors.New("cannot abort service because inner connection not exists")
+	}
+	return (*ctx.tn.innerConn).Close()
+}
+
+func (ctx *Context) AbortUser() error {
+	if ctx.user == nil {
+		return errors.New("cannot abort user because user connection not exists")
+	}
+	return ctx.user.Close()
+}
+
+func (ctx *Context) ServiceInfo() ServiceInfo {
+	return *ctx.tn.service
+}
+
+func (ctx *Context) UserLocalAddress() net.Addr {
+	return ctx.user.LocalAddr()
+}
+
+func (ctx *Context) UserRemoteAddress() net.Addr {
+	return ctx.user.RemoteAddr()
+}
+
+func (ctx *Context) InnerLocalConn() net.Addr {
+	if ctx.tn.innerConn == nil {
+		return nil
+	}
+	return (*ctx.tn.innerConn).LocalAddr()
+}
+
+func (ctx *Context) InnerRemoteConn() net.Addr {
+	if ctx.tn.innerConn == nil {
+		return nil
+	}
+	return (*ctx.tn.innerConn).RemoteAddr()
+}
+
+func (ctx *Context) DataBuffer() []byte {
+	return ctx.data
+}
+
+func (ctx *Context) PutValue(key, value interface{}) {
+	ctx.kvs.Store(key, value)
+}
+
+func (ctx *Context) GetValue(key interface{}) (interface{}, bool) {
+	return ctx.kvs.Load(key)
+}

go.mod

@@ -1,6 +1,6 @@
 module github.com/piaodazhu/proxylite
 
-go 1.12
+go 1.16
 
 require (
 	github.com/cespare/xxhash/v2 v2.2.0

hook.go

@@ -0,0 +1,9 @@
+package proxylite
+
+type HookFunc func(ctx *Context)
+
+func (hf HookFunc) Trigger(ctx *Context) {
+	if hf != nil && ctx != nil {
+		hf(ctx)
+	}
+}

hook_test.go

@@ -0,0 +1,31 @@
+package proxylite
+
+import (
+	"fmt"
+	"testing"
+)
+
+type student struct {
+	say  HookFunc
+	name string
+}
+
+func newSayHello(msg string) HookFunc {
+	return func(ctx *Context) {
+		words = fmt.Sprintf("hello %s", msg)
+	}
+}
+
+var words string
+
+func TestHookBasic(t *testing.T) {
+	alice := student{
+		name: "alice",
+	}
+	alice.say = newSayHello(alice.name)
+	words = ""
+	alice.say.Trigger(&Context{})
+	if words != fmt.Sprintf("hello %s", alice.name) {
+		t.Error("trigger hook func failed")
+	}
+}

proxylite_test.go

@@ -508,3 +508,102 @@ func TestMultiplexMaxConnControl(t *testing.T) {
 		t.Error("count control error")
 	}
 }
+
+func TestSetHook(t *testing.T) {
+
+	logger := logrus.New()
+	logger.Level = logrus.FatalLevel
+
+	proxyServer := NewProxyLiteServer()
+	proxyServer.SetLogger(logger)
+	proxyServer.AddPort(9968, 9968)
+
+	trace := ""
+	proxyServer.OnTunnelCreated(func(ctx *Context) {
+		ctx.PutValue("key1", "v1")
+		trace += "1"
+	})
+	proxyServer.OnTunnelDestroyed(func(ctx *Context) {
+		if v, ok := ctx.GetValue("key1"); !ok || v.(string) != "v1" {
+			panic("kvs doesn't work")
+		}
+		trace += "2"
+	})
+	proxyServer.OnUserComming(func(ctx *Context) {
+		if v, ok := ctx.GetValue("key1"); !ok || v.(string) != "v1" {
+			panic("kvs doesn't work")
+		}
+		trace += "3"
+	})
+	proxyServer.OnUserLeaving(func(ctx *Context) {
+		if v, ok := ctx.GetValue("key1"); !ok || v.(string) != "v1" {
+			panic("kvs doesn't work")
+		}
+		trace += "4"
+	})
+	proxyServer.OnForwardTunnelToUser(func(ctx *Context) {
+		if v, ok := ctx.GetValue("key1"); !ok || v.(string) != "v1" {
+			panic("kvs doesn't work")
+		}
+		trace += "5"
+	})
+	proxyServer.OnForwardUserToTunnel(func(ctx *Context) {
+		if v, ok := ctx.GetValue("key1"); !ok || v.(string) != "v1" {
+			panic("kvs doesn't work")
+		}
+		trace += "6"
+	})
+
+	go func() {
+		proxyServer.Run(":9967")
+	}()
+	time.Sleep(time.Millisecond * 10)
+
+	innerClient := NewProxyLiteClient(":9967")
+	innerClient.SetLogger(logger)
+	cancelFunc, done, err := innerClient.RegisterInnerService(
+		RegisterInfo{
+			OuterPort: 9968,
+			InnerAddr: ":9966",
+			Name:      "Echo",
+			Message:   "TCP Echo Server",
+		},
+		ControlInfo{},
+	)
+	if err != nil {
+		panic(err)
+	}
+	defer func() {
+		cancelFunc()
+		<-done
+		fmt.Println(trace)
+	}()
+	time.Sleep(time.Millisecond * 10)
+
+	user, err := net.Dial("tcp", ":9968")
+	if err != nil {
+		t.Fatal(err)
+	}
+	msg := "hello123"
+	var data []byte
+	for i := 0; i < 10; i++ {
+		err := written(user, []byte(msg), 8)
+		if err != nil {
+			t.Error("write 1, ", err)
+		}
+		data, err = readn(user, 8)
+		if err != nil {
+			t.Error("read 1, ", err)
+		}
+		if string(data) != msg {
+			t.Error("read 3, ", string(data))
+		}
+	}
+	user.Close()
+	time.Sleep(time.Millisecond * 100)
+	select {
+	case <-done:
+		t.Error("unexpected quit")
+	default:
+	}
+}

server.go

@@ -17,13 +17,10 @@ import (
 )
 
 type tunnel struct {
-	empty bool
-	// busy  bool
-	// birth     time.Time
+	empty     bool
 	innerConn *net.Conn
-	// info      *RegisterInfo
-	service *ServiceInfo
-	ctrl    *ControlInfo
+	service   *ServiceInfo
+	ctrl      *ControlInfo
 }
 
 // ProxyLiteServer Public server that forwards traffic between user and inner client.
@@ -32,6 +29,13 @@ type ProxyLiteServer struct {
 	lock   sync.RWMutex
 	used   map[uint32]*tunnel
 	logger *log.Logger
+
+	onTunnelCreated       HookFunc
+	onTunnelDestroyed     HookFunc
+	onUserComming         HookFunc
+	onUserLeaving         HookFunc
+	onForwardUserToTunnel HookFunc
+	onForwardTunnelToUser HookFunc
 }
 
 // NewProxyLiteServer Create a Proxy server with available ports intervals.
@@ -260,6 +264,10 @@ func (s *ProxyLiteServer) startTunnel(tn *tunnel) {
 	doOnce := sync.Once{}
 	binder := newConnUidBinder(0)
 	var totalOut, totalIn uint64
+	kvs := &sync.Map{}
+
+	s.onTunnelCreated.Trigger(makeContext(tn, nil, nil, kvs))
+	defer s.onTunnelDestroyed.Trigger(makeContext(tn, nil, nil, kvs))
 
 	// Now, register is OK, we want to map outer port to the inner client's socket
 	s.logTunnelMessage(tn.service.Name, "REGISTER", fmt.Sprintf("New inner client register port %d ok. Listening for outer client...", tn.service.Port))
@@ -350,6 +358,7 @@ func (s *ProxyLiteServer) startTunnel(tn *tunnel) {
 
 			// forward to the write user. (don't send 4 byte uid)
 			// here can be optimized by goroutines
+			s.onForwardTunnelToUser.Trigger(makeContext(tn, outerConn, data[8:], kvs))
 			n, err = (*outerConn).Write(data[8:])
 			if err != nil {
 				// how to send close thougth tunnel?
@@ -370,6 +379,9 @@ func (s *ProxyLiteServer) startTunnel(tn *tunnel) {
 		var err error
 		var userEnd bool = false
 
+		s.onUserComming.Trigger(makeContext(tn, &outerConn, nil, kvs))
+		defer s.onUserLeaving.Trigger(makeContext(tn, &outerConn, nil, kvs))
+
 		for !userEnd {
 			n, err = outerConn.Read(buf[16:]) // type(4) + length(4) + uid(4) + close(4)
 			if err != nil {
@@ -380,6 +392,7 @@ func (s *ProxyLiteServer) startTunnel(tn *tunnel) {
 				writeUidUnsafe(buf[8:], uid)
 			}
 
+			s.onForwardUserToTunnel.Trigger(makeContext(tn, &outerConn, buf[8:n+8], kvs))
 			err = sendMessageOnBuffer(*tn.innerConn, TypeDataSegment, buf, n+8) // + uid(4) + close(4)
 			if err != nil {
 				// end this tunnel
@@ -445,3 +458,27 @@ func (s *ProxyLiteServer) logProtocolMessage(source, header string, req, rsp int
 func (s *ProxyLiteServer) logTunnelMessage(service, header, msg string) {
 	s.logger.Infof("[%s] [%s] %s", service, header, msg)
 }
+
+func (s *ProxyLiteServer) OnTunnelCreated(f HookFunc) {
+	s.onTunnelCreated = f
+}
+
+func (s *ProxyLiteServer) OnTunnelDestroyed(f HookFunc) {
+	s.onTunnelDestroyed = f
+}
+
+func (s *ProxyLiteServer) OnUserComming(f HookFunc) {
+	s.onUserComming = f
+}
+
+func (s *ProxyLiteServer) OnUserLeaving(f HookFunc) {
+	s.onUserLeaving = f
+}
+
+func (s *ProxyLiteServer) OnForwardTunnelToUser(f HookFunc) {
+	s.onForwardTunnelToUser = f
+}
+
+func (s *ProxyLiteServer) OnForwardUserToTunnel(f HookFunc) {
+	s.onForwardUserToTunnel = f
+}