From b567e5ce54d09c114a2737f3bf650d606a1075ef Mon Sep 17 00:00:00 2001
From: shw <shw@subgraph.com>
Date: Wed, 10 May 2017 14:51:16 +0000
Subject: [PATCH] oz-init pid mappings are now properly destroyed as these
 processes exit.

---
 sgfw/ipc.go    | 15 ++++++++++++++-
 sgfw/policy.go | 11 +++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/sgfw/ipc.go b/sgfw/ipc.go
index efd29f5..4823b38 100644
--- a/sgfw/ipc.go
+++ b/sgfw/ipc.go
@@ -22,7 +22,7 @@ var OzInitPids []OzInitProc = []OzInitProc{}
 
 
 func addInitPid(pid int, name string) {
-	fmt.Println("::::::::::: init pid added: ", pid, " -> ", name)
+fmt.Println("::::::::::: init pid added: ", pid, " -> ", name)
 	for i := 0; i < len(OzInitPids); i++ {
 		if OzInitPids[i].Pid == pid {
 			return
@@ -34,6 +34,7 @@ func addInitPid(pid int, name string) {
 }
 
 func removeInitPid(pid int) {
+fmt.Println("::::::::::: removing PID: ", pid)
 	for i := 0; i < len(OzInitPids); i++ {
 		if OzInitPids[i].Pid == pid {
 			OzInitPids = append(OzInitPids[:i], OzInitPids[i+1:]...)
@@ -174,6 +175,18 @@ func ReceiverLoop(fw *Firewall, c net.Conn) {
 				addInitPid(initpid, ozname)
 				c.Write([]byte("OK.\n"))
 				return
+			} else if tokens[0] == "unregister-init" && len(tokens) == 2 {
+				initp := tokens[1]
+				initpid, err := strconv.Atoi(initp)
+
+				if err != nil {
+					log.Notice("IPC received invalid oz-init pid: ", initp)
+					c.Write([]byte("Bad command: init pid was invalid"))
+					return
+				}
+
+				removeInitPid(initpid)
+				c.Write([]byte("OK.\n"))
 			}
 
 			if len(tokens) != 6 {
diff --git a/sgfw/policy.go b/sgfw/policy.go
index 20029bc..04de998 100644
--- a/sgfw/policy.go
+++ b/sgfw/policy.go
@@ -502,6 +502,7 @@ func findProcessForPacket(pkt *nfqueue.NFQPacket) (*procsnitch.Info, string) {
 		res := procsnitch.LookupTCPSocketProcessAll(srcip, srcp, dstip, dstp, nil)
 
 		if res == nil {
+			removePids := make([]int, 0)
 
 			for i := 0; i < len(OzInitPids); i++ {
 				data := ""
@@ -511,6 +512,11 @@ fmt.Println("XXX: opening: ", fname)
 
 				if err != nil {
 					fmt.Println("Error reading proc data from ", fname, ": ", err)
+
+					if err == syscall.ENOENT {
+						removePids = append(removePids, OzInitPids[i].Pid)
+					}
+
 					continue
 				} else {
 					data = string(bdata)
@@ -538,6 +544,11 @@ fmt.Println("XXX: opening: ", fname)
 				}
 
 			}
+
+			for _, p := range removePids {
+				removeInitPid(p)
+			}
+
 		}
 
 		return res, optstr