From 87f0a26e5ea6221925b536101d9827b918dd4d85 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=A0=E5=BE=B7=E6=BB=A1?= <go_developer@163.com>
Date: Fri, 9 Apr 2021 15:13:22 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=BB=BA=E7=AB=8B=E8=BF=9E?=
 =?UTF-8?q?=E6=8E=A5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 abstract/IWebsocket.go |  2 +-
 construct.go           | 16 +++++++++++++---
 example/server.go      |  3 ++-
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/abstract/IWebsocket.go b/abstract/IWebsocket.go
index a1a4640..053c6a2 100644
--- a/abstract/IWebsocket.go
+++ b/abstract/IWebsocket.go
@@ -22,7 +22,7 @@ type IWebsocket interface {
 	// Author : go_developer@163.com<张德满>
 	//
 	// Date : 6:39 下午 2021/3/27
-	Connect(ctx *context.WSContext)
+	Connect(ctx *context.WSContext) error
 	// Disconnect 断开连接时处理的方法
 	//
 	// Author : go_developer@163.com<张德满>
diff --git a/construct.go b/construct.go
index fdae93d..c7f3880 100644
--- a/construct.go
+++ b/construct.go
@@ -89,10 +89,11 @@ func NewWebsocketServe(wsInstanceList ...abstract.IWebsocket) error {
 				ctxInterface, _ := session.Get("ws_context")
 				ctx := ctxInterface.(*context.WSContext)
 				ctx.Session = session
-				if nil != storage.Connection {
-					storage.Connection.Store(ctx)
+				if err := wsInstance.Connect(ctx); nil == err {
+					if nil != storage.Connection {
+						storage.Connection.Store(ctx)
+					}
 				}
-				wsInstance.Connect(ctx)
 			})
 			// 2. 指令处理的函数
 			ginRouterTable[wsInstance.GetServerPort()].wsServer.HandleMessage(func(session *melody.Session, bytes []byte) {
@@ -106,6 +107,9 @@ func NewWebsocketServe(wsInstanceList ...abstract.IWebsocket) error {
 				ctxInterface, _ := session.Get("ws_context")
 				ctx := ctxInterface.(*context.WSContext)
 				defer func() {
+					if nil == storage.Connection {
+						return
+					}
 					storage.Connection.Del(ctx)
 				}()
 				return wsInstance.Close(ctx, i, s)
@@ -114,6 +118,12 @@ func NewWebsocketServe(wsInstanceList ...abstract.IWebsocket) error {
 			ginRouterTable[wsInstance.GetServerPort()].wsServer.HandleDisconnect(func(session *melody.Session) {
 				ctxInterface, _ := session.Get("ws_context")
 				ctx := ctxInterface.(*context.WSContext)
+				defer func() {
+					if nil == storage.Connection {
+						return
+					}
+					storage.Connection.Del(ctx)
+				}()
 				wsInstance.Disconnect(ctx)
 			})
 			// 注册指令
diff --git a/example/server.go b/example/server.go
index 4e0869e..dc31a50 100644
--- a/example/server.go
+++ b/example/server.go
@@ -24,9 +24,10 @@ func main() {
 type Example struct {
 }
 
-func (e Example) Connect(ctx *context.WSContext) {
+func (e Example) Connect(ctx *context.WSContext) error {
 	fmt.Println("建立连接成功")
 	message.Response(ctx, map[string]interface{}{"say": "hello world!", "cid": ctx.ConnectionID})
+	return nil
 }
 
 func (e Example) Disconnect(ctx *context.WSContext) {