From a3aeca204c7fa5b0574aad81ac6d7dc0e1bd4338 Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Thu, 9 Jun 2016 18:41:01 +0600
Subject: [PATCH] Check dst and src IP's for packets (#295)

* Add IP check

* Check dst and src ips
---
 Makefile                        |  7 ++++-
 input_file.go                   |  1 +
 raw_socket_listener/listener.go | 45 ++++++++++++++++++++++++++++-----
 3 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/Makefile b/Makefile
index 427945a..f91108e 100644
--- a/Makefile
+++ b/Makefile
@@ -1,6 +1,8 @@
 SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go test_input.go elasticsearch.go http_modifier.go http_modifier_settings.go http_client.go middleware.go protocol.go output_file_settings.go
 SOURCE_PATH = /go/src/github.com/buger/gor/
-RUN = docker run -v `pwd`:$(SOURCE_PATH) -p 0.0.0.0:8000:8000 -t -i gor
+PORT = 8000
+FADDR = :8000
+RUN = docker run -v `pwd`:$(SOURCE_PATH) -p 0.0.0.0:$(PORT):$(PORT) -t -i gor
 BENCHMARK = BenchmarkRAWInput
 TEST = TestRawListenerBench
 VERSION = DEV-$(shell date +%s)
@@ -68,6 +70,9 @@ run-2:
 run-3:
 	sudo -E go run $(SOURCE) --input-tcp :27001 --output-stdout
 
+run-arg:
+	sudo -E go run $(SOURCE) $(ARGS)
+
 file-server:
 	go run $(SOURCE) file-server $(FADDR)
 
diff --git a/input_file.go b/input_file.go
index 7f59a49..b562a99 100644
--- a/input_file.go
+++ b/input_file.go
@@ -137,6 +137,7 @@ func (i *FileInput) emit() {
 				if e := i.updateFile(); e != nil {
 					if _, ok := e.(*NextFileNotFound); ok && i.loop {
 						// Start from the first file
+						i.Close()
 						i.currentFile = nil
 						i.currentReader = nil
 						lastTime = 0
diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go
index 141285b..6351cf8 100644
--- a/raw_socket_listener/listener.go
+++ b/raw_socket_listener/listener.go
@@ -298,16 +298,27 @@ func (t *Listener) readPcap() {
 			t.mu.Lock()
 			t.pcapHandles = append(t.pcapHandles, handle)
 
+			var bpfDstHost, bpfSrcHost string
+			for i, addr := range device.Addresses {
+				bpfDstHost += "dst host " + addr.IP.String()
+				bpfSrcHost += "src host " + addr.IP.String()
+				if i != len(device.Addresses) - 1 {
+					bpfDstHost += " or "
+					bpfSrcHost += " or "
+				}
+			}
+
 			if bpfSupported {
 				var bpf string
 
 				if t.trackResponse {
-					bpf = "tcp port " + strconv.Itoa(int(t.port))
+					bpf = "(tcp dst port " + strconv.Itoa(int(t.port)) + " and (" + bpfDstHost + ")) or (" + "tcp src port " + strconv.Itoa(int(t.port)) + " and (" + bpfSrcHost + "))"
 				} else {
-					bpf = "tcp dst port " + strconv.Itoa(int(t.port))
+					bpf = "tcp dst port " + strconv.Itoa(int(t.port)) + " and (" + bpfDstHost + ")"
 				}
+
 				if err := handle.SetBPFFilter(bpf); err != nil {
-					log.Println("BPF filter error:", err, "Device:", device.Name)
+					log.Println("BPF filter error:", err, "Device:", device.Name, bpf)
 					wg.Done()
 					return
 				}
@@ -321,7 +332,7 @@ func (t *Listener) readPcap() {
 
 			wg.Done()
 
-			var data, srcIP []byte
+			var data, srcIP, dstIP []byte
 
 			for {
 				packet, err := source.NextPacket()
@@ -350,6 +361,7 @@ func (t *Listener) readPcap() {
 					}
 
 					srcIP = data[12:16]
+					dstIP = data[16:20]
 					data = data[ihl*4:]
 				} else {
 					// Truncated IP info
@@ -358,6 +370,7 @@ func (t *Listener) readPcap() {
 					}
 
 					srcIP = data[8:24]
+					dstIP = data[24:40]
 
 					data = data[40:]
 				}
@@ -376,9 +389,29 @@ func (t *Listener) readPcap() {
 						destPort := binary.BigEndian.Uint16(data[2:4])
 						srcPort := binary.BigEndian.Uint16(data[0:2])
 
-						// log.Println(t.port, destPort, srcPort, packet)
+						var addrCheck []byte
 
-						if !(destPort == t.port || (t.trackResponse && srcPort == t.port)) {
+						if destPort == t.port {
+							addrCheck = dstIP
+						}
+
+						if t.trackResponse && srcPort == t.port {
+							addrCheck = srcIP
+						}
+
+						if len(addrCheck) == 0 {
+							continue
+						}
+
+						addrMatched := false
+						for _, a := range device.Addresses {
+							if a.IP.Equal(net.IP(addrCheck)) {
+								addrMatched = true
+								break
+							}
+						}
+
+						if !addrMatched {
 							continue
 						}
 					}