Add method for retrieving experiment option values from UDP endpoints

PiperOrigin-RevId: 805075889

Author: manninglucas

pkg/tcpip/stack/packet_buffer.go

@@ -768,3 +768,52 @@ func BufferSince(h PacketHeader) buffer.Buffer {
 	clone.TrimFront(int64(offset))
 	return clone
 }
+
+// ExperimentOptionValue returns the experiment option value from the packet
+// and a bool indicating whether an experiment option value was found.
+func (pk *PacketBuffer) ExperimentOptionValue() (uint16, bool) {
+	switch pk.NetworkProtocolNumber {
+	case header.IPv4ProtocolNumber:
+		h := header.IPv4(pk.NetworkHeader().Slice())
+		opts := h.Options()
+		iter := opts.MakeIterator()
+		for {
+			opt, done, err := iter.Next()
+			if err != nil {
+				return 0, false
+			}
+			if done {
+				return 0, false
+			}
+			if opt.Type() == header.IPv4OptionExperimentType {
+				return opt.(*header.IPv4OptionExperiment).Value(), true
+			}
+		}
+	case header.IPv6ProtocolNumber:
+		h := header.IPv6(pk.NetworkHeader().Slice())
+		v := pk.NetworkHeader().View()
+		if v != nil {
+			v.TrimFront(header.IPv6MinimumSize)
+		}
+		buf := buffer.MakeWithView(v)
+		buf.Append(pk.TransportHeader().View())
+		dataBuf := pk.Data().ToBuffer()
+		buf.Merge(&dataBuf)
+		it := header.MakeIPv6PayloadIterator(header.IPv6ExtensionHeaderIdentifier(h.NextHeader()), buf)
+
+		for {
+			hdr, done, err := it.Next()
+			if done || err != nil {
+				break
+			}
+			if h, ok := hdr.(header.IPv6ExperimentExtHdr); ok {
+				hdr.Release()
+				return h.Value, true
+			}
+			hdr.Release()
+		}
+	default:
+		panic(fmt.Sprintf("Unexpected network protocol number %d", pk.NetworkProtocolNumber))
+	}
+	return 0, false
+}

pkg/tcpip/tcpip.go

@@ -712,6 +712,10 @@ type ReadOptions struct {
 	// NeedLinkPacketInfo indicates whether to return the link-layer information,
 	// if supported.
 	NeedLinkPacketInfo bool
+
+	// NeedRecvdExperimentOption indicates whether to return the experiment
+	// option value from the last received packet, if supported.
+	NeedReceivedExperimentOption bool
 }
 
 // ReadResult represents result for a successful Endpoint.Read.
@@ -732,6 +736,10 @@ type ReadResult struct {
 	// LinkPacketInfo is the link-layer information of the received packet if
 	// ReadOptions.NeedLinkPacketInfo is true.
 	LinkPacketInfo LinkPacketInfo
+
+	// ReceivedExperimentOption is the experiment option value from the last
+	// received packet if ReadOptions.NeedReceivedExperimentOption is true.
+	ReceivedExperimentOption uint16
 }
 
 // Endpoint is the interface implemented by transport protocols (e.g., tcp, udp)

pkg/tcpip/transport/tcp/forwarder.go

@@ -15,9 +15,6 @@
 package tcp
 
 import (
-	"fmt"
-
-	"gvisor.dev/gvisor/pkg/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/seqnum"
@@ -179,49 +176,5 @@ func (r *ForwarderRequest) CreateEndpoint(queue *waiter.Queue) (tcpip.Endpoint,
 func (r *ForwarderRequest) ForwardedPacketExperimentOption() (uint16, bool) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
-
-	switch r.segment.pkt.NetworkProtocolNumber {
-	case header.IPv4ProtocolNumber:
-		h := header.IPv4(r.segment.pkt.NetworkHeader().Slice())
-		opts := h.Options()
-		iter := opts.MakeIterator()
-		for {
-			opt, done, err := iter.Next()
-			if err != nil {
-				return 0, false
-			}
-			if done {
-				return 0, false
-			}
-			if opt.Type() == header.IPv4OptionExperimentType {
-				return opt.(*header.IPv4OptionExperiment).Value(), true
-			}
-		}
-	case header.IPv6ProtocolNumber:
-		h := header.IPv6(r.segment.pkt.NetworkHeader().Slice())
-		v := r.segment.pkt.NetworkHeader().View()
-		if v != nil {
-			v.TrimFront(header.IPv6MinimumSize)
-		}
-		buf := buffer.MakeWithView(v)
-		buf.Append(r.segment.pkt.TransportHeader().View())
-		dataBuf := r.segment.pkt.Data().ToBuffer()
-		buf.Merge(&dataBuf)
-		it := header.MakeIPv6PayloadIterator(header.IPv6ExtensionHeaderIdentifier(h.NextHeader()), buf)
-
-		for {
-			hdr, done, err := it.Next()
-			if done || err != nil {
-				break
-			}
-			if h, ok := hdr.(header.IPv6ExperimentExtHdr); ok {
-				hdr.Release()
-				return h.Value, true
-			}
-			hdr.Release()
-		}
-	default:
-		panic(fmt.Sprintf("Unexpected network protocol number %d", r.segment.pkt.NetworkProtocolNumber))
-	}
-	return 0, false
+	return r.segment.pkt.ExperimentOptionValue()
 }

pkg/tcpip/transport/udp/endpoint.go

@@ -294,6 +294,10 @@ func (e *endpoint) Read(dst io.Writer, opts tcpip.ReadOptions) (tcpip.ReadResult
 	if opts.NeedRemoteAddr {
 		res.RemoteAddr = p.senderAddress
 	}
+	if opts.NeedReceivedExperimentOption {
+		expOptVal, _ := p.pkt.ExperimentOptionValue()
+		res.ReceivedExperimentOption = expOptVal
+	}
 
 	n, err := p.pkt.Data().ReadTo(dst, opts.Peek)
 	if n == 0 && err != nil {