From f1d34a1fd6f318cd686cb7bb39eacb10e3415ba7 Mon Sep 17 00:00:00 2001
From: Maximilian Fridrich <m.fridrich@commend.com>
Date: Tue, 26 Nov 2024 07:45:40 +0100
Subject: [PATCH] tcp,udp: set TOS (TCLASS) for IPv6 sockets

---
 src/tcp/tcp.c | 32 ++++++++++++++++++++++++++++++--
 src/udp/udp.c | 16 +++++++++++++++-
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/src/tcp/tcp.c b/src/tcp/tcp.c
index 889713e51..9c2f4ac10 100644
--- a/src/tcp/tcp.c
+++ b/src/tcp/tcp.c
@@ -1450,12 +1450,25 @@ int tcp_settos(struct tcp_sock *ts, uint32_t tos)
 {
 	int err = 0;
 	int v = tos;
+	struct sa sa;
 
 	if (!ts)
 		return EINVAL;
 
 	ts->tos = tos;
-	err = tcp_sock_setopt(ts, IPPROTO_IP, IP_TOS, &v, sizeof(v));
+	err = tcp_local_get(ts, &sa);
+	if (err)
+		return err;
+
+	if (sa_af(&sa) == AF_INET) {
+		err = tcp_sock_setopt(ts, IPPROTO_IP, IP_TOS, &v, sizeof(v));
+	}
+#if defined(IPV6_TCLASS) && !defined(WIN32)
+	else  if (sa_af(&sa) == AF_INET6) {
+		err = tcp_sock_setopt(ts, IPPROTO_IPV6, IPV6_TCLASS, &v,
+				      sizeof(v));
+	}
+#endif
 
 	return err;
 }
@@ -1465,16 +1478,31 @@ int tcp_conn_settos(struct tcp_conn *tc, uint32_t tos)
 {
 	int err = 0;
 	int v = tos;
+	struct sa sa;
 
 	if (!tc)
 		return EINVAL;
 
 	tc->tos = tos;
-	if (tc->fdc != RE_BAD_SOCK) {
+	if (tc->fdc == RE_BAD_SOCK)
+		return err;
+
+	err = tcp_conn_local_get(tc, &sa);
+	if (err)
+		return err;
+
+	if (sa_af(&sa) == AF_INET) {
 		if (0 != setsockopt(tc->fdc, IPPROTO_IP, IP_TOS,
 					BUF_CAST &v, sizeof(v)))
 			err = RE_ERRNO_SOCK;
 	}
+#if defined(IPV6_TCLASS) && !defined(WIN32)
+	else  if (sa_af(&sa) == AF_INET6) {
+		if (0 != setsockopt(tc->fdc, IPPROTO_IPV6, IPV6_TCLASS,
+					BUF_CAST &v, sizeof(v)))
+			err = RE_ERRNO_SOCK;
+	}
+#endif
 
 	return err;
 }
diff --git a/src/udp/udp.c b/src/udp/udp.c
index 7a7f9cb52..8bfbca67d 100644
--- a/src/udp/udp.c
+++ b/src/udp/udp.c
@@ -628,6 +628,7 @@ int udp_settos(struct udp_sock *us, uint8_t tos)
 {
 	int err = 0;
 	int v = tos;
+	struct sa sa;
 #ifdef WIN32
 	QOS_VERSION qos_version = { 1 , 0 };
 	QOS_TRAFFIC_TYPE qos_type = QOSTrafficTypeBestEffort;
@@ -660,7 +661,20 @@ int udp_settos(struct udp_sock *us, uint8_t tos)
 			return WSAGetLastError();
 	}
 #endif
-	err = udp_setsockopt(us, IPPROTO_IP, IP_TOS, &v, sizeof(v));
+	err = udp_local_get(us, &sa);
+	if (err)
+		return err;
+
+	if (sa_af(&sa) == AF_INET) {
+		err = udp_setsockopt(us, IPPROTO_IP, IP_TOS, &v, sizeof(v));
+	}
+#if defined(IPV6_TCLASS) && !defined(WIN32)
+	else if (sa_af(&sa) == AF_INET6) {
+		err = udp_setsockopt(us, IPPROTO_IPV6, IPV6_TCLASS, &v,
+				     sizeof(v));
+	}
+#endif
+
 	return err;
 }