From f242675595a473bff3731641537dd8bfa94559bf Mon Sep 17 00:00:00 2001 From: NicholasChan1025 <29789197+NicholasChan1025@users.noreply.github.com> Date: Tue, 8 Dec 2020 10:59:25 +0800 Subject: [PATCH 1/5] Update env.go MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 添加握手验证hook --- internal/env/env.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/internal/env/env.go b/internal/env/env.go index 4722904e..27df173c 100644 --- a/internal/env/env.go +++ b/internal/env/env.go @@ -38,6 +38,7 @@ var ( CheckOrigin func(*http.Request) bool // check origin when websocket enabled Debug bool // enable Debug WSPath string // WebSocket path(eg: ws://127.0.0.1/WSPath) + HandshakeValidator func([]byte) error // When you need to verify the custom data of the handshake request // timerPrecision indicates the precision of timer, default is time.Second TimerPrecision = time.Second @@ -49,6 +50,8 @@ var ( Serializer serialize.Serializer GrpcOptions = []grpc.DialOption{grpc.WithInsecure()} + + ) func init() { @@ -56,5 +59,7 @@ func init() { Heartbeat = 30 * time.Second Debug = false CheckOrigin = func(_ *http.Request) bool { return true } + HandshakeValidator func(_ []byte) error {return nil} Serializer = protobuf.NewSerializer() + } From 38c7fb76843ce9d1cc45d7d7d4e0837ef9f56d83 Mon Sep 17 00:00:00 2001 From: NicholasChan1025 <29789197+NicholasChan1025@users.noreply.github.com> Date: Tue, 8 Dec 2020 11:02:38 +0800 Subject: [PATCH 2/5] Update handler.go MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit processPacket ，添加校验握手数据 --- cluster/handler.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cluster/handler.go b/cluster/handler.go index 98cb2854..a2406a18 100644 --- a/cluster/handler.go +++ b/cluster/handler.go @@ -252,6 +252,12 @@ func (h *LocalHandler) handle(conn net.Conn) { func (h *LocalHandler) processPacket(agent *agent, p *packet.Packet) error { switch p.Type { case packet.Handshake: + + //Verify handshake data + if err := env.HandshakeValidator(p.Data); err != nil { + return err + } + if _, err := agent.conn.Write(hrd); err != nil { return err } From 34b899b50cbcbf6350e825ef108dd59945dcf268 Mon Sep 17 00:00:00 2001 From: NicholasChan1025 <29789197+NicholasChan1025@users.noreply.github.com> Date: Tue, 8 Dec 2020 11:07:00 +0800 Subject: [PATCH 3/5] Update options.go MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增WithHandshakeValidator 方法 --- options.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/options.go b/options.go index c9996c91..13521cdc 100644 --- a/options.go +++ b/options.go @@ -145,3 +145,9 @@ func WithLogger(l log.Logger) Option { log.SetLogger(l) } } +// WithHandshakeValidator sets the function that Verify `handshake` data +func WithHandshakeValidator(fn func([]byte) error) Option { + return func(opt *cluster.Options) { + env.HandshakeValidator = fn + } +} From fa397f45e710c23349c56d481e92f321ad2c7666 Mon Sep 17 00:00:00 2001 From: NicholasChan1025 <29789197+NicholasChan1025@users.noreply.github.com> Date: Tue, 8 Dec 2020 17:55:51 +0800 Subject: [PATCH 4/5] Update env.go MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 修正‘=’缺失 --- internal/env/env.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/env/env.go b/internal/env/env.go index 27df173c..b925387e 100644 --- a/internal/env/env.go +++ b/internal/env/env.go @@ -59,7 +59,7 @@ func init() { Heartbeat = 30 * time.Second Debug = false CheckOrigin = func(_ *http.Request) bool { return true } - HandshakeValidator func(_ []byte) error {return nil} + HandshakeValidator = func(_ []byte) error {return nil} Serializer = protobuf.NewSerializer() } From f5456842c3dded87a7835499ad477d82fe80c9c3 Mon Sep 17 00:00:00 2001 From: black Date: Wed, 9 Dec 2020 15:52:48 +0800 Subject: [PATCH 5/5] Format Code --- cluster/handler.go | 8 +++----- internal/env/env.go | 19 ++++++++----------- options.go | 3 ++- 3 files changed, 13 insertions(+), 17 deletions(-) diff --git a/cluster/handler.go b/cluster/handler.go index a2406a18..40298ef2 100644 --- a/cluster/handler.go +++ b/cluster/handler.go @@ -252,12 +252,10 @@ func (h *LocalHandler) handle(conn net.Conn) { func (h *LocalHandler) processPacket(agent *agent, p *packet.Packet) error { switch p.Type { case packet.Handshake: - - //Verify handshake data - if err := env.HandshakeValidator(p.Data); err != nil { - return err + if err := env.HandshakeValidator(p.Data); err != nil { + return err } - + if _, err := agent.conn.Write(hrd); err != nil { return err } diff --git a/internal/env/env.go b/internal/env/env.go index b925387e..971ebdaa 100644 --- a/internal/env/env.go +++ b/internal/env/env.go @@ -32,13 +32,13 @@ import ( ) var ( - Wd string // working path - Die chan bool // wait for end application - Heartbeat time.Duration // Heartbeat internal - CheckOrigin func(*http.Request) bool // check origin when websocket enabled - Debug bool // enable Debug - WSPath string // WebSocket path(eg: ws://127.0.0.1/WSPath) - HandshakeValidator func([]byte) error // When you need to verify the custom data of the handshake request + Wd string // working path + Die chan bool // wait for end application + Heartbeat time.Duration // Heartbeat internal + CheckOrigin func(*http.Request) bool // check origin when websocket enabled + Debug bool // enable Debug + WSPath string // WebSocket path(eg: ws://127.0.0.1/WSPath) + HandshakeValidator func([]byte) error // When you need to verify the custom data of the handshake request // timerPrecision indicates the precision of timer, default is time.Second TimerPrecision = time.Second @@ -50,8 +50,6 @@ var ( Serializer serialize.Serializer GrpcOptions = []grpc.DialOption{grpc.WithInsecure()} - - ) func init() { @@ -59,7 +57,6 @@ func init() { Heartbeat = 30 * time.Second Debug = false CheckOrigin = func(_ *http.Request) bool { return true } - HandshakeValidator = func(_ []byte) error {return nil} + HandshakeValidator = func(_ []byte) error { return nil } Serializer = protobuf.NewSerializer() - } diff --git a/options.go b/options.go index 13521cdc..5ae6f615 100644 --- a/options.go +++ b/options.go @@ -145,7 +145,8 @@ func WithLogger(l log.Logger) Option { log.SetLogger(l) } } -// WithHandshakeValidator sets the function that Verify `handshake` data + +// WithHandshakeValidator sets the function that Verify `handshake` data func WithHandshakeValidator(fn func([]byte) error) Option { return func(opt *cluster.Options) { env.HandshakeValidator = fn