Add example for getting bandwidth stats

This commit is contained in:
Vladimir Stoilov 2023-06-30 16:56:39 +03:00
parent 95716dbb47
commit 83479acc95
8 changed files with 154 additions and 41 deletions

View file

@ -104,6 +104,7 @@ func apiAuthenticator(r *http.Request, s *http.Server) (token *api.AuthToken, er
SrcPort: remotePort, // source as in the process we are looking for
Dst: localIP,
DstPort: localPort,
PID: process.UndefinedProcessID,
},
)
if !retry {

View file

@ -28,6 +28,9 @@ func start(ch chan packet.Packet) error {
go windowskext.Handler(ch)
// Example worker for the bandwidth data stats. Not ment for production.
// windowskext.StartBandwidthWorker()
return nil
}

View file

@ -0,0 +1,76 @@
package windowskext
// This file contains example code how to read bandwidth stats from the kext. Its not ment to be used in production.
import (
"time"
"github.com/safing/portbase/log"
)
type Rxtxdata struct {
rx uint64
tx uint64
}
type Key struct {
localIP [4]uint32
remoteIP [4]uint32
localPort uint16
remotePort uint16
ipv6 bool
protocol uint8
}
var m = make(map[Key]Rxtxdata)
func StartBandwidthWorker() {
go func() {
ticker := time.NewTicker(2 * time.Second)
defer ticker.Stop()
for range ticker.C {
conns, err := GetConnectionsStats()
if err != nil {
continue
}
for _, conn := range conns {
if conn.receivedBytes == 0 && conn.transmittedBytes == 0 {
continue
}
key := Key{
localIP: conn.localIP,
remoteIP: conn.remoteIP,
localPort: conn.localPort,
remotePort: conn.remotePort,
ipv6: conn.ipV6 == 1,
protocol: conn.protocol,
}
// First we get a "copy" of the entry
if entry, ok := m[key]; ok {
// Then we modify the copy
entry.rx += conn.receivedBytes
entry.tx += conn.transmittedBytes
// Then we reassign map entry
m[key] = entry
} else {
m[key] = Rxtxdata{
rx: conn.receivedBytes,
tx: conn.transmittedBytes,
}
}
}
log.Debug("----------------------------------")
for key, value := range m {
if key.ipv6 {
log.Debugf("Conn: %d %s:%d %s:%d rx:%d tx:%d", key.protocol, convertIPv6(key.localIP), key.localPort, convertIPv6(key.remoteIP), key.remotePort, value.rx, value.tx)
} else {
log.Debugf("Conn: %d %s:%d %s:%d rx:%d tx:%d", key.protocol, convertIPv4(key.localIP), key.localPort, convertIPv4(key.remoteIP), key.remotePort, value.rx, value.tx)
}
}
}
}()
}

View file

@ -68,6 +68,17 @@ type VerdictUpdateInfo struct {
verdict uint8 // New verdict
}
type ConnectionStat struct {
localIP [4]uint32 //Source Address, only srcIP[0] if IPv4
remoteIP [4]uint32 //Destination Address
localPort uint16 //Source Port
remotePort uint16 //Destination port
receivedBytes uint64 //Number of bytes recived on this connection
transmittedBytes uint64 //Number of bytes transsmited from this connection
ipV6 uint8 //True: IPv6, False: IPv4
protocol uint8 //Protocol (UDP, TCP, ...)
}
type VersionInfo struct {
major uint8
minor uint8

View file

@ -288,6 +288,26 @@ func GetVersion() (*VersionInfo, error) {
return version, nil
}
func GetConnectionsStats() ([]ConnectionStat, error) {
kextLock.RLock()
defer kextLock.RUnlock()
// Check if driver is initialized
if kextHandle == winInvalidHandleValue {
log.Error("kext: failed to clear the cache: kext not ready")
return nil, ErrKextNotReady
}
var data [100]ConnectionStat
size := len(data)
_, err := deviceIOControl(kextHandle, IOCTL_GET_CONNECTIONS_STAT, asByteArray(&size), asByteArray(&data))
if err != nil {
return nil, err
}
return data[:], nil
}
func openDriver(filename string) (windows.Handle, error) {
u16filename, err := syscall.UTF16FromString(filename)
if err != nil {

View file

@ -13,46 +13,47 @@ const (
SIOCTL_TYPE = 40000
)
var (
IOCTL_VERSION = ctlCode(SIOCTL_TYPE, 0x800, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_SHUTDOWN_REQUEST = ctlCode(SIOCTL_TYPE, 0x801, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_RECV_VERDICT_REQ = ctlCode(SIOCTL_TYPE, 0x802, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_SET_VERDICT = ctlCode(SIOCTL_TYPE, 0x803, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_GET_PAYLOAD = ctlCode(SIOCTL_TYPE, 0x804, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_CLEAR_CACHE = ctlCode(SIOCTL_TYPE, 0x805, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_UPDATE_VERDICT = ctlCode(SIOCTL_TYPE, 0x806, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
)
func ctlCode(device_type, function, method, access uint32) uint32 {
return (device_type << 16) | (access << 14) | (function << 2) | method
}
func deviceIOControlAsync(handle windows.Handle, code uint32, inData []byte, outData []byte) (*windows.Overlapped, error) {
var inDataPtr *byte = nil
var inDataSize uint32 = 0
if inData != nil {
inDataPtr = &inData[0]
inDataSize = uint32(len(inData))
}
var outDataPtr *byte = nil
var outDataSize uint32 = 0
if outData != nil {
outDataPtr = &outData[0]
outDataSize = uint32(len(outData))
}
overlapped := &windows.Overlapped{}
err := windows.DeviceIoControl(handle,
code,
inDataPtr, inDataSize,
outDataPtr, outDataSize,
nil, overlapped)
if err != nil {
return nil, err
}
var (
IOCTL_VERSION = ctlCode(SIOCTL_TYPE, 0x800, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_SHUTDOWN_REQUEST = ctlCode(SIOCTL_TYPE, 0x801, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_RECV_VERDICT_REQ = ctlCode(SIOCTL_TYPE, 0x802, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_SET_VERDICT = ctlCode(SIOCTL_TYPE, 0x803, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_GET_PAYLOAD = ctlCode(SIOCTL_TYPE, 0x804, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_CLEAR_CACHE = ctlCode(SIOCTL_TYPE, 0x805, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_UPDATE_VERDICT = ctlCode(SIOCTL_TYPE, 0x806, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
IOCTL_GET_CONNECTIONS_STAT = ctlCode(SIOCTL_TYPE, 0x807, METHOD_BUFFERED, windows.FILE_READ_DATA|windows.FILE_WRITE_DATA)
)
func ctlCode(device_type, function, method, access uint32) uint32 {
return (device_type << 16) | (access << 14) | (function << 2) | method
}
func deviceIOControlAsync(handle windows.Handle, code uint32, inData []byte, outData []byte) (*windows.Overlapped, error) {
var inDataPtr *byte = nil
var inDataSize uint32 = 0
if inData != nil {
inDataPtr = &inData[0]
inDataSize = uint32(len(inData))
}
var outDataPtr *byte = nil
var outDataSize uint32 = 0
if outData != nil {
outDataPtr = &outData[0]
outDataSize = uint32(len(outData))
}
overlapped := &windows.Overlapped{}
err := windows.DeviceIoControl(handle,
code,
inDataPtr, inDataSize,
outDataPtr, outDataSize,
nil, overlapped)
if err != nil {
return nil, err
}
return overlapped, nil

View file

@ -66,6 +66,7 @@ func cleanConnections() (activePIDs map[int]struct{}) {
SrcPort: conn.LocalPort,
Dst: conn.Entity.IP,
DstPort: conn.Entity.Port,
PID: process.UndefinedProcessID,
}, now)
activePIDs[conn.process.Pid] = struct{}{}

View file

@ -126,7 +126,7 @@ func GetProcessByRequestOrigin(ar *api.Request) (*Process, error) {
Protocol: packet.TCP,
Src: remoteIP, // source as in the process we are looking for
SrcPort: remotePort, // source as in the process we are looking for
PID: -1,
PID: UndefinedProcessID,
}
pid, _, err := GetPidOfConnection(ar.Context(), pkt)