Merge pull request #810 from buger/afpacket_linux_only

fix #809 and minor bug in pool

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

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")
+}

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
+}

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")
 )

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
 	}

tcp/tcp_test.go

@@ -76,8 +76,8 @@ func TestMessageParserWithHint(t *testing.T) {
 		return
 	case m = <-mssg:
 	}
-	if len(m.packets) != 3 {
-		t.Errorf("expected to have 3 packets got %d", len(m.packets))
+	if len(m.packets) != 7 {
+		t.Errorf("expected to have 7 packets got %d", len(m.packets))
 	}
 	if !bytes.HasSuffix(m.Data(), []byte("\n7\r\nNetwork\r\n0\r\n\r\n")) {
 		t.Errorf("expected to %q to have suffix %q", m.Data(), []byte("\n7\r\nNetwork\r\n0\r\n\r\n"))
@@ -89,8 +89,8 @@ func TestMessageParserWithHint(t *testing.T) {
 		return
 	case m = <-mssg:
 	}
-	if len(m.packets) != 3 {
-		t.Errorf("expected to have 3 packets got %d", len(m.packets))
+	if len(m.packets) != 7 {
+		t.Errorf("expected to have 7 packets got %d", len(m.packets))
 	}
 	if !bytes.HasSuffix(m.Data(), []byte("Network")) {
 		t.Errorf("expected to %q to have suffix %q", m.Data(), []byte("Network"))