win32api/iphlpapi.go

package win32api

import (
	"fmt"
	"strings"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

func GetAdaptersAddresses(family, flags uint32) ([]AdapterAddressInfo, error) {
	size := uint32(15 * 1024)
	for i := 0; i < 4; i++ {
		buf := make([]byte, size)
		head := (*windows.IpAdapterAddresses)(unsafe.Pointer(&buf[0]))
		err := windows.GetAdaptersAddresses(family, flags, 0, head, &size)
		if err == nil {
			return collectAdapterAddressInfo(head), nil
		}
		if errno, ok := err.(syscall.Errno); ok && errno == syscall.ERROR_BUFFER_OVERFLOW {
			continue
		}
		return nil, err
	}
	return nil, fmt.Errorf("GetAdaptersAddresses exceeded retry limit")
}

func GetIfTable2() ([]MIB_IF_ROW2, error) {
	proc, err := getProcAddr("iphlpapi.dll", "GetIfTable2")
	if err != nil {
		return nil, err
	}

	var table *MIB_IF_TABLE2
	r, _, _ := syscall.Syscall(proc, 1, uintptr(unsafe.Pointer(&table)), 0, 0)
	if r != 0 {
		return nil, syscall.Errno(r)
	}
	if table == nil {
		return nil, fmt.Errorf("GetIfTable2 returned nil table")
	}
	defer FreeMibTable(unsafe.Pointer(table))

	count := int(table.NumEntries)
	rows := make([]MIB_IF_ROW2, 0, count)
	rowSize := unsafe.Sizeof(table.Table[0])
	base := uintptr(unsafe.Pointer(&table.Table[0]))
	for i := 0; i < count; i++ {
		row := (*MIB_IF_ROW2)(unsafe.Pointer(base + uintptr(i)*rowSize))
		rows = append(rows, *row)
	}
	return rows, nil
}

func GetIfEntry2(row *MIB_IF_ROW2) error {
	if row == nil {
		return syscall.EINVAL
	}
	proc, err := getProcAddr("iphlpapi.dll", "GetIfEntry2")
	if err != nil {
		return err
	}
	r, _, _ := syscall.Syscall(proc, 1, uintptr(unsafe.Pointer(row)), 0, 0)
	if r != 0 {
		return syscall.Errno(r)
	}
	return nil
}

func FreeMibTable(memory unsafe.Pointer) {
	if memory == nil {
		return
	}
	proc, err := getProcAddr("iphlpapi.dll", "FreeMibTable")
	if err != nil {
		return
	}
	syscall.Syscall(proc, 1, uintptr(memory), 0, 0)
}

func GetExtendedTcp4Table(order bool, tableClass TCP_TABLE_CLASS) ([]MIB_TCPROW_OWNER_PID, error) {
	if err := validateTCPOwnerPIDTableClass(tableClass); err != nil {
		return nil, err
	}
	buf, err := getExtendedTCPTable(order, AF_INET, tableClass)
	if err != nil {
		return nil, err
	}
	if len(buf) == 0 {
		return nil, nil
	}
	table := (*MIB_TCPTABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
	count := int(table.NumEntries)
	rows := make([]MIB_TCPROW_OWNER_PID, 0, count)
	rowSize := unsafe.Sizeof(table.Table[0])
	base := uintptr(unsafe.Pointer(&table.Table[0]))
	for i := 0; i < count; i++ {
		row := (*MIB_TCPROW_OWNER_PID)(unsafe.Pointer(base + uintptr(i)*rowSize))
		rows = append(rows, *row)
	}
	return rows, nil
}

func GetExtendedTcp6Table(order bool, tableClass TCP_TABLE_CLASS) ([]MIB_TCP6ROW_OWNER_PID, error) {
	if err := validateTCPOwnerPIDTableClass(tableClass); err != nil {
		return nil, err
	}
	buf, err := getExtendedTCPTable(order, AF_INET6, tableClass)
	if err != nil {
		return nil, err
	}
	if len(buf) == 0 {
		return nil, nil
	}
	table := (*MIB_TCP6TABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
	count := int(table.NumEntries)
	rows := make([]MIB_TCP6ROW_OWNER_PID, 0, count)
	rowSize := unsafe.Sizeof(table.Table[0])
	base := uintptr(unsafe.Pointer(&table.Table[0]))
	for i := 0; i < count; i++ {
		row := (*MIB_TCP6ROW_OWNER_PID)(unsafe.Pointer(base + uintptr(i)*rowSize))
		rows = append(rows, *row)
	}
	return rows, nil
}

func GetExtendedUdp4Table(order bool, tableClass UDP_TABLE_CLASS) ([]MIB_UDPROW_OWNER_PID, error) {
	if err := validateUDPOwnerPIDTableClass(tableClass); err != nil {
		return nil, err
	}
	buf, err := getExtendedUDPTable(order, AF_INET, tableClass)
	if err != nil {
		return nil, err
	}
	if len(buf) == 0 {
		return nil, nil
	}
	table := (*MIB_UDPTABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
	count := int(table.NumEntries)
	rows := make([]MIB_UDPROW_OWNER_PID, 0, count)
	rowSize := unsafe.Sizeof(table.Table[0])
	base := uintptr(unsafe.Pointer(&table.Table[0]))
	for i := 0; i < count; i++ {
		row := (*MIB_UDPROW_OWNER_PID)(unsafe.Pointer(base + uintptr(i)*rowSize))
		rows = append(rows, *row)
	}
	return rows, nil
}

func GetExtendedUdp6Table(order bool, tableClass UDP_TABLE_CLASS) ([]MIB_UDP6ROW_OWNER_PID, error) {
	if err := validateUDPOwnerPIDTableClass(tableClass); err != nil {
		return nil, err
	}
	buf, err := getExtendedUDPTable(order, AF_INET6, tableClass)
	if err != nil {
		return nil, err
	}
	if len(buf) == 0 {
		return nil, nil
	}
	table := (*MIB_UDP6TABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
	count := int(table.NumEntries)
	rows := make([]MIB_UDP6ROW_OWNER_PID, 0, count)
	rowSize := unsafe.Sizeof(table.Table[0])
	base := uintptr(unsafe.Pointer(&table.Table[0]))
	for i := 0; i < count; i++ {
		row := (*MIB_UDP6ROW_OWNER_PID)(unsafe.Pointer(base + uintptr(i)*rowSize))
		rows = append(rows, *row)
	}
	return rows, nil
}

func (row MIB_TCPROW_OWNER_PID) LocalPortHost() uint16 {
	return Ntohs(uint16(row.LocalPort))
}

func (row MIB_TCPROW_OWNER_PID) RemotePortHost() uint16 {
	return Ntohs(uint16(row.RemotePort))
}

func (row MIB_TCP6ROW_OWNER_PID) LocalPortHost() uint16 {
	return Ntohs(uint16(row.LocalPort))
}

func (row MIB_TCP6ROW_OWNER_PID) RemotePortHost() uint16 {
	return Ntohs(uint16(row.RemotePort))
}

func (row MIB_UDPROW_OWNER_PID) LocalPortHost() uint16 {
	return Ntohs(uint16(row.LocalPort))
}

func (row MIB_UDP6ROW_OWNER_PID) LocalPortHost() uint16 {
	return Ntohs(uint16(row.LocalPort))
}

func validateTCPOwnerPIDTableClass(tableClass TCP_TABLE_CLASS) error {
	switch tableClass {
	case TCP_TABLE_OWNER_PID_LISTENER, TCP_TABLE_OWNER_PID_CONNECTIONS, TCP_TABLE_OWNER_PID_ALL:
		return nil
	default:
		return fmt.Errorf("unsupported TCP owner-pid table class: %d", tableClass)
	}
}

func validateUDPOwnerPIDTableClass(tableClass UDP_TABLE_CLASS) error {
	if tableClass == UDP_TABLE_OWNER_PID {
		return nil
	}
	return fmt.Errorf("unsupported UDP owner-pid table class: %d", tableClass)
}

func getExtendedTCPTable(order bool, family int, tableClass TCP_TABLE_CLASS) ([]byte, error) {
	proc, err := getProcAddr("iphlpapi.dll", "GetExtendedTcpTable")
	if err != nil {
		return nil, err
	}
	return getExtendedIPTable(proc, order, family, uint32(tableClass))
}

func getExtendedUDPTable(order bool, family int, tableClass UDP_TABLE_CLASS) ([]byte, error) {
	proc, err := getProcAddr("iphlpapi.dll", "GetExtendedUdpTable")
	if err != nil {
		return nil, err
	}
	return getExtendedIPTable(proc, order, family, uint32(tableClass))
}

func getExtendedIPTable(proc uintptr, order bool, family int, tableClass uint32) ([]byte, error) {
	size := uint32(0)
	r, _, _ := syscall.Syscall6(proc, 6,
		0,
		uintptr(unsafe.Pointer(&size)),
		boolToUintptr(order),
		uintptr(family),
		uintptr(tableClass),
		0,
	)
	if r != 0 && syscall.Errno(r) != syscall.ERROR_INSUFFICIENT_BUFFER {
		return nil, syscall.Errno(r)
	}
	if size == 0 {
		return nil, nil
	}

	for i := 0; i < 4; i++ {
		buf := make([]byte, size)
		r, _, _ = syscall.Syscall6(proc, 6,
			uintptr(unsafe.Pointer(&buf[0])),
			uintptr(unsafe.Pointer(&size)),
			boolToUintptr(order),
			uintptr(family),
			uintptr(tableClass),
			0,
		)
		if r == 0 {
			return buf, nil
		}
		if syscall.Errno(r) != syscall.ERROR_INSUFFICIENT_BUFFER {
			return nil, syscall.Errno(r)
		}
	}
	return nil, fmt.Errorf("GetExtended IP table exceeded retry limit")
}

func boolToUintptr(v bool) uintptr {
	if v {
		return 1
	}
	return 0
}

func collectAdapterAddressInfo(head *windows.IpAdapterAddresses) []AdapterAddressInfo {
	out := make([]AdapterAddressInfo, 0, 8)
	for a := head; a != nil; a = a.Next {
		item := AdapterAddressInfo{
			IfIndex:               a.IfIndex,
			AdapterName:           windows.BytePtrToString(a.AdapterName),
			FriendlyName:          windows.UTF16PtrToString(a.FriendlyName),
			Description:           windows.UTF16PtrToString(a.Description),
			DNSSuffix:             windows.UTF16PtrToString(a.DnsSuffix),
			OperStatus:            a.OperStatus,
			Mtu:                   a.Mtu,
			MACAddress:            formatMACAddress(a.PhysicalAddress[:], a.PhysicalAddressLength),
			PhysicalAddressLength: a.PhysicalAddressLength,
			TransmitLinkSpeed:     a.TransmitLinkSpeed,
			ReceiveLinkSpeed:      a.ReceiveLinkSpeed,
			UnicastIPs:            collectUnicastIPs(a.FirstUnicastAddress),
			DNSServers:            collectDNSServerIPs(a.FirstDnsServerAddress),
			Gateways:              collectGatewayIPs(a.FirstGatewayAddress),
		}
		out = append(out, item)
	}
	return out
}

func collectUnicastIPs(addr *windows.IpAdapterUnicastAddress) []string {
	out := make([]string, 0, 4)
	for ua := addr; ua != nil; ua = ua.Next {
		ip := ua.Address.IP()
		if ip == nil {
			continue
		}
		out = append(out, ip.String())
	}
	return out
}

func collectDNSServerIPs(addr *windows.IpAdapterDnsServerAdapter) []string {
	out := make([]string, 0, 2)
	for cur := addr; cur != nil; cur = cur.Next {
		ip := cur.Address.IP()
		if ip == nil {
			continue
		}
		out = append(out, ip.String())
	}
	return out
}

func collectGatewayIPs(addr *windows.IpAdapterGatewayAddress) []string {
	out := make([]string, 0, 2)
	for cur := addr; cur != nil; cur = cur.Next {
		ip := cur.Address.IP()
		if ip == nil {
			continue
		}
		out = append(out, ip.String())
	}
	return out
}

func formatMACAddress(raw []byte, n uint32) string {
	if n == 0 || len(raw) == 0 {
		return ""
	}
	if int(n) > len(raw) {
		n = uint32(len(raw))
	}

	var b strings.Builder
	for i := 0; i < int(n); i++ {
		if i > 0 {
			b.WriteByte(':')
		}
		_, _ = fmt.Fprintf(&b, "%02X", raw[i])
	}
	return b.String()
}