From a931e1f2ce7a8f56f18101b4092ec0df4a6e611d Mon Sep 17 00:00:00 2001 From: piaodazhu <2283079466@qq.com> Date: Tue, 22 Aug 2023 14:45:46 +0800 Subject: [PATCH] support hookfunc --- README.md | 4 ++ README_ZH.md | 4 ++ TODO | 4 +- context.go | 83 +++++++++++++++++++++++++++++++++++++++ go.mod | 2 +- hook.go | 9 +++++ hook_test.go | 31 +++++++++++++++ proxylite_test.go | 99 +++++++++++++++++++++++++++++++++++++++++++++++ server.go | 49 ++++++++++++++++++++--- 9 files changed, 276 insertions(+), 9 deletions(-) create mode 100644 context.go create mode 100644 hook.go create mode 100644 hook_test.go diff --git a/README.md b/README.md index c265bd8..12645e1 100644 --- a/README.md +++ b/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) + + [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. diff --git a/README_ZH.md b/README_ZH.md index 8d63c12..b8c9085 100644 --- a/README_ZH.md +++ b/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) + + [English](./README.md)|[中文](./README_ZH.md) ---- 以下中文文档由ChatGPT生成 ---- diff --git a/TODO b/TODO index 9b7837b..5974a41 100644 --- a/TODO +++ b/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 \ No newline at end of file diff --git a/context.go b/context.go new file mode 100644 index 0000000..1a1e2bd --- /dev/null +++ b/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) +} diff --git a/go.mod b/go.mod index 67e993b..cd376f6 100644 --- a/go.mod +++ b/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 diff --git a/hook.go b/hook.go new file mode 100644 index 0000000..1aa92b8 --- /dev/null +++ b/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) + } +} diff --git a/hook_test.go b/hook_test.go new file mode 100644 index 0000000..a188cf3 --- /dev/null +++ b/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") + } +} diff --git a/proxylite_test.go b/proxylite_test.go index 00e4ad1..c64e48f 100644 --- a/proxylite_test.go +++ b/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: + } +} diff --git a/server.go b/server.go index 9577425..d4f566b 100644 --- a/server.go +++ b/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 +}