[net] Implements windows net package function Connections and ConnectionsPid

net/net_windows.go

@@ -5,8 +5,11 @@ package net
 import (
 	"context"
 	"errors"
+	"fmt"
 	"net"
 	"os"
+	"syscall"
+	"unsafe"
 
 	"github.com/shirou/gopsutil/internal/common"
 	"golang.org/x/sys/windows"
@@ -30,6 +33,46 @@ const (
 	TCPTableOwnerModuleAll
 )
 
+type netConnectionKindType struct {
+	family   uint32
+	sockType uint32
+	filename string
+}
+
+var kindTCP4 = netConnectionKindType{
+	family:   syscall.AF_INET,
+	sockType: syscall.SOCK_STREAM,
+	filename: "tcp",
+}
+var kindTCP6 = netConnectionKindType{
+	family:   syscall.AF_INET6,
+	sockType: syscall.SOCK_STREAM,
+	filename: "tcp6",
+}
+var kindUDP4 = netConnectionKindType{
+	family:   syscall.AF_INET,
+	sockType: syscall.SOCK_DGRAM,
+	filename: "udp",
+}
+var kindUDP6 = netConnectionKindType{
+	family:   syscall.AF_INET6,
+	sockType: syscall.SOCK_DGRAM,
+	filename: "udp6",
+}
+
+var netConnectionKindMap = map[string][]netConnectionKindType{
+	"all":   {kindTCP4, kindTCP6, kindUDP4, kindUDP6},
+	"tcp":   {kindTCP4, kindTCP6},
+	"tcp4":  {kindTCP4},
+	"tcp6":  {kindTCP6},
+	"udp":   {kindUDP4, kindUDP6},
+	"udp4":  {kindUDP4},
+	"udp6":  {kindUDP6},
+	"inet":  {kindTCP4, kindTCP6, kindUDP4, kindUDP6},
+	"inet4": {kindTCP4, kindUDP4},
+	"inet6": {kindTCP6, kindUDP6},
+}
+
 func IOCounters(pernic bool) ([]IOCountersStat, error) {
 	return IOCountersWithContext(context.Background(), pernic)
 }
@@ -78,15 +121,71 @@ func IOCountersByFileWithContext(ctx context.Context, pernic bool, filename stri
 	return IOCounters(pernic)
 }
 
-// Return a list of network connections opened by a process
+// Return a list of network connections
+// Available kind:
+//   reference to netConnectionKindMap
 func Connections(kind string) ([]ConnectionStat, error) {
 	return ConnectionsWithContext(context.Background(), kind)
 }
 
 func ConnectionsWithContext(ctx context.Context, kind string) ([]ConnectionStat, error) {
-	var ret []ConnectionStat
+	return ConnectionsPidWithContext(ctx, kind, 0)
+}
 
-	return ret, common.ErrNotImplementedError
+// ConnectionsPid Return a list of network connections opened by a process
+func ConnectionsPid(kind string, pid int32) ([]ConnectionStat, error) {
+	return ConnectionsPidWithContext(context.Background(), kind, pid)
+}
+
+func ConnectionsPidWithContext(ctx context.Context, kind string, pid int32) ([]ConnectionStat, error) {
+	tmap, ok := netConnectionKindMap[kind]
+	if !ok {
+		return nil, fmt.Errorf("invalid kind, %s", kind)
+	}
+	return getProcInet(tmap, pid)
+}
+
+func getProcInet(kinds []netConnectionKindType, pid int32) ([]ConnectionStat, error) {
+	stats := make([]ConnectionStat, 0)
+
+	for _, kind := range kinds {
+		s, err := getNetStatWithKind(kind)
+		if err != nil {
+			continue
+		}
+
+		if pid == 0 {
+			stats = append(stats, s...)
+		} else {
+			for _, ns := range s {
+				if ns.Pid != pid {
+					continue
+				}
+				stats = append(stats, ns)
+			}
+		}
+	}
+
+	return stats, nil
+}
+
+func getNetStatWithKind(kindType netConnectionKindType) ([]ConnectionStat, error) {
+	if kindType.filename == "" {
+		return nil, fmt.Errorf("kind filename must be required")
+	}
+
+	switch kindType.filename {
+	case kindTCP4.filename:
+		return getTCPConnections(kindTCP4.family)
+	case kindTCP6.filename:
+		return getTCPConnections(kindTCP6.family)
+	case kindUDP4.filename:
+		return getUDPConnections(kindUDP4.family)
+	case kindUDP6.filename:
+		return getUDPConnections(kindUDP6.family)
+	}
+
+	return nil, fmt.Errorf("invalid kind filename, %s", kindType.filename)
 }
 
 // Return a list of network connections opened returning at most `max`
@@ -118,3 +217,428 @@ func ProtoCounters(protocols []string) ([]ProtoCountersStat, error) {
 func ProtoCountersWithContext(ctx context.Context, protocols []string) ([]ProtoCountersStat, error) {
 	return nil, errors.New("NetProtoCounters not implemented for windows")
 }
+
+func getTableUintptr(family uint32, buf []byte) uintptr {
+	var (
+		pmibTCPTable  PMIB_TCPTABLE_OWNER_PID_ALL
+		pmibTCP6Table PMIB_TCP6TABLE_OWNER_PID_ALL
+
+		p uintptr
+	)
+	switch family {
+	case kindTCP4.family:
+		if len(buf) > 0 {
+			pmibTCPTable = (*MIB_TCPTABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
+			p = uintptr(unsafe.Pointer(pmibTCPTable))
+		} else {
+			p = uintptr(unsafe.Pointer(pmibTCPTable))
+		}
+	case kindTCP6.family:
+		if len(buf) > 0 {
+			pmibTCP6Table = (*MIB_TCP6TABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
+			p = uintptr(unsafe.Pointer(pmibTCP6Table))
+		} else {
+			p = uintptr(unsafe.Pointer(pmibTCP6Table))
+		}
+	}
+	return p
+}
+
+func getTableInfo(filename string, table interface{}) (index, step, length int) {
+	switch filename {
+	case kindTCP4.filename:
+		index = int(unsafe.Sizeof(table.(PMIB_TCPTABLE_OWNER_PID_ALL).DwNumEntries))
+		step = int(unsafe.Sizeof(table.(PMIB_TCPTABLE_OWNER_PID_ALL).Table))
+		length = int(table.(PMIB_TCPTABLE_OWNER_PID_ALL).DwNumEntries)
+	case kindTCP6.filename:
+		index = int(unsafe.Sizeof(table.(PMIB_TCP6TABLE_OWNER_PID_ALL).DwNumEntries))
+		step = int(unsafe.Sizeof(table.(PMIB_TCP6TABLE_OWNER_PID_ALL).Table))
+		length = int(table.(PMIB_TCP6TABLE_OWNER_PID_ALL).DwNumEntries)
+	case kindUDP4.filename:
+		index = int(unsafe.Sizeof(table.(PMIB_UDPTABLE_OWNER_PID).DwNumEntries))
+		step = int(unsafe.Sizeof(table.(PMIB_UDPTABLE_OWNER_PID).Table))
+		length = int(table.(PMIB_UDPTABLE_OWNER_PID).DwNumEntries)
+	case kindUDP6.filename:
+		index = int(unsafe.Sizeof(table.(PMIB_UDP6TABLE_OWNER_PID).DwNumEntries))
+		step = int(unsafe.Sizeof(table.(PMIB_UDP6TABLE_OWNER_PID).Table))
+		length = int(table.(PMIB_UDP6TABLE_OWNER_PID).DwNumEntries)
+	}
+
+	return
+}
+
+func getTCPConnections(family uint32) ([]ConnectionStat, error) {
+	var (
+		p    uintptr
+		buf  []byte
+		size uint32
+
+		pmibTCPTable  PMIB_TCPTABLE_OWNER_PID_ALL
+		pmibTCP6Table PMIB_TCP6TABLE_OWNER_PID_ALL
+	)
+
+	if family == 0 {
+		return nil, fmt.Errorf("faimly must be required")
+	}
+
+	for {
+		switch family {
+		case kindTCP4.family:
+			if len(buf) > 0 {
+				pmibTCPTable = (*MIB_TCPTABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
+				p = uintptr(unsafe.Pointer(pmibTCPTable))
+			} else {
+				p = uintptr(unsafe.Pointer(pmibTCPTable))
+			}
+		case kindTCP6.family:
+			if len(buf) > 0 {
+				pmibTCP6Table = (*MIB_TCP6TABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
+				p = uintptr(unsafe.Pointer(pmibTCP6Table))
+			} else {
+				p = uintptr(unsafe.Pointer(pmibTCP6Table))
+			}
+		}
+
+		err := getExtendedTcpTable(p,
+			&size,
+			true,
+			family,
+			TCP_TABLE_OWNER_PID_ALL,
+			0)
+		if err == nil {
+			break
+		}
+		if err != windows.ERROR_INSUFFICIENT_BUFFER {
+			return nil, err
+		}
+		buf = make([]byte, size)
+	}
+
+	var (
+		index, step int
+		length      int
+	)
+
+	stats := make([]ConnectionStat, 0)
+	switch family {
+	case kindTCP4.family:
+		index, step, length = getTableInfo(kindTCP4.filename, pmibTCPTable)
+	case kindTCP6.family:
+		index, step, length = getTableInfo(kindTCP6.filename, pmibTCP6Table)
+	}
+
+	if length == 0 {
+		return nil, nil
+	}
+
+	for i := 0; i < length; i++ {
+		switch family {
+		case kindTCP4.family:
+			mibs := (*MIB_TCPROW_OWNER_PID)(unsafe.Pointer(&buf[index]))
+			ns := mibs.convertToConnectionStat()
+			stats = append(stats, ns)
+		case kindTCP6.family:
+			mibs := (*MIB_TCP6ROW_OWNER_PID)(unsafe.Pointer(&buf[index]))
+			ns := mibs.convertToConnectionStat()
+			stats = append(stats, ns)
+		}
+
+		index += step
+	}
+	return stats, nil
+}
+
+func getUDPConnections(family uint32) ([]ConnectionStat, error) {
+	var (
+		p    uintptr
+		buf  []byte
+		size uint32
+
+		pmibUDPTable  PMIB_UDPTABLE_OWNER_PID
+		pmibUDP6Table PMIB_UDP6TABLE_OWNER_PID
+	)
+
+	if family == 0 {
+		return nil, fmt.Errorf("faimly must be required")
+	}
+
+	for {
+		switch family {
+		case kindUDP4.family:
+			if len(buf) > 0 {
+				pmibUDPTable = (*MIB_UDPTABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
+				p = uintptr(unsafe.Pointer(pmibUDPTable))
+			} else {
+				p = uintptr(unsafe.Pointer(pmibUDPTable))
+			}
+		case kindUDP6.family:
+			if len(buf) > 0 {
+				pmibUDP6Table = (*MIB_UDP6TABLE_OWNER_PID)(unsafe.Pointer(&buf[0]))
+				p = uintptr(unsafe.Pointer(pmibUDP6Table))
+			} else {
+				p = uintptr(unsafe.Pointer(pmibUDP6Table))
+			}
+		}
+
+		err := getExtendedUdpTable(
+			p,
+			&size,
+			true,
+			family,
+			UDP_TABLE_OWNER_PID,
+			0,
+		)
+		if err == nil {
+			break
+		}
+		if err != windows.ERROR_INSUFFICIENT_BUFFER {
+			return nil, err
+		}
+		buf = make([]byte, size)
+	}
+
+	var (
+		index, step, length int
+	)
+
+	stats := make([]ConnectionStat, 0)
+	switch family {
+	case kindUDP4.family:
+		index, step, length = getTableInfo(kindUDP4.filename, pmibUDPTable)
+	case kindUDP6.family:
+		index, step, length = getTableInfo(kindUDP6.filename, pmibUDP6Table)
+	}
+
+	if length == 0 {
+		return nil, nil
+	}
+
+	for i := 0; i < length; i++ {
+		switch family {
+		case kindUDP4.family:
+			mibs := (*MIB_UDPROW_OWNER_PID)(unsafe.Pointer(&buf[index]))
+			ns := mibs.convertToConnectionStat()
+			stats = append(stats, ns)
+		case kindUDP4.family:
+			mibs := (*MIB_UDP6ROW_OWNER_PID)(unsafe.Pointer(&buf[index]))
+			ns := mibs.convertToConnectionStat()
+			stats = append(stats, ns)
+		}
+
+		index += step
+	}
+	return stats, nil
+}
+
+// tcpStatuses https://msdn.microsoft.com/en-us/library/windows/desktop/bb485761(v=vs.85).aspx
+var tcpStatuses = map[MIB_TCP_STATE]string{
+	1:  "CLOSED",
+	2:  "LISTEN",
+	3:  "SYN_SENT",
+	4:  "SYN_RECEIVED",
+	5:  "ESTABLISHED",
+	6:  "FIN_WAIT_1",
+	7:  "FIN_WAIT_2",
+	8:  "CLOSE_WAIT",
+	9:  "CLOSING",
+	10: "LAST_ACK",
+	11: "TIME_WAIT",
+	12: "DELETE",
+}
+
+func getExtendedTcpTable(pTcpTable uintptr, pdwSize *uint32, bOrder bool, ulAf uint32, tableClass TCP_TABLE_CLASS, reserved uint32) (errcode error) {
+	r1, _, _ := syscall.Syscall6(procGetExtendedTCPTable.Addr(), 6, pTcpTable, uintptr(unsafe.Pointer(pdwSize)), getUintptrFromBool(bOrder), uintptr(ulAf), uintptr(tableClass), uintptr(reserved))
+	if r1 != 0 {
+		errcode = syscall.Errno(r1)
+	}
+	return
+}
+
+func getExtendedUdpTable(pUdpTable uintptr, pdwSize *uint32, bOrder bool, ulAf uint32, tableClass UDP_TABLE_CLASS, reserved uint32) (errcode error) {
+	r1, _, _ := syscall.Syscall6(procGetExtendedUDPTable.Addr(), 6, pUdpTable, uintptr(unsafe.Pointer(pdwSize)), getUintptrFromBool(bOrder), uintptr(ulAf), uintptr(tableClass), uintptr(reserved))
+	if r1 != 0 {
+		errcode = syscall.Errno(r1)
+	}
+	return
+}
+
+func getUintptrFromBool(b bool) uintptr {
+	if b {
+		return 1
+	}
+	return 0
+}
+
+const ANY_SIZE = 1
+
+type MIB_TCP_STATE int32
+
+type TCP_TABLE_CLASS int32
+
+const (
+	TCP_TABLE_BASIC_LISTENER TCP_TABLE_CLASS = iota
+	TCP_TABLE_BASIC_CONNECTIONS
+	TCP_TABLE_BASIC_ALL
+	TCP_TABLE_OWNER_PID_LISTENER
+	TCP_TABLE_OWNER_PID_CONNECTIONS
+	TCP_TABLE_OWNER_PID_ALL
+	TCP_TABLE_OWNER_MODULE_LISTENER
+	TCP_TABLE_OWNER_MODULE_CONNECTIONS
+	TCP_TABLE_OWNER_MODULE_ALL
+)
+
+type UDP_TABLE_CLASS int32
+
+const (
+	UDP_TABLE_BASIC UDP_TABLE_CLASS = iota
+	UDP_TABLE_OWNER_PID
+	UDP_TABLE_OWNER_MODULE
+)
+
+// TCP
+
+type MIB_TCPROW_OWNER_PID struct {
+	DwState      uint32
+	DwLocalAddr  uint32
+	DwLocalPort  uint32
+	DwRemoteAddr uint32
+	DwRemotePort uint32
+	DwOwningPid  uint32
+}
+
+func (m *MIB_TCPROW_OWNER_PID) convertToConnectionStat() ConnectionStat {
+	ns := ConnectionStat{
+		Family: kindTCP4.family,
+		Type:   kindTCP4.sockType,
+		Laddr: Addr{
+			IP:   parseIPv4HexString(m.DwLocalAddr),
+			Port: uint32(decodePort(m.DwLocalPort)),
+		},
+		Raddr: Addr{
+			IP:   parseIPv4HexString(m.DwRemoteAddr),
+			Port: uint32(decodePort(m.DwRemotePort)),
+		},
+		Pid:    int32(m.DwOwningPid),
+		Status: tcpStatuses[MIB_TCP_STATE(m.DwState)],
+	}
+
+	return ns
+}
+
+type MIB_TCPTABLE_OWNER_PID struct {
+	DwNumEntries uint32
+	Table        [ANY_SIZE]MIB_TCPROW_OWNER_PID
+}
+
+type MIB_TCP6ROW_OWNER_PID struct {
+	UcLocalAddr     [16]byte
+	DwLocalScopeId  uint32
+	DwLocalPort     uint32
+	UcRemoteAddr    [16]byte
+	DwRemoteScopeId uint32
+	DwRemotePort    uint32
+	DwState         uint32
+	DwOwningPid     uint32
+}
+
+func (m *MIB_TCP6ROW_OWNER_PID) convertToConnectionStat() ConnectionStat {
+	ns := ConnectionStat{
+		Family: kindTCP6.family,
+		Type:   kindTCP6.sockType,
+		Laddr: Addr{
+			IP:   parseIPv6HexString(m.UcLocalAddr),
+			Port: uint32(decodePort(m.DwLocalPort)),
+		},
+		Raddr: Addr{
+			IP:   parseIPv6HexString(m.UcRemoteAddr),
+			Port: uint32(decodePort(m.DwRemotePort)),
+		},
+		Pid:    int32(m.DwOwningPid),
+		Status: tcpStatuses[MIB_TCP_STATE(m.DwState)],
+	}
+
+	return ns
+}
+
+type MIB_TCP6TABLE_OWNER_PID struct {
+	DwNumEntries uint32
+	Table        [ANY_SIZE]MIB_TCP6ROW_OWNER_PID
+}
+
+type PMIB_TCPTABLE_OWNER_PID_ALL *MIB_TCPTABLE_OWNER_PID
+type PMIB_TCP6TABLE_OWNER_PID_ALL *MIB_TCP6TABLE_OWNER_PID
+
+// UDP
+
+type MIB_UDPROW_OWNER_PID struct {
+	DwLocalAddr uint32
+	DwLocalPort uint32
+	DwOwningPid uint32
+}
+
+func (m *MIB_UDPROW_OWNER_PID) convertToConnectionStat() ConnectionStat {
+	ns := ConnectionStat{
+		Family: kindUDP4.family,
+		Type:   kindUDP4.sockType,
+		Laddr: Addr{
+			IP:   parseIPv4HexString(m.DwLocalAddr),
+			Port: uint32(decodePort(m.DwLocalPort)),
+		},
+		Pid: int32(m.DwOwningPid),
+	}
+
+	return ns
+}
+
+type MIB_UDPTABLE_OWNER_PID struct {
+	DwNumEntries uint32
+	Table        [ANY_SIZE]MIB_UDPROW_OWNER_PID
+}
+
+type MIB_UDP6ROW_OWNER_PID struct {
+	UcLocalAddr    [16]byte
+	DwLocalScopeId uint32
+	DwLocalPort    uint32
+	DwOwningPid    uint32
+}
+
+func (m *MIB_UDP6ROW_OWNER_PID) convertToConnectionStat() ConnectionStat {
+	ns := ConnectionStat{
+		Family: kindUDP6.family,
+		Type:   kindUDP6.sockType,
+		Laddr: Addr{
+			IP:   parseIPv6HexString(m.UcLocalAddr),
+			Port: uint32(decodePort(m.DwLocalPort)),
+		},
+		Pid: int32(m.DwOwningPid),
+	}
+
+	return ns
+}
+
+type MIB_UDP6TABLE_OWNER_PID struct {
+	DwNumEntries uint32
+	Table        [ANY_SIZE]MIB_UDP6ROW_OWNER_PID
+}
+
+type PMIB_UDPTABLE_OWNER_PID *MIB_UDPTABLE_OWNER_PID
+type PMIB_UDP6TABLE_OWNER_PID *MIB_UDP6TABLE_OWNER_PID
+
+func decodePort(port uint32) uint16 {
+	return syscall.Ntohs(uint16(port))
+}
+
+func parseIPv4HexString(addr uint32) string {
+	return fmt.Sprintf("%d.%d.%d.%d", addr&255, addr>>8&255, addr>>16&255, addr>>24&255)
+}
+
+func parseIPv6HexString(addr [16]byte) string {
+	var ret [16]byte
+	for i := 0; i < 16; i++ {
+		ret[i] = uint8(addr[i])
+	}
+
+	// convert []byte to net.IP
+	ip := net.IP(ret[:])
+	return ip.String()
+}