fix: Call ConnContext before reading first packet

loopholelabs · Oct 14, 2023 · fe588a8 · fe588a8
1 parent 870bd1d
commit fe588a8
Showing 1 changed file with 3 additions and 9 deletions.
diff --git a/server.go b/server.go
@@ -309,9 +309,6 @@ func (s *Server) handleSinglePacket(frisbeeConn *Async, connCtx context.Context)
 		s.onClosed(frisbeeConn, err)
 		return
 	}
-	if s.ConnContext != nil {
-		connCtx = s.ConnContext(connCtx, frisbeeConn)
-	}
 	for {
 		handlerFunc = s.handlerTable[p.Metadata.Operation]
 		if handlerFunc != nil {
@@ -361,9 +358,6 @@ func (s *Server) handleUnlimitedPacket(frisbeeConn *Async, connCtx context.Conte
 		s.onClosed(frisbeeConn, err)
 		return
 	}
-	if s.ConnContext != nil {
-		connCtx = s.ConnContext(connCtx, frisbeeConn)
-	}
 	wg := new(sync.WaitGroup)
 	closed := atomic.NewBool(false)
 	connCtx, cancel := context.WithCancel(connCtx)
@@ -391,9 +385,6 @@ func (s *Server) handleLimitedPacket(frisbeeConn *Async, connCtx context.Context
 		s.onClosed(frisbeeConn, err)
 		return
 	}
-	if s.ConnContext != nil {
-		connCtx = s.ConnContext(connCtx, frisbeeConn)
-	}
 	wg := new(sync.WaitGroup)
 	closed := atomic.NewBool(false)
 	connCtx, cancel := context.WithCancel(connCtx)
@@ -465,6 +456,9 @@ func (s *Server) serveConn(newConn net.Conn) {
 	}
 	s.connections[frisbeeConn] = struct{}{}
 	s.connectionsMu.Unlock()
+	if s.ConnContext != nil {
+		connCtx = s.ConnContext(connCtx, frisbeeConn)
+	}
 	if s.concurrency == 0 {
 		s.handleUnlimitedPacket(frisbeeConn, connCtx)
 	} else if s.concurrency == 1 {