TUN-8236: Add write timeout to quic and tcp connections

## Summary
To prevent bad eyeballs and severs to be able to exhaust the quic
control flows we are adding the possibility of having a timeout
for a write operation to be acknowledged. This will prevent hanging
connections from exhausting the quic control flows, creating a DDoS.

quic/safe_stream.go

@@ -1,20 +1,28 @@
 package quic
 
 import (
+	"errors"
+	"net"
 	"sync"
 	"time"
 
 	"github.com/quic-go/quic-go"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
 )
 
 type SafeStreamCloser struct {
-	lock   sync.Mutex
-	stream quic.Stream
+	lock         sync.Mutex
+	stream       quic.Stream
+	writeTimeout time.Duration
+	log          *zerolog.Logger
 }
 
-func NewSafeStreamCloser(stream quic.Stream) *SafeStreamCloser {
+func NewSafeStreamCloser(stream quic.Stream, writeTimeout time.Duration, log *zerolog.Logger) *SafeStreamCloser {
 	return &SafeStreamCloser{
-		stream: stream,
+		stream:       stream,
+		writeTimeout: writeTimeout,
+		log:          log,
 	}
 }
 
@@ -25,7 +33,29 @@ func (s *SafeStreamCloser) Read(p []byte) (n int, err error) {
 func (s *SafeStreamCloser) Write(p []byte) (n int, err error) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
-	return s.stream.Write(p)
+	if s.writeTimeout > 0 {
+		err = s.stream.SetWriteDeadline(time.Now().Add(s.writeTimeout))
+		if err != nil {
+			log.Err(err).Msg("Error setting write deadline for QUIC stream")
+		}
+	}
+	nBytes, err := s.stream.Write(p)
+	if err != nil {
+		s.handleTimeout(err)
+	}
+
+	return nBytes, err
+}
+
+// Handles the timeout error in case it happened, by canceling the stream write.
+func (s *SafeStreamCloser) handleTimeout(err error) {
+	var netErr net.Error
+	if errors.As(err, &netErr) {
+		if netErr.Timeout() {
+			s.log.Error().Err(netErr).Msg("Closing quic stream due to timeout while writing")
+			s.stream.CancelWrite(0)
+		}
+	}
 }
 
 func (s *SafeStreamCloser) Close() error {

quic/safe_stream_test.go

@@ -9,6 +9,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/rs/zerolog"
+
 	"github.com/quic-go/quic-go"
 	"github.com/stretchr/testify/require"
 )
@@ -70,7 +72,8 @@ func quicClient(t *testing.T, addr net.Addr) {
 		go func(iter int) {
 			defer wg.Done()
 
-			stream := NewSafeStreamCloser(quicStream)
+			log := zerolog.Nop()
+			stream := NewSafeStreamCloser(quicStream, 30*time.Second, &log)
 			defer stream.Close()
 
 			// Do a bunch of round trips over this stream that should work.
@@ -107,7 +110,8 @@ func quicServer(t *testing.T, serverReady *sync.WaitGroup, conn net.PacketConn)
 		go func(iter int) {
 			defer wg.Done()
 
-			stream := NewSafeStreamCloser(quicStream)
+			log := zerolog.Nop()
+			stream := NewSafeStreamCloser(quicStream, 30*time.Second, &log)
 			defer stream.Close()
 
 			// Do a bunch of round trips over this stream that should work.