Merge pull request #47 from pouriyajamshidi/feat/add_ip_and_port_filt…

…ering

feat/implement IP and port filtering

README.md

@@ -8,7 +8,22 @@ This repo is the companion to my blog posts about eBPF at <https://thegraynode.i
 
 ## Running The Program
 
-Clone the repository.
+You can install **flat** in two ways.
+
+1. Download the [pre-compiled binary](#download-the-pre-compiled-binary)
+2. Compile from [source](#compile-from-source)
+
+### Download The Pre-compiled Binary
+
+```bash
+wget https://github.com/pouriyajamshidi/flat/releases/latest/download/flat
+```
+
+Then check out the [examples](#examples).
+
+### Compile From Source
+
+Clone the repository:
 
 ```bash
 git clone https://github.com/pouriyajamshidi/flat .
@@ -32,13 +47,34 @@ Compile the **Go** program:
 go build -ldflags "-s -w" -o flat cmd/flat.go
 ```
 
+### Examples
+
 Run it with elevated privileges:
 
 ```bash
 # Replace eth0 with your desired interface name
 sudo ./flat -i eth0
+# Or
+sudo ./flat -i eth0 -ip 1.1.1.1
+# Or
+sudo ./flat -i eth0 -port 53
+# Or
+sudo ./flat -i eth0 -ip 1.1.1.1 -port 53
 ```
 
+## Flags
+
+**flat** supports four flags at the moment:
+
+| flag  | Description                         |
+| ----- | ----------------------------------- |
+| -i    | interface to attach the probe to    |
+| -ip   | IP address to filter on (optional)  |
+| -port | Port number to filter on (optional) |
+| -h    | Show help message                   |
+
+---
+
 ## Acknolegments
 
 Heavily inspired by [flowlat](https://github.com/markpash/flowlat).

cmd/flat.go

@@ -6,11 +6,13 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"net/netip"
 	"os"
 	"os/signal"
 	"syscall"
 
 	"github.com/pouriyajamshidi/flat/internal/probe"
+	"github.com/pouriyajamshidi/flat/internal/types"
 	"github.com/vishvananda/netlink"
 )
 
@@ -41,23 +43,58 @@ func displayInterfaces() {
 	os.Exit(1)
 }
 
-func main() {
+// getUserInput gets and validates user input
+func getUserInput() types.UserInput {
 	ifaceFlag := flag.String("i", "eth0", "interface to attach the probe to")
+	ipFlag := flag.String("ip", "", "IP address to track (optional)")
+	portFlag := flag.Uint("port", 0, "Port number to track (optional)")
 	flag.Parse()
 
+	var userInput types.UserInput
+
 	iface, err := netlink.LinkByName(*ifaceFlag)
 
 	if err != nil {
 		log.Printf("Could not find interface %v: %v", *ifaceFlag, err)
 		displayInterfaces()
 	}
 
+	userInput.Interface = iface
+
+	if *ipFlag != "" {
+		userInput.IP, err = netip.ParseAddr(*ipFlag)
+
+		if err != nil {
+			log.Printf("Could not parse IP address %v: %v", *ipFlag, err)
+			os.Exit(1)
+		}
+
+		log.Printf("Filtering results on IP %v", userInput.IP)
+	}
+
+	if *portFlag != 0 {
+		if *portFlag < 1 || *portFlag > 65535 {
+			log.Printf("Could not parse port %v: %v", *portFlag, err)
+			os.Exit(1)
+		}
+
+		userInput.Port = uint16(*portFlag)
+
+		log.Printf("Filtering results on port %d", userInput.Port)
+	}
+
+	return userInput
+}
+
+func main() {
+	userInput := getUserInput()
+
 	ctx := context.Background()
 	ctx, cancel := context.WithCancel(ctx)
 
 	signalHandler(cancel)
 
-	if err := probe.Run(ctx, iface); err != nil {
+	if err := probe.Run(ctx, userInput); err != nil {
 		log.Fatalf("Failed running the probe: %v", err)
 	}
 }

go.mod

@@ -1,21 +1,21 @@
 module github.com/pouriyajamshidi/flat
 
-go 1.22
+go 1.23.2
 
 require (
 	github.com/cilium/ebpf v0.16.0
 	github.com/google/gopacket v1.1.19
 	github.com/gookit/color v1.5.4
 	github.com/stretchr/testify v1.9.0
 	github.com/vishvananda/netlink v1.3.0
-	golang.org/x/sys v0.26.0
+	golang.org/x/sys v0.27.0
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/vishvananda/netns v0.0.4 // indirect
+	github.com/vishvananda/netns v0.0.5 // indirect
 	github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
-	golang.org/x/exp v0.0.0-20240416160154-fe59bbe5cc7f // indirect
+	golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -22,14 +22,15 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/vishvananda/netlink v1.3.0 h1:X7l42GfcV4S6E4vHTsw48qbrV+9PVojNfIhZcwQdrZk=
 github.com/vishvananda/netlink v1.3.0/go.mod h1:i6NetklAujEcC6fK0JPjT8qSwWyO0HLn4UKG+hGqeJs=
-github.com/vishvananda/netns v0.0.4 h1:Oeaw1EM2JMxD51g9uhtC0D7erkIjgmj8+JZc26m1YX8=
 github.com/vishvananda/netns v0.0.4/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM=
+github.com/vishvananda/netns v0.0.5 h1:DfiHV+j8bA32MFM7bfEunvT8IAqQ/NzSJHtcmW5zdEY=
+github.com/vishvananda/netns v0.0.5/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM=
 github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
 github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
-golang.org/x/exp v0.0.0-20240416160154-fe59bbe5cc7f h1:99ci1mjWVBWwJiEKYY6jWa4d2nTQVIEhZIptnrVb1XY=
-golang.org/x/exp v0.0.0-20240416160154-fe59bbe5cc7f/go.mod h1:/lliqkxwWAhPjf5oSOIJup2XcqJaw8RGS6k3TGEc7GI=
+golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c h1:7dEasQXItcW1xKJ2+gg5VOiBnqWrJc+rq0DPKyvvdbY=
+golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c/go.mod h1:NQtJDoLvd6faHhE7m4T/1IY708gDefGGjR/iUW8yQQ8=
 golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
 golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
@@ -39,8 +40,8 @@ golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo=
-golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.27.0 h1:wBqf8DvsY9Y/2P8gAfPDEYNuS30J4lPHJxXSb/nJZ+s=
+golang.org/x/sys v0.27.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

internal/probe/probe.go

@@ -8,6 +8,7 @@ import (
 	"github.com/pouriyajamshidi/flat/clsact"
 	"github.com/pouriyajamshidi/flat/internal/flowtable"
 	"github.com/pouriyajamshidi/flat/internal/packet"
+	"github.com/pouriyajamshidi/flat/internal/types"
 	"github.com/vishvananda/netlink"
 	"golang.org/x/sys/unix"
 )
@@ -180,7 +181,7 @@ func (p *probe) Close() error {
 
 // Run attaches the probe, reads from the eBPF map
 // as well as calculating and displaying the flow latencies
-func Run(ctx context.Context, iface netlink.Link) error {
+func Run(ctx context.Context, userInput types.UserInput) error {
 	log.Println("Starting up the probe")
 
 	if err := setRlimit(); err != nil {
@@ -196,7 +197,7 @@ func Run(ctx context.Context, iface netlink.Link) error {
 		}
 	}()
 
-	probe, err := newProbe(iface)
+	probe, err := newProbe(userInput.Interface)
 
 	if err != nil {
 		return err
@@ -235,7 +236,16 @@ func Run(ctx context.Context, iface netlink.Link) error {
 				log.Printf("Could not unmarshall packet: %+v", pkt)
 				continue
 			}
-			packet.CalcLatency(packetAttrs, flowtable)
+			if !userInput.IP.IsValid() && userInput.Port == 0 {
+				packet.CalcLatency(packetAttrs, flowtable)
+				// FIXME:
+				// } else if (userInput.Port == packetAttrs.DstPort || userInput.Port == packetAttrs.SrcPort) && (userInput.IP == packetAttrs.DstIP.Unmap() || userInput.IP == packetAttrs.SrcIP.Unmap()) {
+				// 	packet.CalcLatency(packetAttrs, flowtable)
+			} else if userInput.IP == packetAttrs.DstIP.Unmap() || userInput.IP == packetAttrs.SrcIP.Unmap() {
+				packet.CalcLatency(packetAttrs, flowtable)
+			} else if userInput.Port == packetAttrs.DstPort || userInput.Port == packetAttrs.SrcPort {
+				packet.CalcLatency(packetAttrs, flowtable)
+			}
 		}
 	}
 }

internal/types/types.go

@@ -0,0 +1,14 @@
+package types
+
+import (
+	"net/netip"
+
+	"github.com/vishvananda/netlink"
+)
+
+// UserInput holds the information provided through flags
+type UserInput struct {
+	Interface netlink.Link
+	IP        netip.Addr
+	Port      uint16
+}