chore: Move private packages to internal (#1664)

internal/headers/iphlpapi/const.go

@@ -0,0 +1,6 @@
+package iphlpapi
+
+const (
+	TCPTableClass  uint32 = 5
+	TCP6TableClass uint32 = 5
+)

internal/headers/iphlpapi/iphlpapi.go

@@ -0,0 +1,77 @@
+package iphlpapi
+
+import (
+	"fmt"
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+var (
+	modiphlpapi             = windows.NewLazySystemDLL("iphlpapi.dll")
+	procGetExtendedTcpTable = modiphlpapi.NewProc("GetExtendedTcpTable")
+)
+
+func GetTCPConnectionStates(family uint32) (map[MIB_TCP_STATE]uint32, error) {
+	var size uint32
+
+	stateCounts := make(map[MIB_TCP_STATE]uint32)
+	rowSize := uint32(unsafe.Sizeof(MIB_TCPROW_OWNER_PID{}))
+	tableClass := TCPTableClass
+
+	if family == windows.AF_INET6 {
+		rowSize = uint32(unsafe.Sizeof(MIB_TCP6ROW_OWNER_PID{}))
+		tableClass = TCP6TableClass
+	}
+
+	ret := getExtendedTcpTable(0, &size, true, family, tableClass, 0)
+	if ret != 0 && ret != uintptr(windows.ERROR_INSUFFICIENT_BUFFER) {
+		return nil, fmt.Errorf("getExtendedTcpTable (size query) failed with code %d", ret)
+	}
+
+	buf := make([]byte, size)
+
+	ret = getExtendedTcpTable(uintptr(unsafe.Pointer(&buf[0])), &size, true, family, tableClass, 0)
+	if ret != 0 {
+		return nil, fmt.Errorf("getExtendedTcpTable (data query) failed with code %d", ret)
+	}
+
+	numEntries := *(*uint32)(unsafe.Pointer(&buf[0]))
+
+	for i := range numEntries {
+		var state MIB_TCP_STATE
+
+		if family == windows.AF_INET6 {
+			row := (*MIB_TCP6ROW_OWNER_PID)(unsafe.Pointer(&buf[4+i*rowSize]))
+			state = row.dwState
+		} else {
+			row := (*MIB_TCPROW_OWNER_PID)(unsafe.Pointer(&buf[4+i*rowSize]))
+			state = row.dwState
+		}
+
+		stateCounts[state]++
+	}
+
+	return stateCounts, nil
+}
+
+func getExtendedTcpTable(pTCPTable uintptr, pdwSize *uint32, bOrder bool, ulAf uint32, tableClass uint32, reserved uint32) uintptr {
+	ret, _, _ := procGetExtendedTcpTable.Call(
+		pTCPTable,
+		uintptr(unsafe.Pointer(pdwSize)),
+		uintptr(boolToInt(bOrder)),
+		uintptr(ulAf),
+		uintptr(tableClass),
+		uintptr(reserved),
+	)
+
+	return ret
+}
+
+func boolToInt(b bool) int {
+	if b {
+		return 1
+	}
+
+	return 0
+}

internal/headers/iphlpapi/types.go

@@ -0,0 +1,76 @@
+package iphlpapi
+
+import "fmt"
+
+// MIB_TCPROW_OWNER_PID structure for IPv4.
+// https://learn.microsoft.com/en-us/windows/win32/api/tcpmib/ns-tcpmib-mib_tcprow_owner_pid
+type MIB_TCPROW_OWNER_PID struct {
+	dwState      MIB_TCP_STATE
+	dwLocalAddr  uint32
+	dwLocalPort  uint32
+	dwRemoteAddr uint32
+	dwRemotePort uint32
+	dwOwningPid  uint32
+}
+
+// MIB_TCP6ROW_OWNER_PID structure for IPv6.
+// https://learn.microsoft.com/en-us/windows/win32/api/tcpmib/ns-tcpmib-mib_tcp6row_owner_pid
+type MIB_TCP6ROW_OWNER_PID struct {
+	ucLocalAddr     [16]byte
+	dwLocalScopeId  uint32
+	dwLocalPort     uint32
+	ucRemoteAddr    [16]byte
+	dwRemoteScopeId uint32
+	dwRemotePort    uint32
+	dwState         MIB_TCP_STATE
+	dwOwningPid     uint32
+}
+
+type MIB_TCP_STATE uint32
+
+const (
+	_ MIB_TCP_STATE = iota
+	TCPStateClosed
+	TCPStateListening
+	TCPStateSynSent
+	TCPStateSynRcvd
+	TCPStateEstablished
+	TCPStateFinWait1
+	TCPStateFinWait2
+	TCPStateCloseWait
+	TCPStateClosing
+	TCPStateLastAck
+	TCPStateTimeWait
+	TCPStateDeleteTcb
+)
+
+func (state MIB_TCP_STATE) String() string {
+	switch state {
+	case TCPStateClosed:
+		return "CLOSED"
+	case TCPStateListening:
+		return "LISTENING"
+	case TCPStateSynSent:
+		return "SYN_SENT"
+	case TCPStateSynRcvd:
+		return "SYN_RECEIVED"
+	case TCPStateEstablished:
+		return "ESTABLISHED"
+	case TCPStateFinWait1:
+		return "FIN_WAIT1"
+	case TCPStateFinWait2:
+		return "FIN_WAIT2"
+	case TCPStateCloseWait:
+		return "CLOSE_WAIT"
+	case TCPStateClosing:
+		return "CLOSING"
+	case TCPStateLastAck:
+		return "LAST_ACK"
+	case TCPStateTimeWait:
+		return "TIME_WAIT"
+	case TCPStateDeleteTcb:
+		return "DELETE_TCB"
+	default:
+		return fmt.Sprintf("UNKNOWN_%d", state)
+	}
+}