Add closed handling in heartbeat read

refraction-networking · Feb 11, 2024 · 6ede634 · 6ede634
1 parent 6b779e2
commit 6ede634
Showing 1 changed file with 14 additions and 9 deletions.
diff --git a/pkg/dtls/heartbeat.go b/pkg/dtls/heartbeat.go
@@ -3,6 +3,7 @@ package dtls
 import (
 	"bytes"
 	"errors"
+	"net"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -101,18 +102,22 @@ func (c *hbConn) Write(b []byte) (n int, err error) {
 }
 
 func (c *hbConn) Read(b []byte) (int, error) {
-	readBytes := <-c.recvCh
-	if readBytes.err != nil {
-		return 0, readBytes.err
-	}
+	select {
+	case <-c.closed:
+		return 0, net.ErrClosed
+	case readBytes := <-c.recvCh:
+		if readBytes.err != nil {
+			return 0, readBytes.err
+		}
 
-	if len(b) < len(readBytes.b) {
-		return 0, ErrInsufficientBuffer
-	}
+		if len(b) < len(readBytes.b) {
+			return 0, ErrInsufficientBuffer
+		}
 
-	n := copy(b, readBytes.b)
+		n := copy(b, readBytes.b)
 
-	return n, nil
+		return n, nil
+	}
 }
 
 func (c *hbConn) BufferedAmount() uint64 {