livetrace/src/winapi_windows.go

//go:build windows

package main

import (
	"fmt"
	"golang.org/x/net/ipv4"
	"net"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"syscall"
	"time"
	"unsafe"
)

var (
	modIphlpapi    = syscall.NewLazyDLL("iphlpapi.dll")
	procIcmpCreate = modIphlpapi.NewProc("IcmpCreateFile")
	procIcmpClose  = modIphlpapi.NewProc("IcmpCloseHandle")
	procIcmpSend2  = modIphlpapi.NewProc("IcmpSendEcho2")
	rePingFrom     = regexp.MustCompile(`Reply from ([\d.]+):`)
	rePingTime     = regexp.MustCompile(`time[=<](\d+)ms`)
	rePingTimeout  = regexp.MustCompile(`Request timed out|could not find`)
)

const (
	ipSuccess           uint32 = 0
	ipTTLExpiredTransit uint32 = 11013
	ipReqTimedOut       uint32 = 11010
)

func udpModeWarning() {
	fmt.Fprintln(os.Stderr,
		"Warning: UDP mode on Windows requires an inbound Windows Firewall rule "+
			"to allow ICMP Time Exceeded (Type 11). Consider using -proto winapi instead.")
}

// mirrors IP_OPTION_INFORMATION from <ipexport.h>
type ipOptInfo struct {
	Ttl         uint8
	Tos         uint8
	Flags       uint8
	OptionsSize uint8
	OptionsData uintptr // PUCHAR — platform-sized pointer
}

// mirrors ICMP_ECHO_REPLY from <ipexport.h>
type icmpEchoReply struct {
	Address       uint32 // source IP, bytes in network order inside a LE DWORD
	Status        uint32
	RoundTripTime uint32 // integer ms; we measure our own sub-ms time
	DataSize      uint16
	Reserved      uint16
	Data          uintptr // PVOID — platform-sized pointer
	Options       ipOptInfo
}

// winAPIProbe sends one ICMP echo with the given TTL via the Windows ICMP
// helper API (iphlpapi.dll) and blocks until a reply arrives or timeout
// elapses. Returns (rtt ms, responder IP, IP status code, error).
func winAPIProbe(targetIP net.IP, ttl int, timeout time.Duration, payload []byte) (float64, string, uint32, error) {
	handle, _, e := procIcmpCreate.Call()
	if handle == 0 {
		return 0, "", 0, fmt.Errorf("IcmpCreateFile: %w", e)
	}
	defer procIcmpClose.Call(handle) //nolint:errcheck

	// Pack the destination IP into a DWORD with the bytes in network order.
	// On LE x86/x64: storing [ip[0]..ip[3]] means ip[0] is the LSB of the uint32.
	ip4 := targetIP.To4()
	dst := uint32(ip4[0]) | uint32(ip4[1])<<8 | uint32(ip4[2])<<16 | uint32(ip4[3])<<24

	opts := ipOptInfo{Ttl: uint8(ttl)}

	replyBufSz := uintptr(unsafe.Sizeof(icmpEchoReply{})) + uintptr(len(payload)) + 8
	replyBuf := make([]byte, replyBufSz)

	var dataPtr uintptr
	if len(payload) > 0 {
		dataPtr = uintptr(unsafe.Pointer(&payload[0]))
	}

	start := time.Now()
	ret, _, _ := procIcmpSend2.Call(
		handle,
		0, 0, 0, // Event, ApcRoutine, ApcContext — synchronous mode
		uintptr(dst),
		dataPtr, uintptr(len(payload)),
		uintptr(unsafe.Pointer(&opts)),
		uintptr(unsafe.Pointer(&replyBuf[0])),
		uintptr(replyBufSz),
		uintptr(timeout.Milliseconds()),
	)
	rtt := float64(time.Since(start).Microseconds()) / 1000.0

	if ret == 0 {
		// Timeout or hard error — treat as timeout.
		return 0, "", ipReqTimedOut, nil
	}

	reply := (*icmpEchoReply)(unsafe.Pointer(&replyBuf[0]))

	// Unpack the source IP: Address bytes are in network order inside the DWORD.
	fromIP := net.IP{
		byte(reply.Address),
		byte(reply.Address >> 8),
		byte(reply.Address >> 16),
		byte(reply.Address >> 24),
	}.String()

	return rtt, fromIP, reply.Status, nil
}

// sendProbeWinAPI sends one probe for the given TTL using the Windows ICMP
// API. It returns immediately; the blocking IcmpSendEcho2 call runs in its
// own goroutine and records the result directly into the hop.
func (t *Tracer) sendProbeWinAPI(ttl int) {
	t.hops[ttl-1].recordSent()

	go func() {
		payload := make([]byte, t.payloadSize())
		copy(payload, "livetrace")

		rtt, fromIP, status, err := winAPIProbe(t.targetIP, ttl, t.timeout, payload)

		// Discard the result if the tracer was stopped while we were blocked.
		select {
		case <-t.stopCh:
			return
		default:
		}

		if err != nil || status == ipReqTimedOut {
			t.hops[ttl-1].recordTimeout()
			return
		}

		reached := status == ipSuccess
		t.hops[ttl-1].recordRecv(rtt, fromIP, reached)

		if reached {
			t.mu.Lock()
			if t.maxActive > ttl {
				t.maxActive = ttl
			}
			t.mu.Unlock()
		}
	}()
}

func setUDPSocketTTL(pc *ipv4.PacketConn, conn net.PacketConn, ttl int) error {
	udpConn, ok := conn.(*net.UDPConn)
	if !ok {
		return pc.SetTTL(ttl)
	}
	raw, err := udpConn.SyscallConn()
	if err != nil {
		return pc.SetTTL(ttl)
	}
	var sockErr error
	raw.Control(func(fd uintptr) {
		sockErr = syscall.SetsockoptInt(syscall.Handle(fd),
			syscall.IPPROTO_IP, syscall.IP_TTL, ttl)
	})
	return sockErr
}

func pingProbe(targetIP net.IP, ttl int, timeout time.Duration) (rtt float64, fromIP string, reached bool, timedOut bool) {
	cmd := exec.Command("ping",
		"-n", "1",
		"-i", strconv.Itoa(ttl),
		"-w", strconv.Itoa(int(timeout.Milliseconds())),
		targetIP.String(),
	)
	start := time.Now()
	out, _ := cmd.Output()
	elapsed := float64(time.Since(start).Microseconds()) / 1000.0
	text := string(out)

	if rePingTimeout.MatchString(text) {
		return 0, "", false, true
	}
	m := rePingFrom.FindStringSubmatch(text)
	if m == nil {
		return 0, "", false, true
	}
	fromIP = m[1]
	reached = fromIP == targetIP.String()

	if mt := rePingTime.FindStringSubmatch(text); mt != nil {
		if v, err := strconv.ParseFloat(mt[1], 64); err == nil {
			rtt = v
		}
	} else {
		rtt = elapsed
	}
	return rtt, fromIP, reached, false
}

func (t *Tracer) sendProbeWinPing(ttl int) {
	t.hops[ttl-1].recordSent()
	go func() {
		rtt, fromIP, reached, timedOut := pingProbe(t.targetIP, ttl, t.timeout)
		select {
		case <-t.stopCh:
			return
		default:
		}
		if timedOut {
			t.hops[ttl-1].recordTimeout()
			return
		}
		t.hops[ttl-1].recordRecv(rtt, fromIP, reached)
		if reached {
			t.mu.Lock()
			if t.maxActive > ttl {
				t.maxActive = ttl
			}
			t.mu.Unlock()
		}
	}()
}