From 3fb9243fc2175e7b163e9b0e8aa8e233f2cf9bed Mon Sep 17 00:00:00 2001
From: Lomanic <Lomanic@users.noreply.github.com>
Date: Fri, 31 May 2019 17:39:23 +0200
Subject: [PATCH] [net][windows] Fix #693 use MIB_IF_ROW2/GetIfEntry2 to get
 real uint64 values from win32 API

---
 net/net_windows.go | 127 +++++++++++++++++++++++++++++++++++++--------
 1 file changed, 106 insertions(+), 21 deletions(-)

diff --git a/net/net_windows.go b/net/net_windows.go
index 61eb6ec..111a73c 100644
--- a/net/net_windows.go
+++ b/net/net_windows.go
@@ -19,6 +19,7 @@ var (
 	modiphlpapi             = windows.NewLazySystemDLL("iphlpapi.dll")
 	procGetExtendedTCPTable = modiphlpapi.NewProc("GetExtendedTcpTable")
 	procGetExtendedUDPTable = modiphlpapi.NewProc("GetExtendedUdpTable")
+	procGetIfEntry2         = modiphlpapi.NewProc("GetIfEntry2")
 )
 
 const (
@@ -73,6 +74,65 @@ var netConnectionKindMap = map[string][]netConnectionKindType{
 	"inet6": {kindTCP6, kindUDP6},
 }
 
+// https://github.com/microsoft/ethr/blob/aecdaf923970e5a9b4c461b4e2e3963d781ad2cc/plt_windows.go#L114-L170
+type guid struct {
+	Data1 uint32
+	Data2 uint16
+	Data3 uint16
+	Data4 [8]byte
+}
+
+const (
+	maxStringSize        = 256
+	maxPhysAddressLength = 32
+	pad0for64_4for32     = 0
+)
+
+type mibIfRow2 struct {
+	InterfaceLuid               uint64
+	InterfaceIndex              uint32
+	InterfaceGuid               guid
+	Alias                       [maxStringSize + 1]uint16
+	Description                 [maxStringSize + 1]uint16
+	PhysicalAddressLength       uint32
+	PhysicalAddress             [maxPhysAddressLength]uint8
+	PermanentPhysicalAddress    [maxPhysAddressLength]uint8
+	Mtu                         uint32
+	Type                        uint32
+	TunnelType                  uint32
+	MediaType                   uint32
+	PhysicalMediumType          uint32
+	AccessType                  uint32
+	DirectionType               uint32
+	InterfaceAndOperStatusFlags uint32
+	OperStatus                  uint32
+	AdminStatus                 uint32
+	MediaConnectState           uint32
+	NetworkGuid                 guid
+	ConnectionType              uint32
+	padding1                    [pad0for64_4for32]byte
+	TransmitLinkSpeed           uint64
+	ReceiveLinkSpeed            uint64
+	InOctets                    uint64
+	InUcastPkts                 uint64
+	InNUcastPkts                uint64
+	InDiscards                  uint64
+	InErrors                    uint64
+	InUnknownProtos             uint64
+	InUcastOctets               uint64
+	InMulticastOctets           uint64
+	InBroadcastOctets           uint64
+	OutOctets                   uint64
+	OutUcastPkts                uint64
+	OutNUcastPkts               uint64
+	OutDiscards                 uint64
+	OutErrors                   uint64
+	OutUcastOctets              uint64
+	OutMulticastOctets          uint64
+	OutBroadcastOctets          uint64
+	OutQLen                     uint64
+}
+
 func IOCounters(pernic bool) ([]IOCountersStat, error) {
 	return IOCountersWithContext(context.Background(), pernic)
 }
@@ -82,34 +142,59 @@ func IOCountersWithContext(ctx context.Context, pernic bool) ([]IOCountersStat,
 	if err != nil {
 		return nil, err
 	}
-	var ret []IOCountersStat
+	var counters []IOCountersStat
 
-	for _, ifi := range ifs {
-		c := IOCountersStat{
-			Name: ifi.Name,
+	err = procGetIfEntry2.Find()
+	if err == nil { // Vista+, uint64 values (issue#693)
+		for _, ifi := range ifs {
+			c := IOCountersStat{
+				Name: ifi.Name,
+			}
+
+			row := mibIfRow2{InterfaceIndex: uint32(ifi.Index)}
+			ret, _, err := procGetIfEntry2.Call(uintptr(unsafe.Pointer(&row)))
+			if ret != 0 {
+				return nil, os.NewSyscallError("GetIfEntry2", err)
+			}
+			c.BytesSent = uint64(row.OutOctets)
+			c.BytesRecv = uint64(row.InOctets)
+			c.PacketsSent = uint64(row.OutUcastPkts)
+			c.PacketsRecv = uint64(row.InUcastPkts)
+			c.Errin = uint64(row.InErrors)
+			c.Errout = uint64(row.OutErrors)
+			c.Dropin = uint64(row.InDiscards)
+			c.Dropout = uint64(row.OutDiscards)
+
+			counters = append(counters, c)
 		}
+	} else { // WinXP fallback, uint32 values
+		for _, ifi := range ifs {
+			c := IOCountersStat{
+				Name: ifi.Name,
+			}
 
-		row := windows.MibIfRow{Index: uint32(ifi.Index)}
-		e := windows.GetIfEntry(&row)
-		if e != nil {
-			return nil, os.NewSyscallError("GetIfEntry", e)
+			row := windows.MibIfRow{Index: uint32(ifi.Index)}
+			err = windows.GetIfEntry(&row)
+			if err != nil {
+				return nil, os.NewSyscallError("GetIfEntry", err)
+			}
+			c.BytesSent = uint64(row.OutOctets)
+			c.BytesRecv = uint64(row.InOctets)
+			c.PacketsSent = uint64(row.OutUcastPkts)
+			c.PacketsRecv = uint64(row.InUcastPkts)
+			c.Errin = uint64(row.InErrors)
+			c.Errout = uint64(row.OutErrors)
+			c.Dropin = uint64(row.InDiscards)
+			c.Dropout = uint64(row.OutDiscards)
+
+			counters = append(counters, c)
 		}
-		c.BytesSent = uint64(row.OutOctets)
-		c.BytesRecv = uint64(row.InOctets)
-		c.PacketsSent = uint64(row.OutUcastPkts)
-		c.PacketsRecv = uint64(row.InUcastPkts)
-		c.Errin = uint64(row.InErrors)
-		c.Errout = uint64(row.OutErrors)
-		c.Dropin = uint64(row.InDiscards)
-		c.Dropout = uint64(row.OutDiscards)
-
-		ret = append(ret, c)
 	}
 
-	if pernic == false {
-		return getIOCountersAll(ret)
+	if !pernic {
+		return getIOCountersAll(counters)
 	}
-	return ret, nil
+	return counters, nil
 }
 
 // NetIOCountersByFile is an method which is added just a compatibility for linux.