tcp: use GetExtendedTcpTable to display states Closed, Listening, SynSent, SynRcvd, CloseWait, TimeWait ... (#1638)

Co-authored-by: Jan-Otto Kröpke <github@jkroepke.de>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Jan-Otto Kröpke <joe@cloudeteer.de>

pkg/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
+}