Merge pull request #171 from malbeclabs/snor/pin-source

feat: support setting ICMP source address for outgoing packets

Author: SuperQ

packetconn.go

@@ -22,6 +22,7 @@ type packetConn interface {
 	SetDoNotFragment() error
 	SetBroadcastFlag() error
 	SetIfIndex(ifIndex int)
+	SetSource(source net.IP)
 	SetTrafficClass(uint8) error
 	InstallICMPIDFilter(id int) error
 }
@@ -30,6 +31,7 @@ type icmpConn struct {
 	c       *icmp.PacketConn
 	ttl     int
 	ifIndex int
+	source  net.IP
 }
 
 func (c *icmpConn) Close() error {
@@ -44,6 +46,10 @@ func (c *icmpConn) SetIfIndex(ifIndex int) {
 	c.ifIndex = ifIndex
 }
 
+func (c *icmpConn) SetSource(source net.IP) {
+	c.source = source
+}
+
 func (c *icmpConn) SetReadDeadline(t time.Time) error {
 	return c.c.SetReadDeadline(t)
 }
@@ -86,6 +92,13 @@ func (c *icmpv4Conn) WriteTo(b []byte, dst net.Addr) (int, error) {
 		cm = &ipv4.ControlMessage{IfIndex: c.ifIndex}
 	}
 
+	if c.source != nil {
+		if cm == nil {
+			cm = &ipv4.ControlMessage{}
+		}
+		cm.Src = c.source
+	}
+
 	return c.c.IPv4PacketConn().WriteTo(b, cm, dst)
 }
 
@@ -131,6 +144,13 @@ func (c *icmpV6Conn) WriteTo(b []byte, dst net.Addr) (int, error) {
 		cm = &ipv6.ControlMessage{IfIndex: c.ifIndex}
 	}
 
+	if c.source != nil {
+		if cm == nil {
+			cm = &ipv6.ControlMessage{}
+		}
+		cm.Src = c.source
+	}
+
 	return c.c.IPv6PacketConn().WriteTo(b, cm, dst)
 }
 

ping.go

@@ -565,6 +565,15 @@ func (p *Pinger) RunWithContext(ctx context.Context) error {
 		}
 		conn.SetIfIndex(iface.Index)
 	}
+
+	if p.Source != "" {
+		ip := net.ParseIP(p.Source)
+		if ip == nil {
+			return fmt.Errorf("invalid source address: %s", p.Source)
+		}
+		conn.SetSource(ip)
+	}
+
 	return p.run(ctx, conn)
 }
 

ping_test.go

@@ -498,6 +498,22 @@ func TestSetInterfaceName(t *testing.T) {
 	AssertError(t, err, "device not found")
 }
 
+func TestSetSource(t *testing.T) {
+	pinger := New("localhost")
+	pinger.Count = 1
+	pinger.Timeout = time.Second
+
+	// Set source address
+	pinger.Source = "127.0.0.1"
+	err := pinger.Run()
+	AssertNoError(t, err)
+
+	// Set source to invalid IP
+	pinger.Source = "invalid-ip"
+	err = pinger.Run()
+	AssertError(t, err, "invalid source address: invalid-ip")
+}
+
 // Test helpers
 func makeTestPinger() *Pinger {
 	pinger := New("127.0.0.1")
@@ -679,6 +695,7 @@ func (c testPacketConn) SetDoNotFragment() error           { return nil }
 func (c testPacketConn) SetBroadcastFlag() error           { return nil }
 func (c testPacketConn) InstallICMPIDFilter(id int) error  { return nil }
 func (c testPacketConn) SetIfIndex(ifIndex int)            {}
+func (c testPacketConn) SetSource(source net.IP)           {}
 func (c testPacketConn) SetTrafficClass(uint8) error       { return nil }
 
 func (c testPacketConn) ReadFrom(b []byte) (n int, ttl int, src net.Addr, err error) {