Add read deadline to hbConn.recvLoop() to prevent blocking

refraction-networking · Feb 13, 2024 · 0b6aa3c · 0b6aa3c
1 parent 0f92b17
commit 0b6aa3c
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/pkg/dtls/heartbeat.go b/pkg/dtls/heartbeat.go
@@ -75,6 +75,7 @@ func (c *hbConn) recvLoop() {
 	for {
 		buffer := make([]byte, c.maxMessageSize)
 
+		c.stream.SetReadDeadline(time.Now().Add(c.timeout))
 		n, err := c.stream.Read(buffer)
 
 		if bytes.Equal(c.hb, buffer[:n]) {
@@ -87,7 +88,15 @@ func (c *hbConn) recvLoop() {
 			return
 		}
 
-		c.recvCh <- errBytes{buffer[:n], err}
+		timer := time.NewTimer(c.timeout)
+		select {
+		case c.recvCh <- errBytes{buffer[:n], err}:
+			timer.Stop()
+			continue
+		case <-timer.C:
+			c.Close()
+			return
+		}
 	}
 
 }