Merge branch 'master' of ssh://github.com/subgraph/oz

ipc/ipc.go

@@ -31,6 +31,7 @@ type MsgConn struct {
 }
 
 type MsgServer struct {
+	isClosed bool
 	log      *logging.Logger
 	disp     *msgDispatcher
 	factory  MsgFactory
@@ -66,9 +67,12 @@ func NewServer(address string, factory MsgFactory, log *logging.Logger, handlers
 }
 
 func (s *MsgServer) Run() error {
-	for {
+	for !s.isClosed {
 		conn, err := s.listener.AcceptUnix()
 		if err != nil {
+			if s.isClosed {
+				return nil
+			}
 			return err
 		}
 		if err := setPassCred(conn); err != nil {
@@ -90,6 +94,10 @@ func (s *MsgServer) Run() error {
 }
 
 func (s *MsgServer) Close() error {
+	if s.isClosed {
+		return nil
+	}
+	s.isClosed = true
 	s.disp.close()
 	close(s.done)
 	return s.listener.Close()

oz-daemon/launch.go

@@ -36,7 +36,6 @@ type Sandbox struct {
 	network *network.SandboxNetwork
 }
 
-
 func createInitCommand(name, chroot string, env []string, uid uint32, display int, stn *network.SandboxNetwork, nettype string) *exec.Cmd {
 	cmd := exec.Command(initPath)
 	cmd.Dir = "/"

oz-init/init.go

@@ -32,11 +32,15 @@ type initState struct {
 	profile   *oz.Profile
 	config    *oz.Config
 	launchEnv []string
+	lock      sync.Mutex
+	children  map[int]*exec.Cmd
 	uid       int
 	gid       int
 	user      *user.User
 	display   int
 	fs        *fs.Filesystem
+	ipcServer *ipc.MsgServer
+	xpra      *xpra.Xpra
 	xpraReady sync.WaitGroup
 	network   *network.SandboxNetwork
 }
@@ -146,6 +150,7 @@ func parseArgs() *initState {
 		config:    config,
 		launchEnv: env,
 		profile:   p,
+		children:  make(map[int]*exec.Cmd),
 		uid:       uid,
 		gid:       gid,
 		user:      u,
@@ -202,6 +207,8 @@ func (st *initState) runInit() {
 
 	go st.processSignals(sigs, s)
 
+	st.ipcServer = s
+
 	if err := s.Run(); err != nil {
 		st.log.Warning("MsgServer.Run() return err: %v", err)
 	}
@@ -233,6 +240,7 @@ func (st *initState) startXpraServer() {
 	if err := xpra.Process.Start(); err != nil {
 		st.log.Warning("Failed to start xpra server: %v", err)
 	}
+	st.xpra = xpra
 }
 
 func (st *initState) readXpraOutput(r io.ReadCloser) {
@@ -280,6 +288,8 @@ func (st *initState) launchApplication() {
 		st.log.Warning("Failed to start application (%s): %v", st.profile.Path, err)
 		return
 	}
+	st.addChildProcess(cmd)
+
 	go st.readApplicationOutput(stdout, "stdout")
 	go st.readApplicationOutput(stderr, "stderr")
 }
@@ -345,6 +355,7 @@ func (st *initState) handleRunShell(rs *RunShellMsg, msg *ipc.Message) error {
 	if err != nil {
 		return msg.Respond(&ErrorMsg{err.Error()})
 	}
+	st.addChildProcess(cmd)
 	err = msg.Respond(&OkMsg{}, int(f.Fd()))
 	return err
 }
@@ -370,14 +381,70 @@ func ptyStart(c *exec.Cmd) (ptty *os.File, err error) {
 	return ptty, nil
 }
 
-func (is *initState) handleChildExit(pid int, wstatus syscall.WaitStatus) {
-	is.log.Debug("Child process pid=%d exited with status %d", pid, wstatus.ExitStatus())
+func (st *initState) addChildProcess(cmd *exec.Cmd) {
+	st.lock.Lock()
+	defer st.lock.Unlock()
+	st.children[cmd.Process.Pid] = cmd
+}
+
+func (st *initState) removeChildProcess(pid int) bool {
+	st.lock.Lock()
+	defer st.lock.Unlock()
+	if _, ok := st.children[pid]; ok {
+		delete(st.children, pid)
+		return true
+	}
+	return false
+}
+
+func (st *initState) handleChildExit(pid int, wstatus syscall.WaitStatus) {
+	st.log.Debug("Child process pid=%d exited with status %d", pid, wstatus.ExitStatus())
+	st.removeChildProcess(pid)
 }
 
 func (st *initState) processSignals(c <-chan os.Signal, s *ipc.MsgServer) {
 	for {
 		sig := <-c
 		st.log.Info("Recieved signal (%v)", sig)
-		s.Close()
+		st.shutdown()
+	}
+}
+
+func (st *initState) shutdown() {
+	for _, c := range st.childrenVector() {
+		c.Process.Signal(os.Interrupt)
+	}
+
+	st.shutdownXpra()
+
+	if st.ipcServer != nil {
+		st.ipcServer.Close()
+	}
+}
+
+func (st *initState) shutdownXpra() {
+	if st.xpra == nil {
+		return
+	}
+	out, err := st.xpra.Stop()
+	if err != nil {
+		st.log.Warning("Error running xpra stop: %v", err)
+		return
+	}
+
+	for _, line := range strings.Split(string(out), "\n") {
+		if len(line) > 0 {
+			st.log.Debug("(xpra stop) %s", line)
+		}
+	}
+}
+
+func (st *initState) childrenVector() []*exec.Cmd {
+	st.lock.Lock()
+	defer st.lock.Unlock()
+	cs := make([]*exec.Cmd, 0, len(st.children))
+	for _, v := range st.children {
+		cs = append(cs, v)
 	}
+	return cs
 }

xpra/xpra.go

@@ -1,6 +1,7 @@
 package xpra
 
 import (
+	"fmt"
 	"github.com/subgraph/oz"
 	"os/exec"
 )
@@ -52,3 +53,13 @@ func getDefaultArgs(config *oz.XServerConf) []string {
 
 	return args
 }
+
+func (x *Xpra) Stop() ([]byte, error) {
+	cmd := exec.Command("/usr/bin/xpra",
+		"--socket-dir="+x.WorkDir,
+		"stop",
+		fmt.Sprintf(":%d", x.Display),
+	)
+	cmd.Env = []string{"TMPDIR=" + x.WorkDir}
+	return cmd.Output()
+}