From 082cfa105aa1a2d5aee980a0ee8aa661681ec708 Mon Sep 17 00:00:00 2001
From: xSmurf <xsmurf+github@smurfturf.net>
Date: Sat, 20 Jun 2015 02:58:17 +0000
Subject: [PATCH] Added `all` option to kill command

---
 oz-daemon/client.go |  8 ++++++--
 oz-daemon/daemon.go | 20 ++++++++++++++------
 oz-daemon/launch.go |  8 +++++++-
 oz-init/init.go     |  2 +-
 oz/main.go          | 12 +++++++++---
 5 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/oz-daemon/client.go b/oz-daemon/client.go
index 4b75276..cd953a9 100644
--- a/oz-daemon/client.go
+++ b/oz-daemon/client.go
@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"os"
 	"strconv"
-	
+
 	"github.com/subgraph/oz/ipc"
 )
 
@@ -59,7 +59,7 @@ func Launch(arg, cpath string, args, env []string, noexec bool) error {
 		return err
 	}
 	pwd, _ := os.Getwd()
-	
+
 	resp, err := clientSend(&LaunchMsg{
 		Index:  idx,
 		Name:   name,
@@ -106,6 +106,10 @@ func Clean(arg string) error {
 	}
 }
 
+func KillAllSanodboxes() error {
+	return KillSandbox(-1)
+}
+
 func KillSandbox(id int) error {
 	resp, err := clientSend(&KillSandboxMsg{Id: id})
 	if err != nil {
diff --git a/oz-daemon/daemon.go b/oz-daemon/daemon.go
index ea290b0..6123be2 100644
--- a/oz-daemon/daemon.go
+++ b/oz-daemon/daemon.go
@@ -189,12 +189,20 @@ func (d *daemonState) sanitizeEnvironment(p *oz.Profile, oldEnv []string) ([]str
 }
 
 func (d *daemonState) handleKillSandbox(msg *KillSandboxMsg, m *ipc.Message) error {
-	sbox := d.sandboxById(msg.Id)
-	if sbox == nil {
-		return m.Respond(&ErrorMsg{fmt.Sprintf("no sandbox found with id = %d", msg.Id)})
-	}
-	if err := sbox.init.Process.Signal(os.Interrupt); err != nil {
-		return m.Respond(&ErrorMsg{fmt.Sprintf("failed to send interrupt signal: %v", err)})
+	if msg.Id == -1 {
+		for _, sb := range d.sandboxes {
+			if err := sb.init.Process.Signal(os.Interrupt); err != nil {
+				return m.Respond(&ErrorMsg{fmt.Sprintf("failed to send interrupt signal: %v", err)})
+			}
+		}
+	} else {
+		sbox := d.sandboxById(msg.Id)
+		if sbox == nil {
+			return m.Respond(&ErrorMsg{fmt.Sprintf("no sandbox found with id = %d", msg.Id)})
+		}
+		if err := sbox.init.Process.Signal(os.Interrupt); err != nil {
+			return m.Respond(&ErrorMsg{fmt.Sprintf("failed to send interrupt signal: %v", err)})
+		}
 	}
 	return m.Respond(&OkMsg{})
 }
diff --git a/oz-daemon/launch.go b/oz-daemon/launch.go
index d181b94..3538cde 100644
--- a/oz-daemon/launch.go
+++ b/oz-daemon/launch.go
@@ -120,12 +120,15 @@ func (d *daemonState) launch(p *oz.Profile, msg *LaunchMsg, uid, gid uint32, log
 		init:    cmd,
 		cred:    &syscall.Credential{Uid: uid, Gid: gid},
 		fs:      fs,
-		addr:    path.Join(fs.Root(), "tmp", "oz-init-control"),
+		addr:    path.Join(fs.Root(), ozinit.SocketAddress),
 		stderr:  pp,
 		network: stn,
 	}
 
+	wgNet := new(sync.WaitGroup)
 	if p.Networking.Nettype == network.TYPE_BRIDGE {
+		defer wgNet.Done()
+		wgNet.Add(1)
 		if err := network.NetInit(stn, d.network, cmd.Process.Pid, log); err != nil {
 			cmd.Process.Kill()
 			fs.Cleanup()
@@ -138,6 +141,8 @@ func (d *daemonState) launch(p *oz.Profile, msg *LaunchMsg, uid, gid uint32, log
 
 	if p.Networking.Nettype != network.TYPE_HOST && len(p.Networking.Sockets) > 0 {
 		go func() {
+			defer wgNet.Done()
+			wgNet.Add(1)
 			sbox.ready.Wait()
 			err := network.ProxySetup(sbox.init.Process.Pid, p.Networking.Sockets, d.log, sbox.ready)
 			if err != nil {
@@ -149,6 +154,7 @@ func (d *daemonState) launch(p *oz.Profile, msg *LaunchMsg, uid, gid uint32, log
 	if !msg.Noexec {
 		go func () {
 			sbox.ready.Wait()
+			wgNet.Wait()
 			go sbox.launchProgram(msg.Path, msg.Pwd, msg.Args, log)
 		}()
 	}
diff --git a/oz-init/init.go b/oz-init/init.go
index 45258da..da9d0c5 100644
--- a/oz-init/init.go
+++ b/oz-init/init.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"os/exec"
 	"os/user"
+	"os/signal"
 	"strconv"
 	"strings"
 	"sync"
@@ -21,7 +22,6 @@ import (
 
 	"github.com/kr/pty"
 	"github.com/op/go-logging"
-	"os/signal"
 )
 
 const SocketAddress = "/tmp/oz-init-control"
diff --git a/oz/main.go b/oz/main.go
index 88f0c5c..5229776 100644
--- a/oz/main.go
+++ b/oz/main.go
@@ -214,16 +214,22 @@ func handleClean(c *cli.Context) {
 
 func handleKill(c *cli.Context) {
 	if len(c.Args()) == 0 {
-		fmt.Println("Need a sandbox id to kill\n")
+		fmt.Errorf("Need a sandbox id to kill\n")
 		os.Exit(1)
 	}
+	if c.Args()[0] == "all" {
+		if err := daemon.KillAllSanodboxes(); err != nil {
+			fmt.Errorf("Kill command failed:", err)
+		}
+		return
+	}
 	id, err := strconv.Atoi(c.Args()[0])
 	if err != nil {
-		fmt.Printf("Could not parse id value %s\n", c.Args()[0])
+		fmt.Errorf("Could not parse id value %s\n", c.Args()[0])
 		os.Exit(1)
 	}
 	if err := daemon.KillSandbox(id); err != nil {
-		fmt.Println("Kill command failed:", err)
+		fmt.Errorf("Kill command failed:", err)
 	}
 
 }