From 571487b1355b3f1211f178b35d3970612bc5495c Mon Sep 17 00:00:00 2001
From: Urban Ishimwe <urbainishimwe@gmail.com>
Date: Wed, 26 Aug 2020 17:46:00 +0200
Subject: [PATCH] dispatch on rst flag

---
 capture/sock.go    |  8 ++++++--
 tcp/tcp_message.go | 17 ++++++++++++++---
 tcp/tcp_packet.go  | 15 +++++++++++++++
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/capture/sock.go b/capture/sock.go
index 796c9d0..fed8cd0 100644
--- a/capture/sock.go
+++ b/capture/sock.go
@@ -120,15 +120,19 @@ func (sock *SockRaw) ReadPacketData() (buf []byte, ci gopacket.CaptureInfo, err
 read:
 	i = int(sock.frame * FRAMESIZE)
 	tpHdr = (*unix.Tpacket2Hdr)(unsafe.Pointer(&sock.buf[i]))
+	sock.frame = (sock.frame + 1) % FRAMENR
 
-	for tpHdr.Status&unix.TP_STATUS_USER == 0 {
+	if tpHdr.Status&unix.TP_STATUS_USER == 0 {
 		_, _, e := unix.Syscall(unix.SYS_POLL, uintptr(unsafe.Pointer(poll)), 1, sock.pollTimeout)
 		if e != 0 && e != unix.EINTR {
 			return buf, ci, e
 		}
+		// it might be some other frame with data!
+		if tpHdr.Status&unix.TP_STATUS_USER == 0 {
+			goto read
+		}
 	}
 
-	sock.frame = (sock.frame + 1) % FRAMENR
 	tpHdr.Status = unix.TP_STATUS_KERNEL
 	sockAddr := (*unix.RawSockaddrLinklayer)(unsafe.Pointer(&sock.buf[i+tpacket2hdrlen]))
 
diff --git a/tcp/tcp_message.go b/tcp/tcp_message.go
index e705c32..0e05bda 100644
--- a/tcp/tcp_message.go
+++ b/tcp/tcp_message.go
@@ -139,7 +139,7 @@ func NewMessagePool(maxSize size.Size, messageExpire time.Duration, debugger Deb
 func (pool *MessagePool) Handler(packet gopacket.Packet) {
 	var in, out bool
 	pckt, err := ParsePacket(packet)
-	if err != nil {
+	if err != nil || pckt == nil {
 		go pool.say(4, fmt.Sprintf("error decoding packet(%dBytes):%s\n", packet.Metadata().CaptureLength, err))
 		return
 	}
@@ -151,6 +151,19 @@ func (pool *MessagePool) Handler(packet gopacket.Packet) {
 	if !ok {
 		m, ok = pool.pool[dstKey]
 	}
+	if pckt.RST {
+		if ok {
+			<-m.done
+		}
+		if m, ok = pool.pool[pckt.Dst()]; !ok {
+			m, ok = pool.pool[pckt.Dst()+"="+srcKey]
+		}
+		if ok {
+			<-m.done
+		}
+		go pool.say(4, fmt.Sprintf("RST flag from %s to %s at %s\n", pckt.Src(), pckt.Dst(), pckt.Timestamp))
+		return
+	}
 	switch {
 	case ok:
 		pool.addPacket(m, pckt)
@@ -201,8 +214,6 @@ func (pool *MessagePool) addPacket(m *Message, pckt *Packet) {
 	case trunc >= 0:
 	case pool.End != nil && pool.End(m):
 	case pckt.FIN:
-	case pckt.RST:
-		go pool.say(4, fmt.Sprintf("RST flag from %s to %s at %s\n", pckt.Src(), pckt.Dst(), pckt.Timestamp))
 	default:
 		return
 	}
diff --git a/tcp/tcp_packet.go b/tcp/tcp_packet.go
index c2aeea4..a1a4968 100644
--- a/tcp/tcp_packet.go
+++ b/tcp/tcp_packet.go
@@ -146,6 +146,19 @@ func (pckt *Packet) SYNOptions() (mss uint16, windowscale byte) {
 	return
 }
 
+// LinkInfo returns info about the link layer
+func (pckt *Packet) LinkInfo() string {
+	if l, ok := pckt.LinkLayer.(*layers.Ethernet); ok {
+		return fmt.Sprintf(
+			"Source Mac: %s\nDestination Mac: %s\nProtocol: %s",
+			l.SrcMAC,
+			l.DstMAC,
+			l.EthernetType,
+		)
+	}
+	return "<Not Ethernet>"
+}
+
 // Flag returns formatted tcp flags
 func (pckt *Packet) Flag() (flag string) {
 	if pckt.FIN {
@@ -175,6 +188,7 @@ func (pckt *Packet) Flag() (flag string) {
 // String output for a TCP Packet
 func (pckt *Packet) String() string {
 	return fmt.Sprintf(`Time: %s
+%s
 Source: %s
 Destination: %s
 IHL: %d
@@ -188,6 +202,7 @@ Options: %s
 Data Size: %d
 Lost Data: %d`,
 		pckt.Timestamp.Format(time.StampNano),
+		pckt.LinkInfo(),
 		pckt.Src(),
 		pckt.Dst(),
 		pckt.IHL(),