Merge pull request #36 from chazlever/tightloop

Faster Packet Loop
chazlever · Jul 11, 2020 · ead4311 · ead4311
2 parents ec1cd18 + 5da2269
commit ead4311
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 61 deletions.
diff --git a/main.go b/main.go
@@ -58,7 +58,7 @@ func main() {
 	app := cli.NewApp()
 	app.Name = "rickybobby"
 	app.Usage = "Parsing DNS packets when you wanna GO fast!"
-	app.Version = "1.0.0"
+	app.Version = "1.0.1"
 	app.Compiled = time.Now()
 
 	app.Authors = []cli.Author{

diff --git a/parser/parse.go b/parser/parse.go
@@ -3,13 +3,15 @@ package parser
 import (
 	"crypto/sha256"
 	"fmt"
+	"io"
+	"os"
+	"time"
+
 	"github.com/google/gopacket"
 	"github.com/google/gopacket/layers"
 	"github.com/google/gopacket/pcap"
 	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
-	"os"
-	"time"
 )
 
 const (
@@ -84,6 +86,7 @@ func ParseDns(handle *pcap.Handle) {
 		ip6    *layers.IPv6
 		tcp    *layers.TCP
 		udp    *layers.UDP
+		msg    *dns.Msg
 	)
 
 	// Set the source and sensor for packet source
@@ -93,77 +96,88 @@ func ParseDns(handle *pcap.Handle) {
 	// Use the handle as a packet source to process all packets
 	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
 	packetSource.NoCopy = true
+	packetSource.Lazy = true
 
 PACKETLOOP:
-	for packet := range packetSource.Packets() {
+	for {
+		packet, err := packetSource.NextPacket()
+		if err == io.EOF {
+			break
+		}
 		stats.PacketTotal += 1
 
-		// Let's analyze decoded layers
-		var msg *dns.Msg
-		for _, curLayer := range packet.Layers() {
-			switch curLayer.LayerType() {
-			case layers.LayerTypeIPv4:
-				ip4 = curLayer.(*layers.IPv4)
-				schema.SourceAddress = ip4.SrcIP.String()
-				schema.DestinationAddress = ip4.DstIP.String()
-				schema.Ipv4 = true
-				stats.PacketIPv4 += 1
-			case layers.LayerTypeIPv6:
-				ip6 = curLayer.(*layers.IPv6)
-				schema.SourceAddress = ip6.SrcIP.String()
-				schema.DestinationAddress = ip6.DstIP.String()
-				schema.Ipv4 = false
-				stats.PacketIPv6 += 1
-			case layers.LayerTypeTCP:
-				tcp = curLayer.(*layers.TCP)
-				stats.PacketTcp += 1
-
-				if !DoParseTcp {
-					continue PACKETLOOP
-				}
+		if err != nil {
+			log.Errorf("Error decoding some part of the packet: %v\n", err)
+			stats.PacketErrors += 1
+			continue
+		}
 
-				msg = new(dns.Msg)
-				if err := msg.Unpack(tcp.Payload); err != nil {
-					log.Errorf("Could not decode DNS: %v\n", err)
-					stats.PacketErrors += 1
-					continue PACKETLOOP
-				}
-				stats.PacketDns += 1
-
-				schema.SourcePort = uint16(tcp.SrcPort)
-				schema.DestinationPort = uint16(tcp.DstPort)
-				schema.Udp = false
-				schema.Sha256 = fmt.Sprintf("%x", sha256.Sum256(tcp.Payload))
-			case layers.LayerTypeUDP:
-				udp = curLayer.(*layers.UDP)
-				stats.PacketUdp += 1
-
-				msg = new(dns.Msg)
-				if err := msg.Unpack(udp.Payload); err != nil {
-					log.Errorf("Could not decode DNS: %v\n", err)
-					stats.PacketErrors += 1
-					continue PACKETLOOP
-				}
-				stats.PacketDns += 1
+		// Parse DNS and transport layer information
+		msg = nil
+		transportLayer := packet.TransportLayer()
+		switch transportLayer.LayerType() {
+		case layers.LayerTypeTCP:
+			tcp = transportLayer.(*layers.TCP)
+			stats.PacketTcp += 1
 
-				schema.SourcePort = uint16(udp.SrcPort)
-				schema.DestinationPort = uint16(udp.DstPort)
-				schema.Udp = true
-				schema.Sha256 = fmt.Sprintf("%x", sha256.Sum256(udp.Payload))
+			if !DoParseTcp {
+				continue PACKETLOOP
 			}
-		}
 
-		// This means we did not attempt to parse a DNS payload
-		if msg == nil {
-			// Let's check if we had any errors decoding any of the packet layers
-			if err := packet.ErrorLayer(); err != nil {
-				log.Debugf("Error decoding some part of the packet:", err)
+			msg = new(dns.Msg)
+			if err := msg.Unpack(tcp.Payload); err != nil {
+				log.Errorf("Could not decode DNS: %v\n", err)
 				stats.PacketErrors += 1
+				continue PACKETLOOP
 			}
+			stats.PacketDns += 1
+
+			schema.SourcePort = uint16(tcp.SrcPort)
+			schema.DestinationPort = uint16(tcp.DstPort)
+			schema.Udp = false
+			schema.Sha256 = fmt.Sprintf("%x", sha256.Sum256(tcp.Payload))
+		case layers.LayerTypeUDP:
+			udp = transportLayer.(*layers.UDP)
+			stats.PacketUdp += 1
+
+			msg = new(dns.Msg)
+			if err := msg.Unpack(udp.Payload); err != nil {
+				log.Errorf("Could not decode DNS: %v\n", err)
+				stats.PacketErrors += 1
+				continue PACKETLOOP
+			}
+			stats.PacketDns += 1
 
+			schema.SourcePort = uint16(udp.SrcPort)
+			schema.DestinationPort = uint16(udp.DstPort)
+			schema.Udp = true
+			schema.Sha256 = fmt.Sprintf("%x", sha256.Sum256(udp.Payload))
+		}
+
+		// This means we did not attempt to parse a DNS payload and
+		// indicates an unexpected transport layer protocol
+		if msg == nil {
+			log.Debug("Unexpected transport layer protocol")
 			continue PACKETLOOP
 		}
 
+		// Parse network layer information
+		networkLayer := packet.NetworkLayer()
+		switch networkLayer.LayerType() {
+		case layers.LayerTypeIPv4:
+			ip4 = networkLayer.(*layers.IPv4)
+			schema.SourceAddress = ip4.SrcIP.String()
+			schema.DestinationAddress = ip4.DstIP.String()
+			schema.Ipv4 = true
+			stats.PacketIPv4 += 1
+		case layers.LayerTypeIPv6:
+			ip6 = networkLayer.(*layers.IPv6)
+			schema.SourceAddress = ip6.SrcIP.String()
+			schema.DestinationAddress = ip6.DstIP.String()
+			schema.Ipv4 = false
+			stats.PacketIPv6 += 1
+		}
+
 		// Ignore questions unless flag set
 		if !msg.Response && !DoParseQuestions && !DoParseQuestionsEcs {
 			continue PACKETLOOP
@@ -208,7 +222,7 @@ PACKETLOOP:
 			schema.Qtype = qr.Qtype
 		}
 
-		// Let's get QUESTION information on if:
+		// Let's get QUESTION information if:
 		//   1. Questions flag is set
 		//   2. QuestionsEcs flag is set and ECS information in question
 		//   3. NXDOMAINs without RRs (i.e., SOA)