package net
 
 import (
 	"encoding/binary"
 	"fmt"
 	"net"
 	"os"
 	"strings"
 	"time"
 
 	"golang.org/x/net/icmp"
 	"golang.org/x/net/ipv4"
 	"golang.org/x/net/ipv6"
 )
 
 const (
 	/* Configuration */
 	defaultMaxHops = 30
 	defaultTimeout = 1 * time.Second
 	probesPerHop   = 3
 	echoPayload    = "tiny-traceroute-djnn-sh"
 
 	/* Network Protocols & Packet sizing */
 	protocolICMPv4   = 1
 	protocolICMPv6   = 58
 	readBufferSize   = 1500
 	minIPv4HeaderLen = 20
 	minIPv6HeaderLen = 40
 	icmpHeaderLen    = 8
 	ipv4WordSize     = 4 /* IHL is measured in 32-bit (4-byte) words */
 	ihlMask          = 0x0f
 	processIDMask    = 0xffff
 )
 
 func Run(target string, ipType string) {
 	destIP, err := net.ResolveIPAddr("ip", target)
 	if err != nil {
 		fmt.Printf("Error resolving host: %v\n", err)
 		os.Exit(1)
 	}
 
 	fmt.Printf("traceroute to %s (%s), %d hops max\n", target, destIP.String(), defaultMaxHops)
 
 	c, pc, err := setupSocket(ipType)
 	if err != nil {
 		fmt.Printf("Socket error: %v\n", err)
 		return
 	}
 	defer c.Close()
 
 	traceID := os.Getpid() & processIDMask
 
 	for ttl := 1; ttl <= defaultMaxHops; ttl++ {
 		fmt.Printf("%2d  ", ttl)
 
 		var prevAddr string
 		destinationReached := false
 
 		for probe := 0; probe < probesPerHop; probe++ {
 			seq := (ttl * 100) + probe
 
 			peerAddr, rtt, err := executeProbe(c, pc, destIP, traceID, seq, ttl, ipType)
 			if err != nil {
 				fmt.Printf("* ")
 			} else {
 				printProbeResult(peerAddr, prevAddr, rtt)
 				prevAddr = peerAddr
 
 				if peerAddr == destIP.String() {
 					destinationReached = true
 				}
 			}
 		}
 
 		fmt.Println()
 
 		if destinationReached {
 			break
 		}
 	}
 }
 
 func setupSocket(ipVersion string) (*icmp.PacketConn, interface{}, error) {
 	protocol := "ip4:icmp"
 	network := "0.0.0.0"
 	if ipVersion == "ip6" {
 		protocol = "ip6:ipv6-icmp"
 		network = "::"
 	}
 
 	c, err := icmp.ListenPacket(protocol, network)
 	if err != nil {
 		return nil, nil, fmt.Errorf("listen packet: %w (ensure sudo/admin)", err)
 	}
 
 	if ipVersion == "ip6" {
 		return c, c.IPv6PacketConn(), nil
 	}
 	return c, c.IPv4PacketConn(), nil
 }
 
 func executeProbe(c *icmp.PacketConn, pc interface{}, dest *net.IPAddr, id, seq, ttl int, ipType string) (string, time.Duration, error) {
 	// Set TTL based on IP version
 	if ipType == "ip6" {
 		if pc6, ok := pc.(*ipv6.PacketConn); ok {
 			if err := pc6.SetHopLimit(ttl); err != nil {
 				return "", 0, err
 			}
 		}
 	} else {
 		if pc4, ok := pc.(*ipv4.PacketConn); ok {
 			if err := pc4.SetTTL(ttl); err != nil {
 				return "", 0, err
 			}
 		}
 	}
 
 	start := time.Now()
 
 	if err := sendEchoRequest(c, dest, id, seq, ipType); err != nil {
 		return "", 0, err
 	}
 
 	peer, err := awaitResponse(c, pc, id, seq, ipType)
 	if err != nil {
 		return "", 0, err
 	}
 
 	return cleanIP(peer.String()), time.Since(start), nil
 }
 
 func sendEchoRequest(c *icmp.PacketConn, dest net.Addr, id, seq int, ipType string) error {
 	var msgType icmp.Type
 	if ipType == "ip6" {
 		msgType = ipv6.ICMPTypeEchoRequest
 	} else {
 		msgType = ipv4.ICMPTypeEcho
 	}
 
 	msg := icmp.Message{
 		Type: msgType, Code: 0,
 		Body: &icmp.Echo{
 			ID: id, Seq: seq,
 			Data: []byte(echoPayload),
 		},
 	}
 	b, err := msg.Marshal(nil)
 	if err != nil {
 		return err
 	}
 	_, err = c.WriteTo(b, dest)
 	return err
 }
 
 func awaitResponse(c *icmp.PacketConn, pc interface{}, expectedID, expectedSeq int, ipType string) (net.Addr, error) {
 	c.SetReadDeadline(time.Now().Add(defaultTimeout))
 	b := make([]byte, readBufferSize)
 
 	protocol := protocolICMPv4
 	if ipType == "ip6" {
 		protocol = protocolICMPv6
 	}
 
 	for {
 		n, peer, err := c.ReadFrom(b)
 		if err != nil {
 			return nil, err
 		}
 
 		rm, err := icmp.ParseMessage(protocol, b[:n])
 		if err != nil {
 			continue
 		}
 
 		// Handle based on IP version
 		if ipType == "ip6" {
 			switch rm.Type {
 			case ipv6.ICMPTypeEchoReply:
 				if echo, ok := rm.Body.(*icmp.Echo); ok {
 					if echo.ID == expectedID && echo.Seq == expectedSeq {
 						return peer, nil
 					}
 				}
 			case ipv6.ICMPTypeTimeExceeded:
 				if te, ok := rm.Body.(*icmp.TimeExceeded); ok {
 					if matchOriginalPacketV6(te.Data, expectedID, expectedSeq) {
 						return peer, nil
 					}
 				}
 			}
 		} else {
 			switch rm.Type {
 			case ipv4.ICMPTypeEchoReply:
 				if echo, ok := rm.Body.(*icmp.Echo); ok {
 					if echo.ID == expectedID && echo.Seq == expectedSeq {
 						return peer, nil
 					}
 				}
 			case ipv4.ICMPTypeTimeExceeded:
 				if te, ok := rm.Body.(*icmp.TimeExceeded); ok {
 					if matchOriginalPacket(te.Data, expectedID, expectedSeq) {
 						return peer, nil
 					}
 				}
 			}
 		}
 	}
 }
 
 func matchOriginalPacket(data []byte, expectedID, expectedSeq int) bool {
 	if len(data) < minIPv4HeaderLen {
 		return false
 	}
 
 	/* Calculate Internet Header Length (IHL) */
 	ihl := int(data[0]&ihlMask) * ipv4WordSize
 
 	/* We need the IP header + 8 bytes of the inner ICMP packet */
 	if len(data) < ihl+icmpHeaderLen {
 		return false
 	}
 
 	innerICMP := data[ihl:]
 	origID, origSeq := extractICMPIDAndSeq(innerICMP)
 	return origID == expectedID && origSeq == expectedSeq
 }
 
 func matchOriginalPacketV6(data []byte, expectedID, expectedSeq int) bool {
 	/* For IPv6, the TimeExceeded message contains the original IPv6 header +
 	   as much of the original packet as possible. IPv6 has a fixed 40-byte header */
 	if len(data) < minIPv6HeaderLen+icmpHeaderLen {
 		return false
 	}
 
 	/* Skip the IPv6 header (40 bytes) to get to the inner ICMPv6 packet */
 	innerICMP := data[minIPv6HeaderLen:]
 	origID, origSeq := extractICMPIDAndSeq(innerICMP)
 	return origID == expectedID && origSeq == expectedSeq
 }
 
 func extractICMPIDAndSeq(icmpData []byte) (id, seq int) {
 	/* ICMP Echo Structure: Type(1) + Code(1) + Checksum(2) + ID(2) + Seq(2) */
 	if len(icmpData) < icmpHeaderLen {
 		return 0, 0
 	}
 	origID := binary.BigEndian.Uint16(icmpData[4:6])
 	origSeq := binary.BigEndian.Uint16(icmpData[6:8])
 	return int(origID), int(origSeq)
 }
 
 func printProbeResult(peerIP, prevAddr string, rtt time.Duration) {
 	if peerIP != prevAddr {
 		fmt.Printf("%s ", resolveHostname(peerIP))
 	}
 	fmt.Printf("%.3f ms  ", float64(rtt.Microseconds())/1000.0)
 }
 
 func resolveHostname(ip string) string {
 	names, _ := net.LookupAddr(ip)
 	if len(names) > 0 {
 		hostname := names[0]
 		if strings.HasSuffix(hostname, ".") {
 			hostname = hostname[:len(hostname)-1]
 		}
 		return fmt.Sprintf("%s (%s)", hostname, ip)
 	}
 	return ip
 }
 
 func cleanIP(addr string) string {
 	if host, _, err := net.SplitHostPort(addr); err == nil {
 		return host
 	}
 	return addr
 }

© 2024 - 2026 djnn, inc • generated with stagit