Improve performance, logging

2019-05-10 11:57:51 +02:00
parent e72ed023db
commit 16db10b84b
9 changed files with 282 additions and 129 deletions
--- a/process/iphelper/get.go
+++ b/process/iphelper/get.go
@@ -10,15 +10,15 @@ import (
 )

 var (
-	tcp4Connections []*connectionEntry
-	tcp4Listeners   []*connectionEntry
-	tcp6Connections []*connectionEntry
-	tcp6Listeners   []*connectionEntry
+	tcp4Connections []*ConnectionEntry
+	tcp4Listeners   []*ConnectionEntry
+	tcp6Connections []*ConnectionEntry
+	tcp6Listeners   []*ConnectionEntry

-	udp4Connections []*connectionEntry
-	udp4Listeners   []*connectionEntry
-	udp6Connections []*connectionEntry
-	udp6Listeners   []*connectionEntry
+	udp4Connections []*ConnectionEntry
+	udp4Listeners   []*ConnectionEntry
+	udp6Connections []*ConnectionEntry
+	udp6Listeners   []*ConnectionEntry

 	ipHelper *IPHelper
 	lock     sync.RWMutex
@@ -34,6 +34,7 @@ func checkIPHelper() (err error) {
 	return nil
 }

+// GetTCP4PacketInfo returns the pid of the given IPv4/TCP connection.
 func GetTCP4PacketInfo(localIP net.IP, localPort uint16, remoteIP net.IP, remotePort uint16, pktDirection bool) (pid int, direction bool, err error) {

 	// search
@@ -69,6 +70,7 @@ func GetTCP4PacketInfo(localIP net.IP, localPort uint16, remoteIP net.IP, remote
 	return -1, pktDirection, nil
 }

+// GetTCP6PacketInfo returns the pid of the given IPv6/TCP connection.
 func GetTCP6PacketInfo(localIP net.IP, localPort uint16, remoteIP net.IP, remotePort uint16, pktDirection bool) (pid int, direction bool, err error) {

 	// search
@@ -104,6 +106,7 @@ func GetTCP6PacketInfo(localIP net.IP, localPort uint16, remoteIP net.IP, remote
 	return -1, pktDirection, nil
 }

+// GetUDP4PacketInfo returns the pid of the given IPv4/UDP connection.
 func GetUDP4PacketInfo(localIP net.IP, localPort uint16, remoteIP net.IP, remotePort uint16, pktDirection bool) (pid int, direction bool, err error) {

 	// search
@@ -139,6 +142,7 @@ func GetUDP4PacketInfo(localIP net.IP, localPort uint16, remoteIP net.IP, remote
 	return -1, pktDirection, nil
 }

+// GetUDP6PacketInfo returns the pid of the given IPv6/UDP connection.
 func GetUDP6PacketInfo(localIP net.IP, localPort uint16, remoteIP net.IP, remotePort uint16, pktDirection bool) (pid int, direction bool, err error) {

 	// search
@@ -174,7 +178,7 @@ func GetUDP6PacketInfo(localIP net.IP, localPort uint16, remoteIP net.IP, remote
 	return -1, pktDirection, nil
 }

-func search(connections, listeners []*connectionEntry, localIP, remoteIP net.IP, localPort, remotePort uint16, pktDirection bool) (pid int, direction bool) {
+func search(connections, listeners []*ConnectionEntry, localIP, remoteIP net.IP, localPort, remotePort uint16, pktDirection bool) (pid int, direction bool) {
 	lock.RLock()
 	defer lock.RUnlock()

@@ -203,7 +207,7 @@ func search(connections, listeners []*connectionEntry, localIP, remoteIP net.IP,
 	return -1, pktDirection
 }

-func searchConnections(list []*connectionEntry, localIP, remoteIP net.IP, localPort, remotePort uint16) (pid int) {
+func searchConnections(list []*ConnectionEntry, localIP, remoteIP net.IP, localPort, remotePort uint16) (pid int) {

 	for _, entry := range list {
 		if localPort == entry.localPort &&
@@ -217,7 +221,7 @@ func searchConnections(list []*connectionEntry, localIP, remoteIP net.IP, localP
 	return -1
 }

-func searchListeners(list []*connectionEntry, localIP net.IP, localPort uint16) (pid int) {
+func searchListeners(list []*ConnectionEntry, localIP net.IP, localPort uint16) (pid int) {

 	for _, entry := range list {
 		if localPort == entry.localPort &&
@@ -230,6 +234,7 @@ func searchListeners(list []*connectionEntry, localIP net.IP, localPort uint16)
 	return -1
 }

+// GetActiveConnectionIDs returns all currently active connection IDs.
 func GetActiveConnectionIDs() (connections []string) {
 	lock.Lock()
 	defer lock.Unlock()
--- a/process/iphelper/iphelper.go
+++ b/process/iphelper/iphelper.go
@@ -14,11 +14,12 @@ var (
 	errInvalid = errors.New("IPHelper not initialzed or broken")
 )

+// IPHelper represents a subset of the Windows iphlpapi.dll.
 type IPHelper struct {
 	dll *windows.LazyDLL

-	getExtendedTcpTable *windows.LazyProc
-	getExtendedUdpTable *windows.LazyProc
+	getExtendedTCPTable *windows.LazyProc
+	getExtendedUDPTable *windows.LazyProc
 	// getOwnerModuleFromTcpEntry  *windows.LazyProc
 	// getOwnerModuleFromTcp6Entry *windows.LazyProc
 	// getOwnerModuleFromUdpEntry  *windows.LazyProc
@@ -27,6 +28,7 @@ type IPHelper struct {
 	valid *abool.AtomicBool
 }

+// New returns a new IPHelper API (with an instance of iphlpapi.dll loaded).
 func New() (*IPHelper, error) {

 	new := &IPHelper{}
@@ -41,13 +43,13 @@ func New() (*IPHelper, error) {
 	}

 	// load functions
-	new.getExtendedTcpTable = new.dll.NewProc("GetExtendedTcpTable")
-	err = new.getExtendedTcpTable.Find()
+	new.getExtendedTCPTable = new.dll.NewProc("GetExtendedTcpTable")
+	err = new.getExtendedTCPTable.Find()
 	if err != nil {
 		return nil, fmt.Errorf("could find proc GetExtendedTcpTable: %s", err)
 	}
-	new.getExtendedUdpTable = new.dll.NewProc("GetExtendedUdpTable")
-	err = new.getExtendedUdpTable.Find()
+	new.getExtendedUDPTable = new.dll.NewProc("GetExtendedUdpTable")
+	err = new.getExtendedUDPTable.Find()
 	if err != nil {
 		return nil, fmt.Errorf("could find proc GetExtendedUdpTable: %s", err)
 	}
--- a/process/iphelper/tables.go
+++ b/process/iphelper/tables.go
@@ -6,18 +6,24 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"sync"
 	"unsafe"

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

+// Windows API constants
 const (
-	iphelper_TCP_TABLE_OWNER_PID_ALL uintptr = 5
-	iphelper_UDP_TABLE_OWNER_PID     uintptr = 1
-	iphelper_TCP_STATE_LISTEN        uint32  = 2
+	iphelperTCPTableOwnerPIDAll uintptr = 5
+	iphelperUDPTableOwnerPID    uintptr = 1
+	iphelperTCPStateListen      uint32  = 2
+
+	winErrInsufficientBuffer = uintptr(windows.ERROR_INSUFFICIENT_BUFFER)
+	winErrInvalidParameter   = uintptr(windows.ERROR_INVALID_PARAMETER)
 )

-type connectionEntry struct {
+// ConnectionEntry describes a connection state table entry.
+type ConnectionEntry struct {
 	localIP    net.IP
 	remoteIP   net.IP
 	localPort  uint16
@@ -25,17 +31,17 @@ type connectionEntry struct {
 	pid        int
 }

-func (entry *connectionEntry) String() string {
+func (entry *ConnectionEntry) String() string {
 	return fmt.Sprintf("PID=%d %s:%d <> %s:%d", entry.pid, entry.localIP, entry.localPort, entry.remoteIP, entry.remotePort)
 }

-type iphelperTcpTable struct {
+type iphelperTCPTable struct {
 	// docs: https://msdn.microsoft.com/en-us/library/windows/desktop/aa366921(v=vs.85).aspx
 	numEntries uint32
-	table      [4096]iphelperTcpRow
+	table      [4096]iphelperTCPRow
 }

-type iphelperTcpRow struct {
+type iphelperTCPRow struct {
 	// docs: https://msdn.microsoft.com/en-us/library/windows/desktop/aa366913(v=vs.85).aspx
 	state      uint32
 	localAddr  uint32
@@ -45,51 +51,52 @@ type iphelperTcpRow struct {
 	owningPid  uint32
 }

-type iphelperTcp6Table struct {
+type iphelperTCP6Table struct {
 	// docs: https://msdn.microsoft.com/en-us/library/windows/desktop/aa366905(v=vs.85).aspx
 	numEntries uint32
-	table      [4096]iphelperTcp6Row
+	table      [4096]iphelperTCP6Row
 }

-type iphelperTcp6Row struct {
+type iphelperTCP6Row struct {
 	// docs: https://msdn.microsoft.com/en-us/library/windows/desktop/aa366896(v=vs.85).aspx
 	localAddr     [16]byte
-	localScopeId  uint32
+	localScopeID  uint32
 	localPort     uint32
 	remoteAddr    [16]byte
-	remoteScopeId uint32
+	remoteScopeID uint32
 	remotePort    uint32
 	state         uint32
 	owningPid     uint32
 }

-type iphelperUdpTable struct {
+type iphelperUDPTable struct {
 	// docs: https://msdn.microsoft.com/en-us/library/windows/desktop/aa366932(v=vs.85).aspx
 	numEntries uint32
-	table      [4096]iphelperUdpRow
+	table      [4096]iphelperUDPRow
 }

-type iphelperUdpRow struct {
+type iphelperUDPRow struct {
 	// docs: https://msdn.microsoft.com/en-us/library/windows/desktop/aa366928(v=vs.85).aspx
 	localAddr uint32
 	localPort uint32
 	owningPid uint32
 }

-type iphelperUdp6Table struct {
+type iphelperUDP6Table struct {
 	// docs: https://msdn.microsoft.com/en-us/library/windows/desktop/aa366925(v=vs.85).aspx
 	numEntries uint32
-	table      [4096]iphelperUdp6Row
+	table      [4096]iphelperUDP6Row
 }

-type iphelperUdp6Row struct {
+type iphelperUDP6Row struct {
 	// docs: https://msdn.microsoft.com/en-us/library/windows/desktop/aa366923(v=vs.85).aspx
 	localAddr    [16]byte
-	localScopeId uint32
+	localScopeID uint32
 	localPort    uint32
 	owningPid    uint32
 }

+// IP and Protocol constants
 const (
 	IPv4 uint8 = 4
 	IPv6 uint8 = 6
@@ -98,7 +105,51 @@ const (
 	UDP uint8 = 17
 )

-func (ipHelper *IPHelper) GetTables(protocol uint8, ipVersion uint8) (connections []*connectionEntry, listeners []*connectionEntry, err error) {
+const (
+	startBufSize = 4096
+	bufSizeUses  = 100
+)
+
+var (
+	bufSize          = startBufSize
+	bufSizeUsageLeft = bufSizeUses
+	bufSizeLock      sync.Mutex
+)
+
+func getBufSize() int {
+	bufSizeLock.Lock()
+	defer bufSizeLock.Unlock()
+
+	// using bufSize
+	bufSizeUsageLeft--
+	// check if we want to reset
+	if bufSizeUsageLeft <= 0 {
+		// reset
+		bufSize = startBufSize
+		bufSizeUsageLeft = bufSizeUses
+	}
+
+	return bufSize
+}
+
+func increaseBufSize() int {
+	bufSizeLock.Lock()
+	defer bufSizeLock.Unlock()
+
+	// increase
+	bufSize = bufSize * 2
+	// not too much
+	if bufSize > 65536 {
+		bufSize = 65536
+	}
+	// reset
+	bufSizeUsageLeft = bufSizeUses
+	// return new bufSize
+	return bufSize
+}
+
+// GetTables returns the current connection state table of Windows of the given protocol and IP version.
+func (ipHelper *IPHelper) GetTables(protocol uint8, ipVersion uint8) (connections []*ConnectionEntry, listeners []*ConnectionEntry, err error) {
 	// docs: https://msdn.microsoft.com/en-us/library/windows/desktop/aa365928(v=vs.85).aspx

 	if !ipHelper.valid.IsSet() {
@@ -115,50 +166,61 @@ func (ipHelper *IPHelper) GetTables(protocol uint8, ipVersion uint8) (connection
 		return nil, nil, errors.New("invalid protocol")
 	}

-	bufSize := 4096
-	buf := make([]byte, bufSize)
-	var r1 uintptr
+	// try max 3 times
+	maxTries := 3
+	bufSize := getBufSize()
+	var buf []byte

-	switch protocol {
-	case TCP:
-		r1, _, err = ipHelper.getExtendedTcpTable.Call(
-			uintptr(unsafe.Pointer(&buf[0])),  // _Out_   PVOID           pTcpTable
-			uintptr(unsafe.Pointer(&bufSize)), // _Inout_ PDWORD          pdwSize
-			0,                                 // _In_    BOOL            bOrder
-			uintptr(afClass),                  // _In_    ULONG           ulAf
-			iphelper_TCP_TABLE_OWNER_PID_ALL,  // _In_    TCP_TABLE_CLASS TableClass
-			0,                                 // _In_    ULONG           Reserved
-		)
-	case UDP:
-		r1, _, err = ipHelper.getExtendedUdpTable.Call(
-			uintptr(unsafe.Pointer(&buf[0])),  // _Out_   PVOID           pUdpTable,
-			uintptr(unsafe.Pointer(&bufSize)), // _Inout_ PDWORD          pdwSize,
-			0,                                 // _In_    BOOL            bOrder,
-			uintptr(afClass),                  // _In_    ULONG           ulAf,
-			iphelper_UDP_TABLE_OWNER_PID,      // _In_    UDP_TABLE_CLASS TableClass,
-			0,                                 // _In_    ULONG           Reserved
-		)
-	}
+	for i := 1; i <= maxTries; i++ {
+		buf = make([]byte, bufSize)
+		var r1 uintptr

-	switch r1 {
-	// case windows.ERROR_INSUFFICIENT_BUFFER:
-	// 	return nil, fmt.Errorf("insufficient buffer error: %s", err)
-	// case windows.ERROR_INVALID_PARAMETER:
-	// 	return nil, fmt.Errorf("invalid parameter: %s", err)
-	case windows.NO_ERROR:
-	default:
-		return nil, nil, fmt.Errorf("unexpected error: %s", err)
+		switch protocol {
+		case TCP:
+			r1, _, err = ipHelper.getExtendedTCPTable.Call(
+				uintptr(unsafe.Pointer(&buf[0])),  // _Out_   PVOID           pTcpTable
+				uintptr(unsafe.Pointer(&bufSize)), // _Inout_ PDWORD          pdwSize
+				0,                                 // _In_    BOOL            bOrder
+				uintptr(afClass),                  // _In_    ULONG           ulAf
+				iphelperTCPTableOwnerPIDAll,       // _In_    TCP_TABLE_CLASS TableClass
+				0,                                 // _In_    ULONG           Reserved
+			)
+		case UDP:
+			r1, _, err = ipHelper.getExtendedUDPTable.Call(
+				uintptr(unsafe.Pointer(&buf[0])),  // _Out_   PVOID           pUdpTable,
+				uintptr(unsafe.Pointer(&bufSize)), // _Inout_ PDWORD          pdwSize,
+				0,                                 // _In_    BOOL            bOrder,
+				uintptr(afClass),                  // _In_    ULONG           ulAf,
+				iphelperUDPTableOwnerPID,          // _In_    UDP_TABLE_CLASS TableClass,
+				0,                                 // _In_    ULONG           Reserved
+			)
+		}
+
+		switch r1 {
+		case winErrInsufficientBuffer:
+			if i >= maxTries {
+				return nil, nil, fmt.Errorf("insufficient buffer error (tried %d times): [NT 0x%X] %s", i, r1, err)
+			}
+			bufSize = increaseBufSize()
+		case winErrInvalidParameter:
+			return nil, nil, fmt.Errorf("invalid parameter: [NT 0x%X] %s", r1, err)
+		case windows.NO_ERROR:
+			// success
+			break
+		default:
+			return nil, nil, fmt.Errorf("unexpected error: [NT 0x%X] %s", r1, err)
+		}
 	}

 	// parse output
 	switch {
 	case protocol == TCP && ipVersion == IPv4:

-		tcpTable := (*iphelperTcpTable)(unsafe.Pointer(&buf[0]))
+		tcpTable := (*iphelperTCPTable)(unsafe.Pointer(&buf[0]))
 		table := tcpTable.table[:tcpTable.numEntries]

 		for _, row := range table {
-			new := &connectionEntry{}
+			new := &ConnectionEntry{}

 			// PID
 			new.pid = int(row.owningPid)
@@ -170,7 +232,7 @@ func (ipHelper *IPHelper) GetTables(protocol uint8, ipVersion uint8) (connection
 			new.localPort = uint16(row.localPort>>8 | row.localPort<<8)

 			// remote
-			if row.state == iphelper_TCP_STATE_LISTEN {
+			if row.state == iphelperTCPStateListen {
 				listeners = append(listeners, new)
 			} else {
 				new.remoteIP = convertIPv4(row.remoteAddr)
@@ -182,11 +244,11 @@ func (ipHelper *IPHelper) GetTables(protocol uint8, ipVersion uint8) (connection

 	case protocol == TCP && ipVersion == IPv6:

-		tcpTable := (*iphelperTcp6Table)(unsafe.Pointer(&buf[0]))
+		tcpTable := (*iphelperTCP6Table)(unsafe.Pointer(&buf[0]))
 		table := tcpTable.table[:tcpTable.numEntries]

 		for _, row := range table {
-			new := &connectionEntry{}
+			new := &ConnectionEntry{}

 			// PID
 			new.pid = int(row.owningPid)
@@ -196,7 +258,7 @@ func (ipHelper *IPHelper) GetTables(protocol uint8, ipVersion uint8) (connection
 			new.localPort = uint16(row.localPort>>8 | row.localPort<<8)

 			// remote
-			if row.state == iphelper_TCP_STATE_LISTEN {
+			if row.state == iphelperTCPStateListen {
 				if new.localIP.Equal(net.IPv6zero) {
 					new.localIP = nil
 				}
@@ -211,11 +273,11 @@ func (ipHelper *IPHelper) GetTables(protocol uint8, ipVersion uint8) (connection

 	case protocol == UDP && ipVersion == IPv4:

-		udpTable := (*iphelperUdpTable)(unsafe.Pointer(&buf[0]))
+		udpTable := (*iphelperUDPTable)(unsafe.Pointer(&buf[0]))
 		table := udpTable.table[:udpTable.numEntries]

 		for _, row := range table {
-			new := &connectionEntry{}
+			new := &ConnectionEntry{}

 			// PID
 			new.pid = int(row.owningPid)
@@ -232,11 +294,11 @@ func (ipHelper *IPHelper) GetTables(protocol uint8, ipVersion uint8) (connection

 	case protocol == UDP && ipVersion == IPv6:

-		udpTable := (*iphelperUdp6Table)(unsafe.Pointer(&buf[0]))
+		udpTable := (*iphelperUDP6Table)(unsafe.Pointer(&buf[0]))
 		table := udpTable.table[:udpTable.numEntries]

 		for _, row := range table {
-			new := &connectionEntry{}
+			new := &ConnectionEntry{}

 			// PID
 			new.pid = int(row.owningPid)