mssql: expose server version info (#1741)

Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

internal/headers/iphlpapi/const.go

@@ -3,6 +3,6 @@
 package iphlpapi
 
 const (
-	TCPTableClass  uint32 = 5
-	TCP6TableClass uint32 = 5
+	TCPTableOwnerPIDAll      uint32 = 5
+	TCPTableOwnerPIDListener uint32 = 3
 )

internal/headers/iphlpapi/iphlpapi.go

@@ -3,6 +3,7 @@
 package iphlpapi
 
 import (
+	"encoding/binary"
 	"fmt"
 	"unsafe"
 
@@ -15,65 +16,99 @@ var (
 )
 
 func GetTCPConnectionStates(family uint32) (map[MIB_TCP_STATE]uint32, error) {
+	stateCounts := make(map[MIB_TCP_STATE]uint32)
+
+	switch family {
+	case windows.AF_INET:
+		table, err := getExtendedTcpTable[MIB_TCPROW_OWNER_PID](family, TCPTableOwnerPIDAll)
+		if err != nil {
+			return nil, err
+		}
+
+		for _, row := range table {
+			stateCounts[row.dwState]++
+		}
+
+		return stateCounts, nil
+	case windows.AF_INET6:
+		table, err := getExtendedTcpTable[MIB_TCP6ROW_OWNER_PID](family, TCPTableOwnerPIDAll)
+		if err != nil {
+			return nil, err
+		}
+
+		for _, row := range table {
+			stateCounts[row.dwState]++
+		}
+
+		return stateCounts, nil
+	default:
+		return nil, fmt.Errorf("unsupported address family %d", family)
+	}
+}
+
+func GetOwnerPIDOfTCPPort(family uint32, tcpPort uint16) (uint32, error) {
+	switch family {
+	case windows.AF_INET:
+		table, err := getExtendedTcpTable[MIB_TCPROW_OWNER_PID](family, TCPTableOwnerPIDListener)
+		if err != nil {
+			return 0, err
+		}
+
+		for _, row := range table {
+			if row.dwLocalPort.uint16() == tcpPort {
+				return row.dwOwningPid, nil
+			}
+		}
+
+		return 0, fmt.Errorf("no process found for port %d", tcpPort)
+	case windows.AF_INET6:
+		table, err := getExtendedTcpTable[MIB_TCP6ROW_OWNER_PID](family, TCPTableOwnerPIDListener)
+		if err != nil {
+			return 0, err
+		}
+
+		for _, row := range table {
+			if row.dwLocalPort.uint16() == tcpPort {
+				return row.dwOwningPid, nil
+			}
+		}
+
+		return 0, fmt.Errorf("no process found for port %d", tcpPort)
+	default:
+		return 0, fmt.Errorf("unsupported address family %d", family)
+	}
+}
+
+func getExtendedTcpTable[T any](ulAf uint32, tableClass uint32) ([]T, error) {
 	var size uint32
 
-	stateCounts := make(map[MIB_TCP_STATE]uint32)
-	rowSize := uint32(unsafe.Sizeof(MIB_TCPROW_OWNER_PID{}))
-	tableClass := TCPTableClass
+	ret, _, _ := procGetExtendedTcpTable.Call(
+		uintptr(0),
+		uintptr(unsafe.Pointer(&size)),
+		uintptr(0),
+		uintptr(ulAf),
+		uintptr(tableClass),
+		uintptr(0),
+	)
 
-	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) {
+	if 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)
+	ret, _, _ = procGetExtendedTcpTable.Call(
+		uintptr(unsafe.Pointer(&buf[0])),
+		uintptr(unsafe.Pointer(&size)),
+		uintptr(0),
+		uintptr(ulAf),
+		uintptr(tableClass),
+		uintptr(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
+	return unsafe.Slice((*T)(unsafe.Pointer(&buf[4])), binary.LittleEndian.Uint32(buf)), nil
 }

internal/headers/iphlpapi/iphlpapi_test.go

@@ -0,0 +1,34 @@
+package iphlpapi_test
+
+import (
+	"net"
+	"os"
+	"testing"
+
+	"github.com/prometheus-community/windows_exporter/internal/headers/iphlpapi"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sys/windows"
+)
+
+func TestGetTCPConnectionStates(t *testing.T) {
+	t.Parallel()
+
+	pid, err := iphlpapi.GetTCPConnectionStates(windows.AF_INET)
+	require.NoError(t, err)
+	require.NotEmpty(t, pid)
+}
+
+func TestGetOwnerPIDOfTCPPort(t *testing.T) {
+	t.Parallel()
+
+	lister, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		require.NoError(t, lister.Close())
+	})
+
+	pid, err := iphlpapi.GetOwnerPIDOfTCPPort(windows.AF_INET, uint16(lister.Addr().(*net.TCPAddr).Port))
+	require.NoError(t, err)
+	require.EqualValues(t, os.Getpid(), pid)
+}

internal/headers/iphlpapi/types.go

@@ -2,28 +2,35 @@
 
 package iphlpapi
 
-import "fmt"
+import (
+	"encoding/binary"
+	"fmt"
+)
 
 // MIB_TCPROW_OWNER_PID structure for IPv4.
 // https://learn.microsoft.com/en-us/windows/win32/api/tcpmib/ns-tcpmib-mib_tcprow_owner_pid
+//
+//nolint:unused
 type MIB_TCPROW_OWNER_PID struct {
 	dwState      MIB_TCP_STATE
-	dwLocalAddr  uint32
-	dwLocalPort  uint32
-	dwRemoteAddr uint32
-	dwRemotePort uint32
+	dwLocalAddr  BigEndianUint32
+	dwLocalPort  BigEndianUint32
+	dwRemoteAddr BigEndianUint32
+	dwRemotePort BigEndianUint32
 	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
+//
+//nolint:unused
 type MIB_TCP6ROW_OWNER_PID struct {
 	ucLocalAddr     [16]byte
 	dwLocalScopeId  uint32
-	dwLocalPort     uint32
+	dwLocalPort     BigEndianUint32
 	ucRemoteAddr    [16]byte
 	dwRemoteScopeId uint32
-	dwRemotePort    uint32
+	dwRemotePort    BigEndianUint32
 	dwState         MIB_TCP_STATE
 	dwOwningPid     uint32
 }
@@ -76,3 +83,12 @@ func (state MIB_TCP_STATE) String() string {
 		return fmt.Sprintf("UNKNOWN_%d", state)
 	}
 }
+
+type BigEndianUint32 uint32
+
+func (b BigEndianUint32) uint16() uint16 {
+	data := make([]byte, 2)
+	binary.BigEndian.PutUint16(data, uint16(b))
+
+	return binary.LittleEndian.Uint16(data)
+}