6 files changed, 42 insertions, 5 deletions
diff --git a/pkg/tcpip/stack/nic.go b/pkg/tcpip/stack/nic.go
index ddd014658..a4556674b 100644
--- a/pkg/tcpip/stack/nic.go
+++ b/pkg/tcpip/stack/nic.go
@@ -575,7 +575,7 @@ func (n *NIC) RemoveAddressRange(subnet tcpip.Subnet) {
 	n.mu.Unlock()
 }
 
-// Subnets returns the Subnets associated with this NIC.
+// AddressRanges returns the Subnets associated with this NIC.
 func (n *NIC) AddressRanges() []tcpip.Subnet {
 	n.mu.RLock()
 	defer n.mu.RUnlock()
diff --git a/pkg/tcpip/stack/stack.go b/pkg/tcpip/stack/stack.go
index 7a9600679..251336224 100644
--- a/pkg/tcpip/stack/stack.go
+++ b/pkg/tcpip/stack/stack.go
@@ -829,7 +829,7 @@ func (s *Stack) CheckNIC(id tcpip.NICID) bool {
 	return false
 }
 
-// NICSubnets returns a map of NICIDs to their associated subnets.
+// NICAddressRanges returns a map of NICIDs to their associated subnets.
 func (s *Stack) NICAddressRanges() map[tcpip.NICID][]tcpip.Subnet {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
diff --git a/pkg/tcpip/tcpip.go b/pkg/tcpip/tcpip.go
index f62fd729f..5c7b2af88 100644
--- a/pkg/tcpip/tcpip.go
+++ b/pkg/tcpip/tcpip.go
@@ -322,7 +322,7 @@ type ControlMessages struct {
 	HasTOS bool
 
 	// TOS is the IPv4 type of service of the associated packet.
-	TOS int8
+	TOS uint8
 
 	// HasTClass indicates whether Tclass is valid/set.
 	HasTClass bool
@@ -666,6 +666,10 @@ type IPv4TOSOption uint8
 // for all subsequent outgoing IPv6 packets from the endpoint.
 type IPv6TrafficClassOption uint8
 
+// ReceiveTOSOption is used by SetSockOpt/GetSockOpt to specify if the TOS
+// ancillary message is passed with incoming packets.
+type ReceiveTOSOption bool
+
 // Route is a row in the routing table. It specifies through which NIC (and
 // gateway) sets of packets should be routed. A row is considered viable if the
 // masked target address matches the destination address in the row.
diff --git a/pkg/tcpip/transport/raw/endpoint.go b/pkg/tcpip/transport/raw/endpoint.go
index 5aafe2615..6d23ab5a1 100644
--- a/pkg/tcpip/transport/raw/endpoint.go
+++ b/pkg/tcpip/transport/raw/endpoint.go
@@ -510,7 +510,7 @@ func (e *endpoint) SetSockOpt(opt interface{}) *tcpip.Error {
 }
 
 // SetSockOptInt implements tcpip.Endpoint.SetSockOptInt.
-func (ep *endpoint) SetSockOptInt(opt tcpip.SockOpt, v int) *tcpip.Error {
+func (e *endpoint) SetSockOptInt(opt tcpip.SockOpt, v int) *tcpip.Error {
 	return tcpip.ErrUnknownProtocolOption
 }
 
diff --git a/pkg/tcpip/transport/udp/endpoint.go b/pkg/tcpip/transport/udp/endpoint.go
index 1ac4705af..269470ed4 100644
--- a/pkg/tcpip/transport/udp/endpoint.go
+++ b/pkg/tcpip/transport/udp/endpoint.go
@@ -32,6 +32,7 @@ type udpPacket struct {
 	senderAddress tcpip.FullAddress
 	data          buffer.VectorisedView `state:".(buffer.VectorisedView)"`
 	timestamp     int64
+	tos           uint8
 }
 
 // EndpointState represents the state of a UDP endpoint.
@@ -114,6 +115,10 @@ type endpoint struct {
 	// applied while sending packets. Defaults to 0 as on Linux.
 	sendTOS uint8
 
+	// receiveTOS determines if the incoming IPv4 TOS header field is passed
+	// as ancillary data to ControlMessages on Read.
+	receiveTOS bool
+
 	// shutdownFlags represent the current shutdown state of the endpoint.
 	shutdownFlags tcpip.ShutdownFlags
 
@@ -244,7 +249,12 @@ func (e *endpoint) Read(addr *tcpip.FullAddress) (buffer.View, tcpip.ControlMess
 		*addr = p.senderAddress
 	}
 
-	return p.data.ToView(), tcpip.ControlMessages{HasTimestamp: true, Timestamp: p.timestamp}, nil
+	return p.data.ToView(), tcpip.ControlMessages{
+		HasTimestamp: true,
+		Timestamp:    p.timestamp,
+		HasTOS:       e.receiveTOS,
+		TOS:          p.tos,
+	}, nil
 }
 
 // prepareForWrite prepares the endpoint for sending data. In particular, it
@@ -656,6 +666,12 @@ func (e *endpoint) SetSockOpt(opt interface{}) *tcpip.Error {
 		e.sendTOS = uint8(v)
 		e.mu.Unlock()
 		return nil
+
+	case tcpip.ReceiveTOSOption:
+		e.mu.Lock()
+		e.receiveTOS = bool(v)
+		e.mu.Unlock()
+		return nil
 	}
 	return nil
 }
@@ -792,6 +808,12 @@ func (e *endpoint) GetSockOpt(opt interface{}) *tcpip.Error {
 		e.mu.RUnlock()
 		return nil
 
+	case *tcpip.ReceiveTOSOption:
+		e.mu.RLock()
+		*o = tcpip.ReceiveTOSOption(e.receiveTOS)
+		e.mu.RUnlock()
+		return nil
+
 	default:
 		return tcpip.ErrUnknownProtocolOption
 	}
@@ -1238,6 +1260,13 @@ func (e *endpoint) HandlePacket(r *stack.Route, id stack.TransportEndpointID, pk
 	e.rcvList.PushBack(packet)
 	e.rcvBufSize += pkt.Data.Size()
 
+	// Save any useful information from the NetworkHeader to the packet.
+	switch r.NetProto {
+	case header.IPv4ProtocolNumber:
+		// This packet has already been validated before being passed up the stack.
+		packet.tos, _ = header.IPv4(pkt.NetworkHeader).TOS()
+	}
+
 	packet.timestamp = e.stack.NowNanoseconds()
 
 	e.rcvMu.Unlock()
diff --git a/pkg/tcpip/transport/udp/udp_state_autogen.go b/pkg/tcpip/transport/udp/udp_state_autogen.go
index 5c4b43214..75b78fa01 100755
--- a/pkg/tcpip/transport/udp/udp_state_autogen.go
+++ b/pkg/tcpip/transport/udp/udp_state_autogen.go
@@ -15,6 +15,7 @@ func (x *udpPacket) save(m state.Map) {
 	m.Save("udpPacketEntry", &x.udpPacketEntry)
 	m.Save("senderAddress", &x.senderAddress)
 	m.Save("timestamp", &x.timestamp)
+	m.Save("tos", &x.tos)
 }
 
 func (x *udpPacket) afterLoad() {}
@@ -22,6 +23,7 @@ func (x *udpPacket) load(m state.Map) {
 	m.Load("udpPacketEntry", &x.udpPacketEntry)
 	m.Load("senderAddress", &x.senderAddress)
 	m.Load("timestamp", &x.timestamp)
+	m.Load("tos", &x.tos)
 	m.LoadValue("data", new(buffer.VectorisedView), func(y interface{}) { x.loadData(y.(buffer.VectorisedView)) })
 }
 
@@ -51,6 +53,7 @@ func (x *endpoint) save(m state.Map) {
 	m.Save("boundBindToDevice", &x.boundBindToDevice)
 	m.Save("boundPortFlags", &x.boundPortFlags)
 	m.Save("sendTOS", &x.sendTOS)
+	m.Save("receiveTOS", &x.receiveTOS)
 	m.Save("shutdownFlags", &x.shutdownFlags)
 	m.Save("multicastMemberships", &x.multicastMemberships)
 	m.Save("effectiveNetProtos", &x.effectiveNetProtos)
@@ -79,6 +82,7 @@ func (x *endpoint) load(m state.Map) {
 	m.Load("boundBindToDevice", &x.boundBindToDevice)
 	m.Load("boundPortFlags", &x.boundPortFlags)
 	m.Load("sendTOS", &x.sendTOS)
+	m.Load("receiveTOS", &x.receiveTOS)
 	m.Load("shutdownFlags", &x.shutdownFlags)
 	m.Load("multicastMemberships", &x.multicastMemberships)
 	m.Load("effectiveNetProtos", &x.effectiveNetProtos)