From 9151a013869e65025d038008c9c7560213aaf564 Mon Sep 17 00:00:00 2001
From: "Gerasimos (Makis) Maropoulos" <kataras2006@hotmail.com>
Date: Fri, 5 Jul 2019 16:10:24 +0300
Subject: [PATCH] add customID to the Upgrade method so Iris side can do its
 magic to set a custom ID Generator as requested at:
 https://github.com/kataras/neffos/issues/1#issuecomment-508689819

---
 server.go | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/server.go b/server.go
index fb45a14..3570a6b 100644
--- a/server.go
+++ b/server.go
@@ -198,8 +198,14 @@ func IsTryingToReconnect(err error) (ok bool) {
 const websocketReconectHeaderKey = "X-Websocket-Reconnect"
 
 // Upgrade handles the connection, same as `ServeHTTP` but it can accept
-// a socket wrapper and it does return the connection or any errors.
-func (s *Server) Upgrade(w http.ResponseWriter, r *http.Request, socketWrapper func(Socket) Socket) (*Conn, error) {
+// a socket wrapper and a "customID" that overrides the server's IDGenerator
+// and it does return the connection or any errors.
+func (s *Server) Upgrade(
+	w http.ResponseWriter,
+	r *http.Request,
+	socketWrapper func(Socket) Socket,
+	customID string,
+) (*Conn, error) {
 	if atomic.LoadUint32(&s.closed) > 0 {
 		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
 		return nil, errServerClosed
@@ -237,7 +243,12 @@ func (s *Server) Upgrade(w http.ResponseWriter, r *http.Request, socketWrapper f
 	}
 
 	c := newConn(socket, s.namespaces)
-	c.id = s.IDGenerator(w, r)
+	if customID != "" {
+		c.id = customID
+	} else {
+		c.id = s.IDGenerator(w, r)
+	}
+
 	c.readTimeout = s.readTimeout
 	c.writeTimeout = s.writeTimeout
 	c.server = s
@@ -298,7 +309,7 @@ func (s *Server) Upgrade(w http.ResponseWriter, r *http.Request, socketWrapper f
 // ServeHTTP completes the `http.Handler` interface, it should be passed on a http server's router
 // to serve this neffos server on a specific endpoint.
 func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	s.Upgrade(w, r, nil)
+	s.Upgrade(w, r, nil, "")
 }
 
 func (s *Server) waitMessage(c *Conn) bool {