Split htons to return uint16 instead of int for correct byte order conversion, add htonsInt wrapper that converts uint16 to int for syscall.Socket protocol parameter, update openSocket to use htonsInt for ethPAll protocol value
339 lines
8.7 KiB
Go
339 lines
8.7 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"net"
|
|
"os"
|
|
"regexp"
|
|
"sort"
|
|
"strings"
|
|
"syscall"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
ethPAll = 0x0003
|
|
ethPIP = 0x0800
|
|
ethP8021Q = 0x8100
|
|
ethP8021AD = 0x88A8
|
|
packetOut = 4
|
|
protoICMP = 1
|
|
protoTCP = 6
|
|
protoUDP = 17
|
|
maxFrameSize = 65535
|
|
)
|
|
|
|
var vmInterfaceRE = regexp.MustCompile(`(?:tap|fwbr|fwln|fwpr)(\d+)`)
|
|
var vmInterfaceDetailRE = regexp.MustCompile(`(?:tap|fwbr|fwln|fwpr)(\d+)i(\d+)`)
|
|
|
|
type flowKey struct {
|
|
VMID string
|
|
NIC string
|
|
Interface string
|
|
Direction string
|
|
SourceIP string
|
|
DestinationIP string
|
|
Protocol string
|
|
SourcePort int
|
|
DestinationPort int
|
|
}
|
|
|
|
type flowValue struct {
|
|
SourceIP string `json:"source_ip"`
|
|
DestinationIP string `json:"destination_ip"`
|
|
Protocol string `json:"protocol"`
|
|
SourcePort *int `json:"source_port,omitempty"`
|
|
DestinationPort *int `json:"destination_port,omitempty"`
|
|
Packets uint64 `json:"packets"`
|
|
Bytes uint64 `json:"bytes"`
|
|
VMID string `json:"vmid,omitempty"`
|
|
NIC string `json:"nic,omitempty"`
|
|
Interface string `json:"interface,omitempty"`
|
|
Direction string `json:"direction,omitempty"`
|
|
State string `json:"state"`
|
|
Collector string `json:"collector"`
|
|
}
|
|
|
|
type diagnostics struct {
|
|
AttachMode string `json:"attach_mode"`
|
|
InterfacesRequested []string `json:"interfaces_requested"`
|
|
InterfacesAttached []string `json:"interfaces_attached"`
|
|
Errors []string `json:"errors"`
|
|
}
|
|
|
|
type payload struct {
|
|
Flows []flowValue `json:"flows"`
|
|
Diagnostics diagnostics `json:"diagnostics"`
|
|
}
|
|
|
|
func htons(value uint16) uint16 {
|
|
return (value<<8)&0xff00 | value>>8
|
|
}
|
|
|
|
func htonsInt(value uint16) int {
|
|
return int(htons(value))
|
|
}
|
|
|
|
func intPtr(value int) *int {
|
|
if value == 0 {
|
|
return nil
|
|
}
|
|
return &value
|
|
}
|
|
|
|
func protocolName(value byte) string {
|
|
switch value {
|
|
case protoICMP:
|
|
return "icmp"
|
|
case protoTCP:
|
|
return "tcp"
|
|
case protoUDP:
|
|
return "udp"
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func parsePacket(packet []byte) (flowKey, int, bool) {
|
|
var key flowKey
|
|
if len(packet) < 34 {
|
|
return key, 0, false
|
|
}
|
|
offset := 12
|
|
ethType := binary.BigEndian.Uint16(packet[offset : offset+2])
|
|
offset = 14
|
|
for ethType == ethP8021Q || ethType == ethP8021AD {
|
|
if len(packet) < offset+4 {
|
|
return key, 0, false
|
|
}
|
|
ethType = binary.BigEndian.Uint16(packet[offset+2 : offset+4])
|
|
offset += 4
|
|
}
|
|
if ethType != ethPIP || len(packet) < offset+20 {
|
|
return key, 0, false
|
|
}
|
|
versionIHL := packet[offset]
|
|
version := versionIHL >> 4
|
|
ihl := int(versionIHL&0x0f) * 4
|
|
if version != 4 || ihl < 20 || len(packet) < offset+ihl {
|
|
return key, 0, false
|
|
}
|
|
totalLength := int(binary.BigEndian.Uint16(packet[offset+2 : offset+4]))
|
|
protocol := protocolName(packet[offset+9])
|
|
if protocol == "" {
|
|
return key, 0, false
|
|
}
|
|
key.SourceIP = net.IP(packet[offset+12 : offset+16]).String()
|
|
key.DestinationIP = net.IP(packet[offset+16 : offset+20]).String()
|
|
key.Protocol = protocol
|
|
transportOffset := offset + ihl
|
|
if protocol == "tcp" || protocol == "udp" {
|
|
if len(packet) < transportOffset+4 {
|
|
return key, 0, false
|
|
}
|
|
key.SourcePort = int(binary.BigEndian.Uint16(packet[transportOffset : transportOffset+2]))
|
|
key.DestinationPort = int(binary.BigEndian.Uint16(packet[transportOffset+2 : transportOffset+4]))
|
|
} else if protocol == "icmp" && len(packet) >= transportOffset+2 {
|
|
key.SourcePort = int(packet[transportOffset])
|
|
key.DestinationPort = int(packet[transportOffset+1])
|
|
}
|
|
if totalLength <= 0 {
|
|
totalLength = len(packet) - offset
|
|
}
|
|
return key, totalLength, true
|
|
}
|
|
|
|
func interfaceMeta(name string) (string, string) {
|
|
vmid := ""
|
|
nic := "0"
|
|
if match := vmInterfaceRE.FindStringSubmatch(name); len(match) > 1 {
|
|
vmid = match[1]
|
|
}
|
|
if match := vmInterfaceDetailRE.FindStringSubmatch(name); len(match) > 2 {
|
|
nic = match[2]
|
|
}
|
|
return vmid, nic
|
|
}
|
|
|
|
func setFd(fd int, set *syscall.FdSet) {
|
|
set.Bits[fd/64] |= 1 << uint(fd%64)
|
|
}
|
|
|
|
func isSet(fd int, set *syscall.FdSet) bool {
|
|
return set.Bits[fd/64]&(1<<uint(fd%64)) != 0
|
|
}
|
|
|
|
func openSocket(interfaceName string) (int, error) {
|
|
iface, err := net.InterfaceByName(interfaceName)
|
|
if err != nil {
|
|
return -1, err
|
|
}
|
|
fd, err := syscall.Socket(syscall.AF_PACKET, syscall.SOCK_RAW, htonsInt(ethPAll))
|
|
if err != nil {
|
|
return -1, err
|
|
}
|
|
addr := &syscall.SockaddrLinklayer{Protocol: htons(ethPAll), Ifindex: iface.Index}
|
|
if err := syscall.Bind(fd, addr); err != nil {
|
|
_ = syscall.Close(fd)
|
|
return -1, err
|
|
}
|
|
if err := syscall.SetNonblock(fd, true); err != nil {
|
|
_ = syscall.Close(fd)
|
|
return -1, err
|
|
}
|
|
return fd, nil
|
|
}
|
|
|
|
func collect(interfaceNames []string, duration time.Duration, limit int) payload {
|
|
result := payload{
|
|
Diagnostics: diagnostics{
|
|
AttachMode: "af_packet_raw_socket",
|
|
InterfacesRequested: interfaceNames,
|
|
Errors: []string{},
|
|
},
|
|
}
|
|
type socketInfo struct {
|
|
name string
|
|
fd int
|
|
}
|
|
sockets := []socketInfo{}
|
|
for _, name := range interfaceNames {
|
|
if strings.TrimSpace(name) == "" {
|
|
continue
|
|
}
|
|
fd, err := openSocket(strings.TrimSpace(name))
|
|
if err != nil {
|
|
result.Diagnostics.Errors = append(result.Diagnostics.Errors, fmt.Sprintf("%s: %v", name, err))
|
|
continue
|
|
}
|
|
sockets = append(sockets, socketInfo{name: strings.TrimSpace(name), fd: fd})
|
|
result.Diagnostics.InterfacesAttached = append(result.Diagnostics.InterfacesAttached, strings.TrimSpace(name))
|
|
}
|
|
defer func() {
|
|
for _, socket := range sockets {
|
|
_ = syscall.Close(socket.fd)
|
|
}
|
|
}()
|
|
if len(sockets) == 0 {
|
|
return result
|
|
}
|
|
|
|
fdToSocket := map[int]socketInfo{}
|
|
maxFd := 0
|
|
for _, socket := range sockets {
|
|
fdToSocket[socket.fd] = socket
|
|
if socket.fd > maxFd {
|
|
maxFd = socket.fd
|
|
}
|
|
}
|
|
flows := map[flowKey]*flowValue{}
|
|
deadline := time.Now().Add(duration)
|
|
buffer := make([]byte, maxFrameSize)
|
|
for time.Now().Before(deadline) {
|
|
var readfds syscall.FdSet
|
|
for _, socket := range sockets {
|
|
setFd(socket.fd, &readfds)
|
|
}
|
|
timeout := syscall.NsecToTimeval(int64(250 * time.Millisecond))
|
|
_, err := syscall.Select(maxFd+1, &readfds, nil, nil, &timeout)
|
|
if err != nil && err != syscall.EINTR {
|
|
result.Diagnostics.Errors = append(result.Diagnostics.Errors, err.Error())
|
|
break
|
|
}
|
|
for fd, socket := range fdToSocket {
|
|
if !isSet(fd, &readfds) {
|
|
continue
|
|
}
|
|
n, from, err := syscall.Recvfrom(fd, buffer, 0)
|
|
if err != nil {
|
|
if err != syscall.EAGAIN && err != syscall.EWOULDBLOCK {
|
|
result.Diagnostics.Errors = append(result.Diagnostics.Errors, fmt.Sprintf("%s: %v", socket.name, err))
|
|
}
|
|
continue
|
|
}
|
|
key, bytes, ok := parsePacket(buffer[:n])
|
|
if !ok {
|
|
continue
|
|
}
|
|
key.Interface = socket.name
|
|
key.VMID, key.NIC = interfaceMeta(socket.name)
|
|
key.Direction = "ingress"
|
|
if link, ok := from.(*syscall.SockaddrLinklayer); ok && link.Pkttype == packetOut {
|
|
key.Direction = "egress"
|
|
}
|
|
current := flows[key]
|
|
if current == nil {
|
|
current = &flowValue{
|
|
SourceIP: key.SourceIP,
|
|
DestinationIP: key.DestinationIP,
|
|
Protocol: key.Protocol,
|
|
SourcePort: intPtr(key.SourcePort),
|
|
DestinationPort: intPtr(key.DestinationPort),
|
|
VMID: key.VMID,
|
|
NIC: key.NIC,
|
|
Interface: key.Interface,
|
|
Direction: key.Direction,
|
|
State: "observed",
|
|
Collector: "ebpf-helper",
|
|
}
|
|
flows[key] = current
|
|
}
|
|
current.Packets++
|
|
current.Bytes += uint64(bytes)
|
|
if len(flows) >= limit {
|
|
break
|
|
}
|
|
}
|
|
if len(flows) >= limit {
|
|
break
|
|
}
|
|
}
|
|
for _, flow := range flows {
|
|
result.Flows = append(result.Flows, *flow)
|
|
}
|
|
sort.Slice(result.Flows, func(i, j int) bool {
|
|
if result.Flows[i].Bytes == result.Flows[j].Bytes {
|
|
return result.Flows[i].Packets > result.Flows[j].Packets
|
|
}
|
|
return result.Flows[i].Bytes > result.Flows[j].Bytes
|
|
})
|
|
if len(result.Flows) > limit {
|
|
result.Flows = result.Flows[:limit]
|
|
}
|
|
return result
|
|
}
|
|
|
|
func main() {
|
|
jsonOutput := flag.Bool("json", false, "print JSON output")
|
|
limit := flag.Int("limit", 1500, "maximum unique flows")
|
|
durationSeconds := flag.Int("duration", 10, "collection duration in seconds")
|
|
interfaces := flag.String("interfaces", "", "comma-separated interface names")
|
|
flag.Parse()
|
|
if !*jsonOutput {
|
|
fmt.Fprintln(os.Stderr, "only --json output is supported")
|
|
os.Exit(2)
|
|
}
|
|
if *limit <= 0 {
|
|
*limit = 1500
|
|
}
|
|
if *durationSeconds <= 0 {
|
|
*durationSeconds = 10
|
|
}
|
|
names := []string{}
|
|
for _, name := range strings.Split(*interfaces, ",") {
|
|
if strings.TrimSpace(name) != "" {
|
|
names = append(names, strings.TrimSpace(name))
|
|
}
|
|
}
|
|
result := collect(names, time.Duration(*durationSeconds)*time.Second, *limit)
|
|
data, err := json.Marshal(result)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "json marshal failed: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
fmt.Println(string(data))
|
|
}
|