From 78590351ea4fa67bc2a51728b91800c8e11dfe4b Mon Sep 17 00:00:00 2001
From: xtaci <daniel820313@gmail.com>
Date: Thu, 8 Aug 2019 12:29:23 +0000
Subject: [PATCH] let SetDSCP invoke the conn's SetDSCP function if there is
 any

---
 sess.go | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/sess.go b/sess.go
index de86ffe..9187785 100644
--- a/sess.go
+++ b/sess.go
@@ -454,6 +454,9 @@ func (s *UDPSession) SetNoDelay(nodelay, interval, resend, nc int) {
 
 // SetDSCP sets the 6bit DSCP field in IPv4 header, or 8bit Traffic Class in IPv6 header.
 //
+// if the underlying connection has implemented `func SetDSCP(int) error`, SetDSCP() will invoke
+// this function instead.
+//
 // It has no effect if it's accepted from Listener.
 func (s *UDPSession) SetDSCP(dscp int) error {
 	s.mu.Lock()
@@ -461,6 +464,14 @@ func (s *UDPSession) SetDSCP(dscp int) error {
 	if s.l != nil {
 		return errInvalidOperation
 	}
+
+	// interface enabled
+	if ts, ok := s.conn.(interface {
+		SetDSCP(int) error
+	}); ok {
+		return ts.SetDSCP(dscp)
+	}
+
 	if nc, ok := s.conn.(net.Conn); ok {
 		var succeed bool
 		if err := ipv4.NewConn(nc).SetTOS(dscp << 2); err == nil {
@@ -813,7 +824,17 @@ func (l *Listener) SetWriteBuffer(bytes int) error {
 }
 
 // SetDSCP sets the 6bit DSCP field in IPv4 header, or 8bit Traffic Class in IPv6 header.
+//
+// if the underlying connection has implemented `func SetDSCP(int) error`, SetDSCP() will invoke
+// this function instead.
 func (l *Listener) SetDSCP(dscp int) error {
+	// interface enabled
+	if ts, ok := l.conn.(interface {
+		SetDSCP(int) error
+	}); ok {
+		return ts.SetDSCP(dscp)
+	}
+
 	if nc, ok := l.conn.(net.Conn); ok {
 		var succeed bool
 		if err := ipv4.NewConn(nc).SetTOS(dscp << 2); err == nil {