Add support for setting SO_MARK socket option on outgoing ICMP packets

Signed-off-by: Jeremiah Millay <jmillay@fastly.com>

README.md

@@ -111,6 +111,10 @@ See [this blog](https://sturmflut.github.io/linux/ubuntu/2015/01/17/unprivileged
 and the Go [x/net/icmp](https://godoc.org/golang.org/x/net/icmp) package
 for more details.
 
+This library also supports setting the `SO_MARK` socket option which is equivalent to the `-m mark`
+flag in standard ping binaries on linux. Setting this option requires the `CAP_NET_ADMIN` capability
+(via `setcap` or elevated privileges). You can set a mark (ex: 100) with `pinger.SetMark(100)` in your code.
+
 ### Windows
 
 You must use `pinger.SetPrivileged(true)`, otherwise you will receive

packetconn.go

@@ -18,6 +18,7 @@ type packetConn interface {
 	SetReadDeadline(t time.Time) error
 	WriteTo(b []byte, dst net.Addr) (int, error)
 	SetTTL(ttl int)
+	SetMark(m uint) error
 	SetTOS(tos int)
 }
 

ping.go

@@ -189,6 +189,9 @@ type Pinger struct {
 	ipaddr *net.IPAddr
 	addr   string
 
+	// Mark is a mark (fwmark) set on outgoing icmp packets
+	mark uint
+
 	// trackerUUIDs is the list of UUIDs being used for sending packets.
 	trackerUUIDs []uuid.UUID
 
@@ -403,6 +406,16 @@ func (p *Pinger) ID() int {
 	return p.id
 }
 
+// SetMark sets a mark intended to be set on outgoing ICMP packets.
+func (p *Pinger) SetMark(m uint) {
+	p.mark = m
+}
+
+// Mark returns the mark to be set on outgoing ICMP packets.
+func (p *Pinger) Mark() uint {
+	return p.mark
+}
+
 // Run runs the pinger. This is a blocking function that will exit when it's
 // done. If Count or Interval are not specified, it will run continuously until
 // it is interrupted.
@@ -423,6 +436,12 @@ func (p *Pinger) Run() error {
 	}
 	defer conn.Close()
 
+	if p.mark != 0 {
+		if err := conn.SetMark(p.mark); err != nil {
+			return fmt.Errorf("error setting mark: %v", err)
+		}
+	}
+
 	conn.SetTTL(p.TTL)
 	conn.SetTOS(p.TOS)
 	return p.run(conn)

ping_test.go

@@ -658,6 +658,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) SetMark(m uint) error              { return nil }
 func (c testPacketConn) SetTOS(t int)                      {}
 
 func (c testPacketConn) ReadFrom(b []byte) (n int, ttl int, src net.Addr, err error) {

utils_linux.go

@@ -3,6 +3,15 @@
 
 package ping
 
+import (
+	"errors"
+	"os"
+	"reflect"
+	"syscall"
+
+	"golang.org/x/net/icmp"
+)
+
 // Returns the length of an ICMP message.
 func (p *Pinger) getMessageLength() int {
 	return p.Size + 8
@@ -12,9 +21,66 @@ func (p *Pinger) getMessageLength() int {
 func (p *Pinger) matchID(ID int) bool {
 	// On Linux we can only match ID if we are privileged.
 	if p.protocol == "icmp" {
-		if ID != p.id {
-			return false
-		}
+		return ID == p.id
 	}
 	return true
 }
+
+// SetMark sets the SO_MARK socket option on outgoing ICMP packets.
+// Setting this option requires CAP_NET_ADMIN.
+func (c *icmpConn) SetMark(mark uint) error {
+	fd, err := getFD(c.c)
+	if err != nil {
+		return err
+	}
+	return os.NewSyscallError(
+		"setsockopt",
+		syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_MARK, int(mark)),
+	)
+}
+
+// SetMark sets the SO_MARK socket option on outgoing ICMP packets.
+// Setting this option requires CAP_NET_ADMIN.
+func (c *icmpv4Conn) SetMark(mark uint) error {
+	fd, err := getFD(c.icmpConn.c)
+	if err != nil {
+		return err
+	}
+	return os.NewSyscallError(
+		"setsockopt",
+		syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_MARK, int(mark)),
+	)
+}
+
+// SetMark sets the SO_MARK socket option on outgoing ICMP packets.
+// Setting this option requires CAP_NET_ADMIN.
+func (c *icmpV6Conn) SetMark(mark uint) error {
+	fd, err := getFD(c.icmpConn.c)
+	if err != nil {
+		return err
+	}
+	return os.NewSyscallError(
+		"setsockopt",
+		syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_MARK, int(mark)),
+	)
+}
+
+// getFD gets the system file descriptor for an icmp.PacketConn
+func getFD(c *icmp.PacketConn) (uintptr, error) {
+	v := reflect.ValueOf(c).Elem().FieldByName("c").Elem()
+	if v.Elem().Kind() != reflect.Struct {
+		return 0, errors.New("invalid type")
+	}
+
+	fd := v.Elem().FieldByName("conn").FieldByName("fd")
+	if fd.Elem().Kind() != reflect.Struct {
+		return 0, errors.New("invalid type")
+	}
+
+	pfd := fd.Elem().FieldByName("pfd")
+	if pfd.Kind() != reflect.Struct {
+		return 0, errors.New("invalid type")
+	}
+
+	return uintptr(pfd.FieldByName("Sysfd").Int()), nil
+}

utils_other.go

@@ -3,6 +3,10 @@
 
 package ping
 
+import (
+	"errors"
+)
+
 // Returns the length of an ICMP message.
 func (p *Pinger) getMessageLength() int {
 	return p.Size + 8
@@ -10,8 +14,23 @@ func (p *Pinger) getMessageLength() int {
 
 // Attempts to match the ID of an ICMP packet.
 func (p *Pinger) matchID(ID int) bool {
-	if ID != p.id {
-		return false
-	}
-	return true
+	return ID == p.id
+}
+
+// SetMark sets the SO_MARK socket option on outgoing ICMP packets.
+// Setting this option requires CAP_NET_ADMIN.
+func (c *icmpConn) SetMark(mark uint) error {
+	return errors.New("setting SO_MARK socket option is not supported on this platform")
+}
+
+// SetMark sets the SO_MARK socket option on outgoing ICMP packets.
+// Setting this option requires CAP_NET_ADMIN.
+func (c *icmpv4Conn) SetMark(mark uint) error {
+	return errors.New("setting SO_MARK socket option is not supported on this platform")
+}
+
+// SetMark sets the SO_MARK socket option on outgoing ICMP packets.
+// Setting this option requires CAP_NET_ADMIN.
+func (c *icmpV6Conn) SetMark(mark uint) error {
+	return errors.New("setting SO_MARK socket option is not supported on this platform")
 }

utils_windows.go

@@ -4,6 +4,8 @@
 package ping
 
 import (
+	"errors"
+
 	"golang.org/x/net/ipv4"
 	"golang.org/x/net/ipv6"
 )
@@ -23,3 +25,21 @@ func (p *Pinger) matchID(ID int) bool {
 	}
 	return true
 }
+
+// SetMark sets the SO_MARK socket option on outgoing ICMP packets.
+// Setting this option requires CAP_NET_ADMIN.
+func (c *icmpConn) SetMark(mark uint) error {
+	return errors.New("setting SO_MARK socket option is not supported on this platform")
+}
+
+// SetMark sets the SO_MARK socket option on outgoing ICMP packets.
+// Setting this option requires CAP_NET_ADMIN.
+func (c *icmpv4Conn) SetMark(mark uint) error {
+	return errors.New("setting SO_MARK socket option is not supported on this platform")
+}
+
+// SetMark sets the SO_MARK socket option on outgoing ICMP packets.
+// Setting this option requires CAP_NET_ADMIN.
+func (c *icmpV6Conn) SetMark(mark uint) error {
+	return errors.New("setting SO_MARK socket option is not supported on this platform")
+}