2 files changed, 23 insertions, 12 deletions
diff --git a/pkg/tcpip/transport/internal/network/endpoint.go b/pkg/tcpip/transport/internal/network/endpoint.go
index 0dce60d89..c5b575e1c 100644
--- a/pkg/tcpip/transport/internal/network/endpoint.go
+++ b/pkg/tcpip/transport/internal/network/endpoint.go
@@ -60,10 +60,8 @@ type Endpoint struct {
 	multicastAddr tcpip.Address
 	// TODO(https://gvisor.dev/issue/6389): Use different fields for IPv4/IPv6.
 	multicastNICID tcpip.NICID
-	// sendTOS represents IPv4 TOS or IPv6 TrafficClass,
-	// applied while sending packets. Defaults to 0 as on Linux.
-	// TODO(https://gvisor.dev/issue/6389): Use different fields for IPv4/IPv6.
-	sendTOS uint8
+	ipv4TOS        uint8
+	ipv6TClass     uint8
 }
 
 // +stateify savable
@@ -267,11 +265,21 @@ func (e *Endpoint) AcquireContextForWrite(opts tcpip.WriteOptions) (WriteContext
 		return WriteContext{}, &tcpip.ErrBroadcastDisabled{}
 	}
 
+	var tos uint8
+	switch netProto := route.NetProto(); netProto {
+	case header.IPv4ProtocolNumber:
+		tos = e.ipv4TOS
+	case header.IPv6ProtocolNumber:
+		tos = e.ipv6TClass
+	default:
+		panic(fmt.Sprintf("invalid protocol number = %d", netProto))
+	}
+
 	return WriteContext{
 		transProto: e.transProto,
 		route:      route,
 		ttl:        calculateTTL(route, e.ttl, e.multicastTTL),
-		tos:        e.sendTOS,
+		tos:        tos,
 		owner:      e.owner,
 	}, nil
 }
@@ -533,12 +541,12 @@ func (e *Endpoint) SetSockOptInt(opt tcpip.SockOptInt, v int) tcpip.Error {
 
 	case tcpip.IPv4TOSOption:
 		e.mu.Lock()
-		e.sendTOS = uint8(v)
+		e.ipv4TOS = uint8(v)
 		e.mu.Unlock()
 
 	case tcpip.IPv6TrafficClassOption:
 		e.mu.Lock()
-		e.sendTOS = uint8(v)
+		e.ipv6TClass = uint8(v)
 		e.mu.Unlock()
 	}
 
@@ -566,13 +574,13 @@ func (e *Endpoint) GetSockOptInt(opt tcpip.SockOptInt) (int, tcpip.Error) {
 
 	case tcpip.IPv4TOSOption:
 		e.mu.RLock()
-		v := int(e.sendTOS)
+		v := int(e.ipv4TOS)
 		e.mu.RUnlock()
 		return v, nil
 
 	case tcpip.IPv6TrafficClassOption:
 		e.mu.RLock()
-		v := int(e.sendTOS)
+		v := int(e.ipv6TClass)
 		e.mu.RUnlock()
 		return v, nil
 
diff --git a/pkg/tcpip/transport/internal/network/network_state_autogen.go b/pkg/tcpip/transport/internal/network/network_state_autogen.go
index 8f1cf9c0d..0ce695bb8 100644
--- a/pkg/tcpip/transport/internal/network/network_state_autogen.go
+++ b/pkg/tcpip/transport/internal/network/network_state_autogen.go
@@ -25,7 +25,8 @@ func (e *Endpoint) StateFields() []string {
 		"multicastTTL",
 		"multicastAddr",
 		"multicastNICID",
-		"sendTOS",
+		"ipv4TOS",
+		"ipv6TClass",
 	}
 }
 
@@ -47,7 +48,8 @@ func (e *Endpoint) StateSave(stateSinkObject state.Sink) {
 	stateSinkObject.Save(10, &e.multicastTTL)
 	stateSinkObject.Save(11, &e.multicastAddr)
 	stateSinkObject.Save(12, &e.multicastNICID)
-	stateSinkObject.Save(13, &e.sendTOS)
+	stateSinkObject.Save(13, &e.ipv4TOS)
+	stateSinkObject.Save(14, &e.ipv6TClass)
 }
 
 func (e *Endpoint) afterLoad() {}
@@ -67,7 +69,8 @@ func (e *Endpoint) StateLoad(stateSourceObject state.Source) {
 	stateSourceObject.Load(10, &e.multicastTTL)
 	stateSourceObject.Load(11, &e.multicastAddr)
 	stateSourceObject.Load(12, &e.multicastNICID)
-	stateSourceObject.Load(13, &e.sendTOS)
+	stateSourceObject.Load(13, &e.ipv4TOS)
+	stateSourceObject.Load(14, &e.ipv6TClass)
 }
 
 func (m *multicastMembership) StateTypeName() string {