From 1c2cebbbc480ec1fc0ab03904e33dc65ed1f7c48 Mon Sep 17 00:00:00 2001 From: pytimer Date: Mon, 9 Jul 2018 10:42:55 +0800 Subject: [PATCH] [net] Implements windows net package function Connections and ConnectionsPid --- net/net_windows.go | 530 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 527 insertions(+), 3 deletions(-) diff --git a/net/net_windows.go b/net/net_windows.go index 7fff20e2..e57c9832 100644 --- a/net/net_windows.go +++ b/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() +}