Fix: Call ConnContext before Reading First Packet

.github/workflows/tests.yaml

@@ -7,11 +7,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Install Go
-        uses: actions/setup-go@v3
+        uses: actions/setup-go@v4
         with:
-          go-version: "1.20"
+          go-version: "1.22"
           check-latest: true
           cache: true
       - name: Run Tests
@@ -20,11 +20,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Install Go
-        uses: actions/setup-go@v3
+        uses: actions/setup-go@v4
         with:
-          go-version: "1.20"
+          go-version: "1.22"
           check-latest: true
           cache: true
       - name: Test with Race Conditions
@@ -34,11 +34,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Install Go
-        uses: actions/setup-go@v3
+        uses: actions/setup-go@v4
         with:
-          go-version: "1.20"
+          go-version: "1.22"
           check-latest: true
           cache: true
       - name: Benchmark
@@ -47,11 +47,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Install Go
-        uses: actions/setup-go@v3
+        uses: actions/setup-go@v4
         with:
-          go-version: "1.20"
+          go-version: "1.22"
           check-latest: true
           cache: true
       - name: Benchmark with Race Conditions

async.go

@@ -183,7 +183,7 @@ func (c *Async) WritePacket(p *packet.Packet) error {
 	if p.Metadata.Operation <= RESERVED9 {
 		return InvalidOperation
 	}
-	return c.writePacket(p)
+	return c.writePacket(p, true)
 }
 
 // ReadPacket is a blocking function that will wait until a Frisbee packet is available and then return it (and its content).
@@ -300,7 +300,7 @@ func (c *Async) Close() error {
 }
 
 // write packet is the internal write packet function that does not check for reserved operations.
-func (c *Async) writePacket(p *packet.Packet) error {
+func (c *Async) writePacket(p *packet.Packet, closeOnErr bool) error {
 	if int(p.Metadata.ContentLength) != len(*p.Content) {
 		return InvalidContentLength
 	}
@@ -323,7 +323,10 @@ func (c *Async) writePacket(p *packet.Packet) error {
 			return ConnectionClosed
 		}
 		c.Logger().Debug().Err(err).Uint16("Packet ID", p.Metadata.Id).Msg("error while setting write deadline before writing packet")
-		return c.closeWithError(err)
+		if closeOnErr {
+			return c.closeWithError(err)
+		}
+		return err
 	}
 	_, err = c.writer.Write(encodedMetadata[:])
 	metadata.PutBuffer(encodedMetadata)
@@ -334,7 +337,10 @@ func (c *Async) writePacket(p *packet.Packet) error {
 			return ConnectionClosed
 		}
 		c.Logger().Debug().Err(err).Uint16("Packet ID", p.Metadata.Id).Msg("error while writing encoded metadata")
-		return c.closeWithError(err)
+		if closeOnErr {
+			return c.closeWithError(err)
+		}
+		return err
 	}
 	if p.Metadata.ContentLength != 0 {
 		_, err = c.writer.Write((*p.Content)[:p.Metadata.ContentLength])
@@ -345,7 +351,10 @@ func (c *Async) writePacket(p *packet.Packet) error {
 				return ConnectionClosed
 			}
 			c.Logger().Debug().Err(err).Uint16("Packet ID", p.Metadata.Id).Msg("error while writing packet content")
-			return c.closeWithError(err)
+			if closeOnErr {
+				return c.closeWithError(err)
+			}
+			return err
 		}
 	}
 
@@ -457,7 +466,7 @@ func (c *Async) pingLoop() {
 			c.wg.Done()
 			return
 		case <-ticker.C:
-			err = c.writePacket(PINGPacket)
+			err = c.writePacket(PINGPacket, false)
 			if err != nil {
 				c.wg.Done()
 				_ = c.closeWithError(err)
@@ -516,7 +525,7 @@ func (c *Async) readLoop() {
 			switch p.Metadata.Operation {
 			case PING:
 				c.Logger().Debug().Msg("PING Packet received by read loop, sending back PONG packet")
-				err = c.writePacket(PONGPacket)
+				err = c.writePacket(PONGPacket, false)
 				if err != nil {
 					c.wg.Done()
 					_ = c.closeWithError(err)
@@ -541,13 +550,13 @@ func (c *Async) readLoop() {
 			default:
 				if p.Metadata.ContentLength > 0 {
 					if n-index < int(p.Metadata.ContentLength) {
-						min := int(p.Metadata.ContentLength) - p.Content.Write(buf[index:n])
+						minSize := int(p.Metadata.ContentLength) - p.Content.Write(buf[index:n])
 						n = 0
-						for cap(buf) < min {
+						for cap(buf) < minSize {
 							buf = append(buf[:cap(buf)], 0)
 						}
 						buf = buf[:cap(buf)]
-						for n < min {
+						for n < minSize {
 							var nn int
 							err = c.conn.SetReadDeadline(time.Now().Add(DefaultDeadline))
 							if err != nil {
@@ -558,16 +567,16 @@ func (c *Async) readLoop() {
 							nn, err = c.conn.Read(buf[n:])
 							n += nn
 							if err != nil {
-								if n < min {
+								if n < minSize {
 									c.wg.Done()
 									_ = c.closeWithError(err)
 									return
 								}
 								break
 							}
 						}
-						p.Content.Write(buf[:min])
-						index = min
+						p.Content.Write(buf[:minSize])
+						index = minSize
 					} else {
 						index += p.Content.Write(buf[index : index+int(p.Metadata.ContentLength)])
 					}
@@ -649,14 +658,14 @@ func (c *Async) readLoop() {
 				index = n
 
 				buf = buf[:cap(buf)]
-				min := metadata.Size - index
-				if len(buf) < min {
+				minSize := metadata.Size - index
+				if len(buf) < minSize {
 					c.wg.Done()
 					_ = c.closeWithError(InvalidBufferLength)
 					return
 				}
 				n = 0
-				for n < min {
+				for n < minSize {
 					var nn int
 					err = c.conn.SetReadDeadline(time.Now().Add(DefaultDeadline))
 					if err != nil {
@@ -667,7 +676,7 @@ func (c *Async) readLoop() {
 					nn, err = c.conn.Read(buf[index+n:])
 					n += nn
 					if err != nil {
-						if n < min {
+						if n < minSize {
 							c.wg.Done()
 							_ = c.closeWithError(err)
 							return

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,11 +456,15 @@ func (s *Server) serveConn(newConn net.Conn) {
 	}
 	s.connections[frisbeeConn] = struct{}{}
 	s.connectionsMu.Unlock()
-	if s.concurrency == 0 {
+	if s.ConnContext != nil {
+		connCtx = s.ConnContext(connCtx, frisbeeConn)
+	}
+	switch s.concurrency {
+	case 0:
 		s.handleUnlimitedPacket(frisbeeConn, connCtx)
-	} else if s.concurrency == 1 {
+	case 1:
 		s.handleSinglePacket(frisbeeConn, connCtx)
-	} else {
+	default:
 		s.handleLimitedPacket(frisbeeConn, connCtx)
 	}
 	s.connectionsMu.Lock()

stream.go

@@ -92,7 +92,7 @@ func (s *Stream) WritePacket(p *packet.Packet) error {
 	}
 	p.Metadata.Id = s.id
 	p.Metadata.Operation = STREAM
-	return s.conn.writePacket(p)
+	return s.conn.writePacket(p, true)
 }
 
 // ID returns the stream's ID.
@@ -116,7 +116,7 @@ func (s *Stream) Close() error {
 		p := packet.Get()
 		p.Metadata.Id = s.id
 		p.Metadata.Operation = STREAM
-		err := s.conn.writePacket(p)
+		err := s.conn.writePacket(p, true)
 		packet.Put(p)
 
 		s.conn.streamsMu.Lock()