From f9b0ae6d551f7c6c8be7c1099a829130f471d8b0 Mon Sep 17 00:00:00 2001
From: Urban Ishimwe <urbainishimwe@gmail.com>
Date: Sun, 30 Aug 2020 12:07:00 +0200
Subject: [PATCH] fix #809 and minor bug in pool

---
 capture/sock.go        | 15 ++-------------
 capture/sock_others.go | 13 +++++++++++++
 capture/sock_raw.go    | 15 +++++++++++++++
 gor.go                 |  1 -
 tcp/tcp_message.go     | 13 +++++++------
 5 files changed, 37 insertions(+), 20 deletions(-)
 create mode 100644 capture/sock_others.go
 create mode 100644 capture/sock_raw.go

diff --git a/capture/sock.go b/capture/sock.go
index fed8cd0..781317e 100644
--- a/capture/sock.go
+++ b/capture/sock.go
@@ -1,9 +1,10 @@
+// +build linux
+
 package capture
 
 import (
 	"fmt"
 	"net"
-	"sync"
 	"time"
 	"unsafe"
 
@@ -31,18 +32,6 @@ const (
 
 var tpacket2hdrlen = tpAlign(int(unsafe.Sizeof(unix.Tpacket2Hdr{})))
 
-// SockRaw is a linux M'maped af_packet socket
-type SockRaw struct {
-	mu          sync.Mutex
-	fd          int
-	ifindex     int
-	snaplen     int
-	pollTimeout uintptr
-	frame       uint32 // current frame
-	buf         []byte // points to the memory space of the ring buffer shared with the kernel.
-	loopIndex   int32  // this field must filled to avoid reading packet twice on a loopback device
-}
-
 // NewSockRaw returns new M'maped sock_raw on packet version 2.
 func NewSockRaw(ifi net.Interface) (*SockRaw, error) {
 	// sock create
diff --git a/capture/sock_others.go b/capture/sock_others.go
new file mode 100644
index 0000000..fb4b519
--- /dev/null
+++ b/capture/sock_others.go
@@ -0,0 +1,13 @@
+// +build !linux
+
+package capture
+
+import (
+	"errors"
+	"net"
+)
+
+// NewSockRaw returns new M'maped sock_raw on packet version 2.
+func NewSockRaw(_ net.Interface) (*SockRaw, error) {
+	return nil, errors.New("afpacket socket is only available on linux")
+}
diff --git a/capture/sock_raw.go b/capture/sock_raw.go
new file mode 100644
index 0000000..8e28e60
--- /dev/null
+++ b/capture/sock_raw.go
@@ -0,0 +1,15 @@
+package capture
+
+import "sync"
+
+// SockRaw is a linux M'maped af_packet socket
+type SockRaw struct {
+	mu          sync.Mutex
+	fd          int
+	ifindex     int
+	snaplen     int
+	pollTimeout uintptr
+	frame       uint32 // current frame
+	buf         []byte // points to the memory space of the ring buffer shared with the kernel.
+	loopIndex   int32  // this field must filled to avoid reading packet twice on a loopback device
+}
diff --git a/gor.go b/gor.go
index f4f29a1..6272b06 100644
--- a/gor.go
+++ b/gor.go
@@ -19,7 +19,6 @@ import (
 )
 
 var (
-	mode       string
 	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
 	memprofile = flag.String("memprofile", "", "write memory profile to this file")
 )
diff --git a/tcp/tcp_message.go b/tcp/tcp_message.go
index 0e05bda..c83b765 100644
--- a/tcp/tcp_message.go
+++ b/tcp/tcp_message.go
@@ -153,12 +153,14 @@ func (pool *MessagePool) Handler(packet gopacket.Packet) {
 	}
 	if pckt.RST {
 		if ok {
+			m.done <- true
 			<-m.done
 		}
 		if m, ok = pool.pool[pckt.Dst()]; !ok {
 			m, ok = pool.pool[pckt.Dst()+"="+srcKey]
 		}
 		if ok {
+			m.done <- true
 			<-m.done
 		}
 		go pool.say(4, fmt.Sprintf("RST flag from %s to %s at %s\n", pckt.Src(), pckt.Dst(), pckt.Timestamp))
@@ -168,13 +170,12 @@ func (pool *MessagePool) Handler(packet gopacket.Packet) {
 	case ok:
 		pool.addPacket(m, pckt)
 		return
-	case pool.Start != nil:
-		if in, out = pool.Start(pckt); in || out {
-			break
-		}
-		return
 	case pckt.SYN:
 		in = !pckt.ACK
+	case pool.Start != nil:
+		if in, out = pool.Start(pckt); !(in || out) {
+			return
+		}
 	default:
 		return
 	}
@@ -212,8 +213,8 @@ func (pool *MessagePool) addPacket(m *Message, pckt *Packet) {
 	m.add(pckt)
 	switch {
 	case trunc >= 0:
-	case pool.End != nil && pool.End(m):
 	case pckt.FIN:
+	case pool.End != nil && pool.End(m):
 	default:
 		return
 	}