From 8ba623111b488b803ed9d6e2e6d05ba2636dc07d Mon Sep 17 00:00:00 2001
From: crusader <crusader@loafle.com>
Date: Thu, 9 Nov 2017 18:47:30 +0900
Subject: [PATCH] ing

---
 server.go          | 4 +++-
 socket_handler.go  | 2 ++
 socket_handlers.go | 7 +++++++
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/server.go b/server.go
index 15047cf..dd0947f 100644
--- a/server.go
+++ b/server.go
@@ -151,6 +151,8 @@ func handleConnection(s *server, soc *Socket, socketHandler SocketHandler) {
 	defer s.stopWg.Done()
 
 	logging.Logger().Debug(fmt.Sprintf("Server: Client[%s] is connected.", soc.RemoteAddr()))
+	socketHandler.OnConnect(soc)
+
 	clientStopChan := make(chan struct{})
 	handleDoneChan := make(chan struct{})
 
@@ -159,11 +161,11 @@ func handleConnection(s *server, soc *Socket, socketHandler SocketHandler) {
 	select {
 	case <-s.stopChan:
 		close(clientStopChan)
-		soc.Close()
 		<-handleDoneChan
 	case <-handleDoneChan:
 		close(clientStopChan)
 		logging.Logger().Debug(fmt.Sprintf("Server: Client[%s] is disconnected.", soc.RemoteAddr()))
+		socketHandler.OnDisconnect(soc)
 		soc.Close()
 	}
 }
diff --git a/socket_handler.go b/socket_handler.go
index 46ebe57..b628c8b 100644
--- a/socket_handler.go
+++ b/socket_handler.go
@@ -10,7 +10,9 @@ type SocketHandler interface {
 	// Handshake do handshake client and server
 	// id is identity of client socket. if id is nil, disallow connection
 	Handshake(ctx *fasthttp.RequestCtx) (id interface{}, extensionsHeader *fasthttp.ResponseHeader)
+	OnConnect(soc *Socket)
 	Handle(soc *Socket, stopChan <-chan struct{}, doneChan chan<- struct{})
+	OnDisconnect(soc *Socket)
 
 	GetMaxMessageSize() int64
 	GetWriteTimeout() time.Duration
diff --git a/socket_handlers.go b/socket_handlers.go
index b8ce65b..8c54112 100644
--- a/socket_handlers.go
+++ b/socket_handlers.go
@@ -30,11 +30,18 @@ type SocketHandlers struct {
 func (sh *SocketHandlers) Handshake(ctx *fasthttp.RequestCtx) (id interface{}, extensionsHeader *fasthttp.ResponseHeader) {
 	return nil, nil
 }
+func (sh *SocketHandlers) OnConnect(soc *Socket) {
+	// no op
+}
 
 func (sh *SocketHandlers) Handle(soc *Socket, stopChan <-chan struct{}, doneChan chan<- struct{}) {
 	// no op
 }
 
+func (sh *SocketHandlers) OnDisconnect(soc *Socket) {
+	// no op
+}
+
 func (sh *SocketHandlers) GetMaxMessageSize() int64 {
 	return sh.MaxMessageSize
 }