Workaround for unconstrained WriteTo loop in quic-go

psiphon/common/quic/quic.go

@@ -43,6 +43,7 @@ package quic
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -119,8 +120,8 @@ func Listen(
 	}
 
 	// This wrapping must be outermost to ensure that all
-	// ReadFrom errors are intercepted and logged.
-	packetConn = newLoggingPacketConn(logger, packetConn)
+	// ReadFrom/WriteTo calls are intercepted.
+	packetConn = newWorkaroundPacketConn(logger, packetConn)
 
 	quicListener, err := quic_go.Listen(
 		packetConn, tlsConfig, quicConfig)
@@ -420,7 +421,7 @@ func isErrorIndicatingClosed(err error) bool {
 	return false
 }
 
-// loggingPacketConn is a workaround for issues in the quic-go server (as of
+// workaroundPacketConn is a workaround for issues in the quic-go server (as of
 // revision ffdfa1).
 //
 // 1. quic-go will shutdown the QUIC server on any error returned from
@@ -440,33 +441,45 @@ func isErrorIndicatingClosed(err error) bool {
 //    packetHandlerMap and its mutex are used by all client sessions, this
 //    effectively hangs the entire server.
 //
-// loggingPacketConn PacketConn ReadFrom errors and returns any usable values
-// or loops and calls ReadFrom again. In practise, due to the nature of UDP
-// sockets, ReadFrom errors are exceptional as they will mosyt likely not
+// 3. In certain cases, quic-go appears to get into a state where it
+//    calls WriteTo in an unconstrained loop, far exceeding the expected
+//    rate for normal outbound traffic. This state pegs the psiphond
+//    CPU and, in the case of obfuscated QUIC, exhausts the 2^38 byte
+//    random padding key stream. To mitigate this, we rate limit
+//    workaroundPacketConn when an excessive WriteTo call rate is
+//    detected.
+//
+// workaroundPacketConn checks PacketConn ReadFrom errors and returns any usable
+// values or loops and calls ReadFrom again. In practise, due to the nature of
+// UDP sockets, ReadFrom errors are exceptional as they will most likely not
 // occur due to network transmission failures. ObfuscatedPacketConn returns
 // errors that could be due to network transmission failures that corrupt
-// packets; these are marked as net.Error.Temporary() and loggingPacketConn
+// packets; these are marked as net.Error.Temporary() and workaroundPacketConn
 // logs these at debug level.
 //
-// loggingPacketConn assumes quic-go revision ffdfa1 behavior and will break
-// other behavior, such as setting deadlines and expecting net.Error.Timeout()
+// workaroundPacketConn assumes specific quic-go behavior and will break other
+// use cases, such as setting deadlines and expecting net.Error.Timeout()
 // errors from ReadFrom.
-type loggingPacketConn struct {
+type workaroundPacketConn struct {
 	net.PacketConn
 	logger common.Logger
+
+	mutex      sync.Mutex
+	currentMS  time.Time
+	callsPerMS int
 }
 
-func newLoggingPacketConn(
+func newWorkaroundPacketConn(
 	logger common.Logger,
-	packetConn net.PacketConn) *loggingPacketConn {
+	packetConn net.PacketConn) *workaroundPacketConn {
 
-	return &loggingPacketConn{
+	return &workaroundPacketConn{
 		PacketConn: packetConn,
 		logger:     logger,
 	}
 }
 
-func (conn *loggingPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
+func (conn *workaroundPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
 
 	for {
 		n, addr, err := conn.PacketConn.ReadFrom(p)
@@ -488,3 +501,22 @@ func (conn *loggingPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
 		}
 	}
 }
+
+func (conn *workaroundPacketConn) WriteTo(p []byte, addr net.Addr) (int, error) {
+
+	conn.mutex.Lock()
+	currentMS := time.Now().Round(time.Millisecond)
+	if currentMS != conn.currentMS {
+		conn.currentMS = currentMS
+		conn.callsPerMS = 0
+	} else {
+		if conn.callsPerMS >= 10000 {
+			conn.mutex.Unlock()
+			return 0, common.ContextError(errors.New("rate limit exceeded"))
+		}
+		conn.callsPerMS += 1
+	}
+	conn.mutex.Unlock()
+
+	return conn.PacketConn.WriteTo(p, addr)
+}