webrtc: fail Write early if deadline has exceeded before the call (#2578

)
libp2p · Sep 21, 2023 · 4ae549a · 4ae549a
1 parent 70d31da
commit 4ae549a
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/p2p/transport/webrtc/stream_write.go b/p2p/transport/webrtc/stream_write.go
@@ -35,6 +35,10 @@ func (s *stream) Write(b []byte) (int, error) {
 		s.readLoopOnce.Do(s.spawnControlMessageReader)
 	}
 
+	if !s.writeDeadline.IsZero() && time.Now().After(s.writeDeadline) {
+		return 0, os.ErrDeadlineExceeded
+	}
+
 	var writeDeadlineTimer *time.Timer
 	defer func() {
 		if writeDeadlineTimer != nil {

diff --git a/p2p/transport/webrtc/transport_test.go b/p2p/transport/webrtc/transport_test.go
@@ -402,10 +402,16 @@ func TestTransportWebRTC_Deadline(t *testing.T) {
 		stream, err := conn.OpenStream(context.Background())
 		require.NoError(t, err)
 
-		stream.SetWriteDeadline(time.Now().Add(200 * time.Millisecond))
+		stream.SetWriteDeadline(time.Now().Add(100 * time.Millisecond))
 		largeBuffer := make([]byte, 2*1024*1024)
 		_, err = stream.Write(largeBuffer)
 		require.ErrorIs(t, err, os.ErrDeadlineExceeded)
+
+		stream.SetWriteDeadline(time.Now().Add(-200 * time.Millisecond))
+		smallBuffer := make([]byte, 1024)
+		_, err = stream.Write(smallBuffer)
+		require.ErrorIs(t, err, os.ErrDeadlineExceeded)
+
 	})
 }