Add AF_PACKET engine

Co-authored-by: Leonid Bugaev <leonsbox@gmail.com>

Added a way to capture multiple ports at the same time, with a single listener. 
Ports should be separated by comma like this: --input-raw :80,8080,3000

.gitignore

@@ -2,6 +2,7 @@
 *.gor
 *.rpm
 *.dep
+*.deb
 *.pkg
 *.exe
 

Makefile

@@ -59,16 +59,16 @@ release-mac:
 	rm -rf /tmp/gor-build
 
 release-windows:
-	echo $(pwd)
 	docker run -it --rm \
 	  -v `pwd`:/go/src/github.com/buger/goreplay \
 	  -w /go/src/github.com/buger/goreplay \
 	  -e CGO_ENABLED=1 \
 	  docker.elastic.co/beats-dev/golang-crossbuild:1.16.4-main \
-	  --build-cmd "VERSION=make build" \
+	  --build-cmd "make VERSION=$(VERSION) build" \
 	  -p "windows/amd64"
-
-	mv ./gor ./gor-$(VERSION)$(PREFIX).exe
+	mv ./gor ./gor.exe
+	zip gor-$(VERSION)$(PREFIX)_windows.zip ./gor.exe
+	rm -rf ./gor.exe
 
 build:
 	go build -o $(BIN_NAME) $(LDFLAGS)

capture/af_packet.go

@@ -0,0 +1,180 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"log"
+	"os"
+	"runtime/pprof"
+	"time"
+
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/afpacket"
+	"github.com/google/gopacket/layers"
+	"github.com/google/gopacket/pcap"
+	"golang.org/x/net/bpf"
+
+	_ "github.com/google/gopacket/layers"
+)
+
+var (
+	iface      = flag.String("i", "any", "Interface to read from")
+	cpuprofile = flag.String("cpuprofile", "", "If non-empty, write CPU profile here")
+	snaplen    = flag.Int("s", 0, "Snaplen, if <= 0, use 65535")
+	bufferSize = flag.Int("b", 8, "Interface buffersize (MB)")
+	filter     = flag.String("f", "port not 22", "BPF filter")
+	count      = flag.Int64("c", -1, "If >= 0, # of packets to capture before returning")
+	verbose    = flag.Int64("log_every", 1, "Write a log every X packets")
+	addVLAN    = flag.Bool("add_vlan", false, "If true, add VLAN header")
+)
+
+type afpacketHandle struct {
+	TPacket *afpacket.TPacket
+}
+
+func newAfpacketHandle(device string, snaplen int, block_size int, num_blocks int,
+	useVLAN bool, timeout time.Duration) (*afpacketHandle, error) {
+
+	h := &afpacketHandle{}
+	var err error
+
+	if device == "any" {
+		h.TPacket, err = afpacket.NewTPacket(
+			afpacket.OptFrameSize(snaplen),
+			afpacket.OptBlockSize(block_size),
+			afpacket.OptNumBlocks(num_blocks),
+			afpacket.OptAddVLANHeader(useVLAN),
+			afpacket.OptPollTimeout(timeout),
+			afpacket.SocketRaw,
+			afpacket.TPacketVersion3)
+	} else {
+		h.TPacket, err = afpacket.NewTPacket(
+			afpacket.OptInterface(device),
+			afpacket.OptFrameSize(snaplen),
+			afpacket.OptBlockSize(block_size),
+			afpacket.OptNumBlocks(num_blocks),
+			afpacket.OptAddVLANHeader(useVLAN),
+			afpacket.OptPollTimeout(timeout),
+			afpacket.SocketRaw,
+			afpacket.TPacketVersion3)
+	}
+	return h, err
+}
+
+// ZeroCopyReadPacketData satisfies ZeroCopyPacketDataSource interface
+func (h *afpacketHandle) ZeroCopyReadPacketData() (data []byte, ci gopacket.CaptureInfo, err error) {
+	return h.TPacket.ZeroCopyReadPacketData()
+}
+
+// SetBPFFilter translates a BPF filter string into BPF RawInstruction and applies them.
+func (h *afpacketHandle) SetBPFFilter(filter string, snaplen int) (err error) {
+	pcapBPF, err := pcap.CompileBPFFilter(layers.LinkTypeEthernet, snaplen, filter)
+	if err != nil {
+		return err
+	}
+	bpfIns := []bpf.RawInstruction{}
+	for _, ins := range pcapBPF {
+		bpfIns2 := bpf.RawInstruction{
+			Op: ins.Code,
+			Jt: ins.Jt,
+			Jf: ins.Jf,
+			K:  ins.K,
+		}
+		bpfIns = append(bpfIns, bpfIns2)
+	}
+	if h.TPacket.SetBPF(bpfIns); err != nil {
+		return err
+	}
+	return h.TPacket.SetBPF(bpfIns)
+}
+
+// LinkType returns ethernet link type.
+func (h *afpacketHandle) LinkType() layers.LinkType {
+	return layers.LinkTypeEthernet
+}
+
+// Close will close afpacket source.
+func (h *afpacketHandle) Close() {
+	h.TPacket.Close()
+}
+
+// SocketStats prints received, dropped, queue-freeze packet stats.
+func (h *afpacketHandle) SocketStats() (as afpacket.SocketStats, asv afpacket.SocketStatsV3, err error) {
+	return h.TPacket.SocketStats()
+}
+
+// afpacketComputeSize computes the block_size and the num_blocks in such a way that the
+// allocated mmap buffer is close to but smaller than target_size_mb.
+// The restriction is that the block_size must be divisible by both the
+// frame size and page size.
+func afpacketComputeSize(targetSizeMb int, snaplen int, pageSize int) (
+	frameSize int, blockSize int, numBlocks int, err error) {
+
+	if snaplen < pageSize {
+		frameSize = pageSize / (pageSize / snaplen)
+	} else {
+		frameSize = (snaplen/pageSize + 1) * pageSize
+	}
+
+	// 128 is the default from the gopacket library so just use that
+	blockSize = frameSize * 128
+	numBlocks = (targetSizeMb * 1024 * 1024) / blockSize
+
+	if numBlocks == 0 {
+		return 0, 0, 0, fmt.Errorf("Interface buffersize is too small")
+	}
+
+	return frameSize, blockSize, numBlocks, nil
+}
+
+func main() {
+	flag.Parse()
+	if *cpuprofile != "" {
+		log.Printf("Writing CPU profile to %q", *cpuprofile)
+		f, err := os.Create(*cpuprofile)
+		if err != nil {
+			log.Fatal(err)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			log.Fatal(err)
+		}
+		defer pprof.StopCPUProfile()
+	}
+	log.Printf("Starting on interface %q", *iface)
+	if *snaplen <= 0 {
+		*snaplen = 65535
+	}
+	szFrame, szBlock, numBlocks, err := afpacketComputeSize(*bufferSize, *snaplen, os.Getpagesize())
+	if err != nil {
+		log.Fatal(err)
+	}
+	afpacketHandle, err := newAfpacketHandle(*iface, szFrame, szBlock, numBlocks, *addVLAN, pcap.BlockForever)
+	if err != nil {
+		log.Fatal(err)
+	}
+	err = afpacketHandle.SetBPFFilter(*filter, *snaplen)
+	if err != nil {
+		log.Fatal(err)
+	}
+	source := gopacket.ZeroCopyPacketDataSource(afpacketHandle)
+	defer afpacketHandle.Close()
+
+	bytes := uint64(0)
+	packets := uint64(0)
+	for ; *count != 0; *count-- {
+		data, _, err := source.ZeroCopyReadPacketData()
+		if err != nil {
+			log.Fatal(err)
+		}
+		bytes += uint64(len(data))
+		packets++
+		if *count%*verbose == 0 {
+			_, afpacketStats, err := afpacketHandle.SocketStats()
+			if err != nil {
+				log.Println(err)
+			}
+			log.Printf("Read in %d bytes in %d packets", bytes, packets)
+			log.Printf("Stats {received dropped queue-freeze}: %d", afpacketStats)
+		}
+	}
+}

capture/capture.go

@@ -43,12 +43,12 @@ type Listener struct {
 	Transport  string       // transport layer default to tcp
 	Activate   func() error // function is used to activate the engine. it must be called before reading packets
 	Handles    map[string]gopacket.ZeroCopyPacketDataSource
-	Interfaces []net.Interface
+	Interfaces []pcap.Interface
 	loopIndex  int
 	Reading    chan bool // this channel is closed when the listener has started reading packets
 	PcapOptions
 	Engine        EngineType
-	port          uint16 // src or/and dst port
+	ports         []uint16 // src or/and dst ports
 	trackResponse bool
 
 	host string // pcap file name or interface (name, hardware addr, index or ip address)
@@ -65,6 +65,7 @@ const (
 	EnginePcap EngineType = 1 << iota
 	EnginePcapFile
 	EngineRawSocket
+	EngineAFPacket
 )
 
 // Set is here so that EngineType can implement flag.Var
@@ -74,8 +75,10 @@ func (eng *EngineType) Set(v string) error {
 		*eng = EnginePcap
 	case "pcap_file":
 		*eng = EnginePcapFile
-	case "raw_socket", "af_packet":
+	case "raw_socket":
 		*eng = EngineRawSocket
+	case "af_packet":
+		*eng = EngineAFPacket
 	default:
 		return fmt.Errorf("invalid engine %s", v)
 	}
@@ -90,6 +93,8 @@ func (eng *EngineType) String() (e string) {
 		e = "libpcap"
 	case EngineRawSocket:
 		e = "raw_socket"
+	case EngineAFPacket:
+		e = "af_packet"
 	default:
 		e = ""
 	}
@@ -99,11 +104,15 @@ func (eng *EngineType) String() (e string) {
 // NewListener creates and initialize a new Listener. if transport or/and engine are invalid/unsupported
 // is "tcp" and "pcap", are assumed. l.Engine and l.Transport can help to get the values used.
 // if there is an error it will be associated with getting network interfaces
-func NewListener(host string, port uint16, transport string, engine EngineType, trackResponse bool) (l *Listener, err error) {
+func NewListener(host string, ports []uint16, transport string, engine EngineType, trackResponse bool) (l *Listener, err error) {
 	l = &Listener{}
 
 	l.host = host
-	l.port = port
+	if l.host == "localhost" {
+		l.host = "127.0.0.1"
+	}
+	l.ports = ports
+
 	l.Transport = "tcp"
 	if transport != "" {
 		l.Transport = transport
@@ -120,11 +129,15 @@ func NewListener(host string, port uint16, transport string, engine EngineType,
 	case EngineRawSocket:
 		l.Engine = EngineRawSocket
 		l.Activate = l.activateRawSocket
+	case EngineAFPacket:
+		l.Engine = EngineAFPacket
+		l.Activate = l.activateAFPacket
 	case EnginePcapFile:
 		l.Engine = EnginePcapFile
 		l.Activate = l.activatePcapFile
 		return
 	}
+
 	err = l.setInterfaces()
 	if err != nil {
 		return nil, err
@@ -168,7 +181,7 @@ func (l *Listener) ListenBackground(ctx context.Context, handler PacketHandler)
 
 // Filter returns automatic filter applied by goreplay
 // to a pcap handle of a specific interface
-func (l *Listener) Filter(ifi net.Interface) (filter string) {
+func (l *Listener) Filter(ifi pcap.Interface) (filter string) {
 	// https://www.tcpdump.org/manpages/pcap-filter.7.html
 
 	hosts := []string{l.host}
@@ -176,7 +189,7 @@ func (l *Listener) Filter(ifi net.Interface) (filter string) {
 		hosts = interfaceAddresses(ifi)
 	}
 
-	filter = portsFilter(l.Transport, "dst", l.port)
+	filter = portsFilter(l.Transport, "dst", l.ports)
 
 	if len(hosts) != 0 {
 		filter = fmt.Sprintf("((%s) and (%s))", filter, hostsFilter("dst", hosts))
@@ -185,7 +198,7 @@ func (l *Listener) Filter(ifi net.Interface) (filter string) {
 	}
 
 	if l.trackResponse {
-		responseFilter := portsFilter(l.Transport, "src", l.port)
+		responseFilter := portsFilter(l.Transport, "src", l.ports)
 
 		if len(hosts) != 0 {
 			responseFilter = fmt.Sprintf("((%s) and (%s))", responseFilter, hostsFilter("src", hosts))
@@ -218,7 +231,7 @@ func (l *Listener) Filter(ifi net.Interface) (filter string) {
 
 // PcapHandle returns new pcap Handle from dev on success.
 // this function should be called after setting all necessary options for this listener
-func (l *Listener) PcapHandle(ifi net.Interface) (handle *pcap.Handle, err error) {
+func (l *Listener) PcapHandle(ifi pcap.Interface) (handle *pcap.Handle, err error) {
 	var inactive *pcap.InactiveHandle
 	inactive, err = pcap.NewInactiveHandle(ifi.Name)
 	if err != nil {
@@ -243,12 +256,22 @@ func (l *Listener) PcapHandle(ifi net.Interface) (handle *pcap.Handle, err error
 			return nil, fmt.Errorf("monitor mode error: %q, interface: %q", err, ifi.Name)
 		}
 	}
+
 	var snap int
-	if l.Snaplen {
-		snap = 64<<10 + 200
-	} else if ifi.MTU > 0 {
-		snap = ifi.MTU + 200
+
+	if !l.Snaplen {
+		infs, _ := net.Interfaces()
+		for _, i := range infs {
+			if i.Name == ifi.Name {
+				snap = i.MTU + 200
+			}
+		}
 	}
+
+	if snap == 0 {
+		snap = 64<<10 + 200
+	}
+
 	err = inactive.SetSnapLen(snap)
 	if err != nil {
 		return nil, fmt.Errorf("snapshot length error: %q, interface: %q", err, ifi.Name)
@@ -281,7 +304,7 @@ func (l *Listener) PcapHandle(ifi net.Interface) (handle *pcap.Handle, err error
 }
 
 // SocketHandle returns new unix ethernet handle associated with this listener settings
-func (l *Listener) SocketHandle(ifi net.Interface) (handle Socket, err error) {
+func (l *Listener) SocketHandle(ifi pcap.Interface) (handle Socket, err error) {
 	handle, err = NewSocket(ifi)
 	if err != nil {
 		return nil, fmt.Errorf("sock raw error: %q, interface: %q", err, ifi.Name)
@@ -388,6 +411,36 @@ func (l *Listener) activatePcap() error {
 	return nil
 }
 
+func (l *Listener) activateAFPacket() error {
+	var e error
+	var msg string
+
+	szFrame, szBlock, numBlocks, err := afpacketComputeSize(*bufferSize, *snaplen, os.Getpagesize())
+	if l.host == "" || l.host == "any" {
+		handle, err := newAfpacketHandle("any", szFrame, szBlock, numBlocks, false, pcap.BlockForever)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
+
+	for _, ifi := range l.Interfaces {
+		var handle *pcap.Handle
+
+		handle, err := newAfpacketHandle(ifi.Name, szFrame, szBlock, numBlocks, false, pcap.BlockForever)
+
+		if err != nil {
+			msg += ("\n" + err.Error())
+			continue
+		}
+		l.Handles[ifi.Name] = handle
+	}
+	if len(l.Handles) == 0 {
+		return fmt.Errorf("pcap handles error:%s", msg)
+	}
+	return nil
+}
+
 func (l *Listener) activateRawSocket() error {
 	if runtime.GOOS != "linux" {
 		return fmt.Errorf("sock_raw is not stabilized on OS other than linux")
@@ -418,7 +471,7 @@ func (l *Listener) activatePcapFile() (err error) {
 
 	tmp := l.host
 	l.host = ""
-	l.BPFFilter = l.Filter(net.Interface{})
+	l.BPFFilter = l.Filter(pcap.Interface{})
 	l.host = tmp
 
 	if e = handle.SetBPFFilter(l.BPFFilter); e != nil {
@@ -430,60 +483,50 @@ func (l *Listener) activatePcapFile() (err error) {
 }
 
 func (l *Listener) setInterfaces() (err error) {
-	var ifis []net.Interface
-	ifis, err = net.Interfaces()
+	var pifis []pcap.Interface
+	pifis, err = pcap.FindAllDevs()
+	ifis, _ := net.Interfaces()
 	if err != nil {
 		return
 	}
-	for i := range ifis {
-		if ifis[i].Flags&net.FlagLoopback != 0 {
-			l.loopIndex = ifis[i].Index
+	for _, pi := range pifis {
+		var ni net.Interface
+		for _, i := range ifis {
+			if i.Name == pi.Name {
+				ni = i
+				break
+			}
 		}
-		if ifis[i].Flags&net.FlagUp == 0 {
+
+		if net.Flags(pi.Flags)&net.FlagLoopback != 0 {
+			l.loopIndex = ni.Index
+		}
+		if net.Flags(pi.Flags)&net.FlagUp == 0 {
 			continue
 		}
-		if isDevice(l.host, ifis[i]) {
-			l.Interfaces = []net.Interface{ifis[i]}
+		if isDevice(l.host, pi) {
+			l.Interfaces = []pcap.Interface{pi}
 			return
 		}
-		addrs, e := ifis[i].Addrs()
-		if e != nil {
-			// don't give up on a failure from a single interface
-			continue
-		}
-		for _, addr := range addrs {
-			if cutMask(addr) == l.host {
-				l.Interfaces = []net.Interface{ifis[i]}
+		for _, addr := range pi.Addresses {
+			if addr.IP.String() == l.host {
+				l.Interfaces = []pcap.Interface{pi}
 				return
 			}
 		}
 	}
-	l.Interfaces = ifis
+	l.Interfaces = pifis
 	return
 }
 
-func cutMask(addr net.Addr) string {
-	mask := addr.String()
-	for i, v := range mask {
-		if v == '/' {
-			return mask[:i]
-		}
-	}
-	return mask
+func isDevice(addr string, ifi pcap.Interface) bool {
+	return addr == ifi.Name
 }
 
-func isDevice(addr string, ifi net.Interface) bool {
-	return addr == ifi.Name || addr == fmt.Sprintf("%d", ifi.Index) || (addr != "" && addr == ifi.HardwareAddr.String())
-}
-
-func interfaceAddresses(ifi net.Interface) []string {
+func interfaceAddresses(ifi pcap.Interface) []string {
 	var hosts []string
-	if addrs, err := ifi.Addrs(); err == nil {
-		for _, addr := range addrs {
-			if ip := addr.(*net.IPNet).IP.To16(); ip != nil {
-				hosts = append(hosts, ip.String())
-			}
-		}
+	for _, addr := range ifi.Addresses {
+		hosts = append(hosts, addr.IP.String())
 	}
 	return hosts
 }
@@ -496,12 +539,16 @@ func listenAll(addr string) bool {
 	return false
 }
 
-func portsFilter(transport string, direction string, port uint16) string {
-	if port == 0 {
+func portsFilter(transport string, direction string, ports []uint16) string {
+	if len(ports) == 0 || ports[0] == 0 {
 		return fmt.Sprintf("%s %s portrange 0-%d", transport, direction, 1<<16-1)
 	}
 
-	return fmt.Sprintf("%s %s port %d", transport, direction, port)
+	var filters []string
+	for _, port := range ports {
+		filters = append(filters, fmt.Sprintf("%s %s port %d", transport, direction, port))
+	}
+	return strings.Join(filters, " or ")
 }
 
 func hostsFilter(direction string, hosts []string) string {

capture/sock_linux.go

@@ -44,7 +44,23 @@ type SockRaw struct {
 }
 
 // NewSocket returns new M'maped sock_raw on packet version 2.
-func NewSocket(ifi net.Interface) (*SockRaw, error) {
+func NewSocket(pifi pcap.Interface) (*SockRaw, error) {
+	var ifi net.Interface
+
+	infs, _ := net.Interfaces()
+	found := false
+	for _, i := range infs {
+		if i.Name == pifi.Name {
+			ifi = i
+			found = true
+			break
+		}
+	}
+
+	if !found {
+		return nil, fmt.Errorf("Can't find matching interface")
+	}
+
 	// sock create
 	fd, err := unix.Socket(unix.AF_PACKET, unix.SOCK_RAW, int(ETHALL))
 	if err != nil {

capture/sock_others.go

@@ -4,10 +4,11 @@ package capture
 
 import (
 	"errors"
-	"net"
+
+	"github.com/google/gopacket/pcap"
 )
 
 // NewSocket returns new M'maped sock_raw on packet version 2.
-func NewSocket(_ net.Interface) (Socket, error) {
+func NewSocket(_ pcap.Interface) (Socket, error) {
 	return nil, errors.New("afpacket socket is only available on linux")
 }

input_raw.go

@@ -6,6 +6,7 @@ import (
 	"log"
 	"net"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -61,7 +62,7 @@ type RAWInputConfig struct {
 	Stats          bool               `json:"input-raw-stats"`
 	quit           chan bool          // Channel used only to indicate goroutine should shutdown
 	host           string
-	port           uint16
+	ports          []uint16
 }
 
 // RAWInput used for intercepting traffic for given address
@@ -81,22 +82,28 @@ func NewRAWInput(address string, config RAWInputConfig) (i *RAWInput) {
 	i.RAWInputConfig = config
 	i.message = make(chan *tcp.Message, 10000)
 	i.quit = make(chan bool)
-	var host, _port string
-	var err error
-	var port int
-	host, _port, err = net.SplitHostPort(address)
+
+	host, _ports, err := net.SplitHostPort(address)
 	if err != nil {
 		log.Fatalf("input-raw: error while parsing address: %s", err)
 	}
-	if _port != "" {
-		port, err = strconv.Atoi(_port)
+
+	var ports []uint16
+	if _ports != "" {
+		portsStr := strings.Split(_ports, ",")
+
+		for _, portStr := range portsStr {
+			port, err := strconv.Atoi(strings.TrimSpace(portStr))
+			if err != nil {
+				log.Fatalf("parsing port error: %v", err)
+			}
+			ports = append(ports, uint16(port))
+
+		}
 	}
 
-	if err != nil {
-		log.Fatalf("parsing port error: %v", err)
-	}
 	i.host = host
-	i.port = uint16(port)
+	i.ports = ports
 
 	i.listen(address)
 
@@ -141,7 +148,7 @@ func (i *RAWInput) PluginRead() (*Message, error) {
 
 func (i *RAWInput) listen(address string) {
 	var err error
-	i.listener, err = capture.NewListener(i.host, i.port, "", i.Engine, i.TrackResponse)
+	i.listener, err = capture.NewListener(i.host, i.ports, "", i.Engine, i.TrackResponse)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -172,7 +179,7 @@ func (i *RAWInput) messageEmitter(m *tcp.Message) {
 }
 
 func (i *RAWInput) String() string {
-	return fmt.Sprintf("Intercepting traffic from: %s:%d", i.host, i.port)
+	return fmt.Sprintf("Intercepting traffic from: %s:%s", i.host, strings.Join(strings.Fields(fmt.Sprint(i.ports)), ","))
 }
 
 // GetStats returns the stats so far and reset the stats

proto/proto.go

@@ -19,6 +19,7 @@ package proto
 import (
 	"bufio"
 	"bytes"
+	_ "fmt"
 	"net/http"
 	"net/textproto"
 	"strings"

tcp/tcp_message.go

@@ -3,6 +3,7 @@ package tcp
 import (
 	"encoding/binary"
 	"encoding/hex"
+	_ "fmt"
 	"sort"
 	"time"
 
@@ -112,11 +113,7 @@ func (m *Message) MissingChunk() bool {
 }
 
 func (m *Message) PacketData() [][]byte {
-	var totalLen int
-	for _, p := range m.packets {
-		totalLen += len(p.Payload)
-	}
-	tmp := make([][]byte, totalLen)
+	tmp := make([][]byte, len(m.packets))
 
 	for i, p := range m.packets {
 		tmp[i] = p.Payload