From 30dfde090f705a96cea18e0d8b68296398c3edd1 Mon Sep 17 00:00:00 2001
From: Constantine Shablia <constantine.shablya@collabora.com>
Date: Wed, 31 Jan 2024 16:17:57 +0200
Subject: [PATCH] introduce an API that lets the user get the current datagram
 size

Closes #4259
---
 connection.go       | 18 ++++++++++++++++++
 interface.go        |  9 +++++++++
 mock_packer_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 packet_packer.go    | 13 +++++++++++++
 4 files changed, 78 insertions(+)

diff --git a/connection.go b/connection.go
index 895c252422d..9ebd0ea763b 100644
--- a/connection.go
+++ b/connection.go
@@ -637,6 +637,8 @@ runLoop:
 		} else {
 			sendQueueAvailable = nil
 		}
+
+		s.updateMaxDatagramDataSize()
 	}
 
 	s.cryptoStreamHandler.Close()
@@ -679,6 +681,22 @@ func (s *connection) ConnectionState() ConnectionState {
 	return s.connState
 }
 
+func (s *connection) updateMaxDatagramDataSize() {
+	if s.peerParams == nil {
+		return
+	}
+	if !s.supportsDatagrams() {
+		return
+	}
+
+	maxDatagramFrameSize := min(s.peerParams.MaxDatagramFrameSize, s.packer.MaxPayloadSize(s.mtuDiscoverer.CurrentSize()))
+	maxDatagramDataSize := (&wire.DatagramFrame{DataLenPresent: true}).MaxDataLen(maxDatagramFrameSize, s.version)
+
+	s.connStateMutex.Lock()
+	s.connState.MaxDatagramSize = int(maxDatagramDataSize)
+	s.connStateMutex.Unlock()
+}
+
 // Time when the connection should time out
 func (s *connection) nextIdleTimeoutTime() time.Time {
 	idleTimeout := max(s.idleTimeout, s.rttStats.PTO(true)*3)
diff --git a/interface.go b/interface.go
index b269d790ef0..3b6269c0e27 100644
--- a/interface.go
+++ b/interface.go
@@ -344,6 +344,15 @@ type ConnectionState struct {
 	// If datagram support was negotiated, datagrams can be sent and received using the
 	// SendDatagram and ReceiveDatagram methods on the Connection.
 	SupportsDatagrams bool
+	// MaxDatagramSize specifies how big a datagram can be sent using
+	// SendDatagram. Datagrams bigger than MaxDatagramSize are silently dropped.
+	//
+	// Note: MaxDatagramSize can grow or shrink at any time. Users should call
+	// ConnectionState every once in a while to get up to date MaxDatagramSize.
+	//
+	// MaxDatagramSize is zero if datagrams are not supported on this
+	// connection.
+	MaxDatagramSize int
 	// Used0RTT says if 0-RTT resumption was used.
 	Used0RTT bool
 	// Version is the QUIC version of the QUIC connection.
diff --git a/mock_packer_test.go b/mock_packer_test.go
index 8ef1c323a01..3d8a49c79a2 100644
--- a/mock_packer_test.go
+++ b/mock_packer_test.go
@@ -79,6 +79,44 @@ func (c *PackerAppendPacketCall) DoAndReturn(f func(*packetBuffer, protocol.Byte
 	return c
 }
 
+// MaxPayloadSize mocks base method.
+func (m *MockPacker) MaxPayloadSize(arg0 protocol.ByteCount) protocol.ByteCount {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "MaxPayloadSize", arg0)
+	ret0, _ := ret[0].(protocol.ByteCount)
+	return ret0
+}
+
+// MaxPayloadSize indicates an expected call of MaxPayloadSize.
+func (mr *MockPackerMockRecorder) MaxPayloadSize(arg0 any) *PackerMaxPayloadSizeCall {
+	mr.mock.ctrl.T.Helper()
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MaxPayloadSize", reflect.TypeOf((*MockPacker)(nil).MaxPayloadSize), arg0)
+	return &PackerMaxPayloadSizeCall{Call: call}
+}
+
+// PackerMaxPayloadSizeCall wrap *gomock.Call
+type PackerMaxPayloadSizeCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *PackerMaxPayloadSizeCall) Return(arg0 protocol.ByteCount) *PackerMaxPayloadSizeCall {
+	c.Call = c.Call.Return(arg0)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *PackerMaxPayloadSizeCall) Do(f func(protocol.ByteCount) protocol.ByteCount) *PackerMaxPayloadSizeCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *PackerMaxPayloadSizeCall) DoAndReturn(f func(protocol.ByteCount) protocol.ByteCount) *PackerMaxPayloadSizeCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
+}
+
 // MaybePackProbePacket mocks base method.
 func (m *MockPacker) MaybePackProbePacket(arg0 protocol.EncryptionLevel, arg1 protocol.ByteCount, arg2 protocol.VersionNumber) (*coalescedPacket, error) {
 	m.ctrl.T.Helper()
diff --git a/packet_packer.go b/packet_packer.go
index 9a97295264e..b9442619f57 100644
--- a/packet_packer.go
+++ b/packet_packer.go
@@ -25,6 +25,7 @@ type packer interface {
 	PackConnectionClose(*qerr.TransportError, protocol.ByteCount, protocol.VersionNumber) (*coalescedPacket, error)
 	PackApplicationClose(*qerr.ApplicationError, protocol.ByteCount, protocol.VersionNumber) (*coalescedPacket, error)
 	PackMTUProbePacket(ping ackhandler.Frame, size protocol.ByteCount, v protocol.VersionNumber) (shortHeaderPacket, *packetBuffer, error)
+	MaxPayloadSize(maxPacketSize protocol.ByteCount) protocol.ByteCount
 
 	SetToken([]byte)
 }
@@ -325,6 +326,18 @@ func (p *packetPacker) initialPaddingLen(frames []ackhandler.Frame, currentSize,
 	return maxPacketSize - currentSize
 }
 
+func (p *packetPacker) MaxPayloadSize(maxPacketSize protocol.ByteCount) protocol.ByteCount {
+	oneRTTSealer, err := p.cryptoSetup.Get1RTTSealer()
+	if err == nil {
+		connID := p.getDestConnID()
+		_, pnLen := p.pnManager.PeekPacketNumber(protocol.Encryption1RTT)
+		hdrLen := wire.ShortHeaderLen(connID, pnLen)
+		return maxPacketSize - hdrLen - protocol.ByteCount(oneRTTSealer.Overhead())
+	}
+
+	return 0
+}
+
 // PackCoalescedPacket packs a new packet.
 // It packs an Initial / Handshake if there is data to send in these packet number spaces.
 // It should only be called before the handshake is confirmed.