From 14b0aea1ca398a39779d992f3284f66600053597 Mon Sep 17 00:00:00 2001
From: crusader <crusader@loafle.com>
Date: Wed, 20 Sep 2017 12:53:06 +0900
Subject: [PATCH] OnConnection changed to OnConnected

---
 server.go          |  7 ++++++-
 server_handler.go  |  3 ++-
 server_handlers.go | 13 ++++++++++++-
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/server.go b/server.go
index 94a63a1..01e22c5 100644
--- a/server.go
+++ b/server.go
@@ -89,6 +89,11 @@ func (s *server) onConnection(ctx *fasthttp.RequestCtx) {
 		return
 	}
 
+	if !s._sh.OnConnection(ctx) {
+		ctx.Error(http.StatusText(http.StatusNotAcceptable), http.StatusNotAcceptable)
+		return
+	}
+
 	s._upgrader.Upgrade(ctx, nil, func(conn *websocket.Conn, err error) {
 		if err != nil {
 			s.onError(ctx, fasthttp.StatusInternalServerError, err)
@@ -97,7 +102,7 @@ func (s *server) onConnection(ctx *fasthttp.RequestCtx) {
 		id := s._sh.OnIDGenerate(ctx)
 		soc := NewSocket(s._ctx, id, path, co, conn)
 		s.addSocket(soc)
-		s._sh.OnConnection(soc)
+		s._sh.OnConnected(soc)
 
 		soc.run()
 	})
diff --git a/server_handler.go b/server_handler.go
index d6bb098..3441234 100644
--- a/server_handler.go
+++ b/server_handler.go
@@ -12,7 +12,8 @@ type ServerHandler interface {
 	GetWriteBufferSize() int
 	GetEnableCompression() bool
 
-	OnConnection(soc Socket)
+	OnConnection(ctx *fasthttp.RequestCtx) bool
+	OnConnected(soc Socket)
 	OnDisconnected(soc Socket)
 	OnCheckOrigin(ctx *fasthttp.RequestCtx) bool
 	OnError(ctx *fasthttp.RequestCtx, status int, reason error)
diff --git a/server_handlers.go b/server_handlers.go
index ece1231..5018ff8 100644
--- a/server_handlers.go
+++ b/server_handlers.go
@@ -29,28 +29,39 @@ type ServerHandlers struct {
 func (sh *ServerHandlers) GetHandshakeTimeout() time.Duration {
 	return sh.HandshakeTimeout
 }
+
 func (sh *ServerHandlers) GetReadBufferSize() int {
 	return sh.ReadBufferSize
 }
+
 func (sh *ServerHandlers) GetWriteBufferSize() int {
 	return sh.WriteBufferSize
 }
+
 func (sh *ServerHandlers) GetEnableCompression() bool {
 	return sh.EnableCompression
 }
 
-func (sh *ServerHandlers) OnConnection(soc Socket) {
+func (sh *ServerHandlers) OnConnection(ctx *fasthttp.RequestCtx) bool {
+	return true
+}
+
+func (sh *ServerHandlers) OnConnected(soc Socket) {
 
 }
+
 func (sh *ServerHandlers) OnDisconnected(soc Socket) {
 
 }
+
 func (sh *ServerHandlers) OnCheckOrigin(ctx *fasthttp.RequestCtx) bool {
 	return true
 }
+
 func (sh *ServerHandlers) OnError(ctx *fasthttp.RequestCtx, status int, reason error) {
 
 }
+
 func (sh *ServerHandlers) OnIDGenerate(ctx *fasthttp.RequestCtx) string {
 	return uuid.NewV4().String()
 }