From 63b43a38b659e653f83b756ef52adf2af4d9bfdf Mon Sep 17 00:00:00 2001
From: Mika Savela <msavela@basen.net>
Date: Mon, 19 Sep 2022 09:57:36 +0000
Subject: [PATCH] Add TOS support

---
 packetconn.go | 12 ++++++++++++
 ping.go       |  8 +++++++-
 ping_test.go  |  1 +
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/packetconn.go b/packetconn.go
index 6a972f5..3c302f0 100644
--- a/packetconn.go
+++ b/packetconn.go
@@ -18,11 +18,13 @@ type packetConn interface {
 	SetReadDeadline(t time.Time) error
 	WriteTo(b []byte, dst net.Addr) (int, error)
 	SetTTL(ttl int)
+	SetTOS(tos int)
 }
 
 type icmpConn struct {
 	c   *icmp.PacketConn
 	ttl int
+	tos int
 }
 
 func (c *icmpConn) Close() error {
@@ -33,6 +35,10 @@ func (c *icmpConn) SetTTL(ttl int) {
 	c.ttl = ttl
 }
 
+func (c *icmpConn) SetTOS(tos int) {
+	c.tos = tos
+}
+
 func (c *icmpConn) SetReadDeadline(t time.Time) error {
 	return c.c.SetReadDeadline(t)
 }
@@ -42,11 +48,17 @@ func (c *icmpConn) WriteTo(b []byte, dst net.Addr) (int, error) {
 		if err := c.c.IPv6PacketConn().SetHopLimit(c.ttl); err != nil {
 			return 0, err
 		}
+		if err := c.c.IPv6PacketConn().SetTrafficClass(c.tos); err != nil {
+			return 0, err
+		}
 	}
 	if c.c.IPv4PacketConn() != nil {
 		if err := c.c.IPv4PacketConn().SetTTL(c.ttl); err != nil {
 			return 0, err
 		}
+		if err := c.c.IPv4PacketConn().SetTOS(c.tos); err != nil {
+			return 0, err
+		}
 	}
 
 	return c.c.WriteTo(b, dst)
diff --git a/ping.go b/ping.go
index ef277ab..a5e09c0 100644
--- a/ping.go
+++ b/ping.go
@@ -49,7 +49,6 @@
 // it calls the OnFinish callback.
 //
 // For a full ping example, see "cmd/ping/ping.go".
-//
 package ping
 
 import (
@@ -107,6 +106,7 @@ func New(addr string) *Pinger {
 		protocol:          "udp",
 		awaitingSequences: firstSequence,
 		TTL:               64,
+		TOS:               0,
 		logger:            StdLogger{Logger: log.New(log.Writer(), log.Prefix(), log.Flags())},
 	}
 }
@@ -205,6 +205,8 @@ type Pinger struct {
 	logger Logger
 
 	TTL int
+
+	TOS int
 }
 
 type packet struct {
@@ -233,6 +235,9 @@ type Packet struct {
 	// TTL is the Time To Live on the packet.
 	Ttl int
 
+	// ToS is the Type of Service on the packet.
+	Tos int
+
 	// ID is the ICMP identifier.
 	ID int
 }
@@ -419,6 +424,7 @@ func (p *Pinger) Run() error {
 	defer conn.Close()
 
 	conn.SetTTL(p.TTL)
+	conn.SetTOS(p.TOS)
 	return p.run(conn)
 }
 
diff --git a/ping_test.go b/ping_test.go
index b8755e7..76cc45b 100644
--- a/ping_test.go
+++ b/ping_test.go
@@ -642,6 +642,7 @@ func (c testPacketConn) ICMPRequestType() icmp.Type        { return ipv4.ICMPTyp
 func (c testPacketConn) SetFlagTTL() error                 { return nil }
 func (c testPacketConn) SetReadDeadline(t time.Time) error { return nil }
 func (c testPacketConn) SetTTL(t int)                      {}
+func (c testPacketConn) SetTOS(t int)                      {}
 
 func (c testPacketConn) ReadFrom(b []byte) (n int, ttl int, src net.Addr, err error) {
 	return 0, 0, nil, nil