Add a Control function to manage raw network connections.

packetconn.go

@@ -1,8 +1,10 @@
 package probing
 
 import (
+	"fmt"
 	"net"
 	"runtime"
+	"syscall"
 	"time"
 
 	"golang.org/x/net/icmp"
@@ -24,6 +26,7 @@ type packetConn interface {
 	SetIfIndex(ifIndex int)
 	SetTrafficClass(uint8) error
 	InstallICMPIDFilter(id int) error
+	Control(f func(uintptr)) error
 }
 
 type icmpConn struct {
@@ -93,6 +96,22 @@ func (c icmpv4Conn) ICMPRequestType() icmp.Type {
 	return ipv4.ICMPTypeEcho
 }
 
+// Control invokes f on the underlying connection's file
+// descriptor or handle.
+func (c *icmpv4Conn) Control(f func(uintptr)) error {
+	syscallConn, ok := c.c.IPv4PacketConn().PacketConn.(syscall.Conn)
+	if ok {
+		sysConn, err := syscallConn.SyscallConn()
+		if err != nil {
+			return err
+		}
+
+		return sysConn.Control(f)
+	}
+
+	return fmt.Errorf("can't find SyscallConn method")
+}
+
 type icmpV6Conn struct {
 	icmpConn
 }
@@ -137,3 +156,19 @@ func (c *icmpV6Conn) WriteTo(b []byte, dst net.Addr) (int, error) {
 func (c icmpV6Conn) ICMPRequestType() icmp.Type {
 	return ipv6.ICMPTypeEchoRequest
 }
+
+// Control invokes f on the underlying connection's file
+// descriptor or handle.
+func (c *icmpV6Conn) Control(f func(uintptr)) error {
+	syscallConn, ok := c.c.IPv6PacketConn().PacketConn.(syscall.Conn)
+	if ok {
+		sysConn, err := syscallConn.SyscallConn()
+		if err != nil {
+			return err
+		}
+
+		return sysConn.Control(f)
+	}
+
+	return fmt.Errorf("can't find SyscallConn method")
+}

ping.go

@@ -221,6 +221,9 @@ type Pinger struct {
 	// mark is a SO_MARK (fwmark) set on outgoing icmp packets
 	mark uint
 
+	// Control the underlying connection's file descriptor or handle.
+	Control func(fd uintptr)
+
 	// df when true sets the do-not-fragment bit in the outer IP or IPv6 header
 	df bool
 
@@ -545,6 +548,12 @@ func (p *Pinger) RunWithContext(ctx context.Context) error {
 		}
 	}
 
+	if p.Control != nil {
+		if err := conn.Control(p.Control); err != nil {
+			return fmt.Errorf("error setting control: %v", err)
+		}
+	}
+
 	if p.df {
 		if err := conn.SetDoNotFragment(); err != nil {
 			return fmt.Errorf("error setting do-not-fragment: %v", err)

ping_test.go

@@ -675,6 +675,7 @@ 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) Control(f func(uintptr)) error     { return nil }
 func (c testPacketConn) SetDoNotFragment() error           { return nil }
 func (c testPacketConn) SetBroadcastFlag() error           { return nil }
 func (c testPacketConn) InstallICMPIDFilter(id int) error  { return nil }
@@ -878,3 +879,12 @@ func TestRunStatisticsConcurrent(t *testing.T) {
 	go p.Statistics()
 	p.Run()
 }
+
+func TestControl(t *testing.T) {
+	p := New("127.0.0.1")
+	err := p.Resolve()
+	AssertNoError(t, err)
+	p.Count = 1
+	p.Control = func(fd uintptr) {}
+	p.Run()
+}