Properly handle operation as init process

pkg/cli/agent/agent.go

@@ -23,9 +23,16 @@ func Run(ctx *cli.Context) error {
 	// database credentials or other secrets.
 	gspt.SetProcTitle(os.Args[0] + " agent")
 
+	// Do init stuff if pid 1.
+	// This must be done before InitLogging as that may reexec in order to capture log output
+	if err := cmds.HandleInit(); err != nil {
+		return err
+	}
+
 	if err := cmds.InitLogging(); err != nil {
 		return err
 	}
+
 	if os.Getuid() != 0 && runtime.GOOS != "windows" {
 		return fmt.Errorf("agent must be ran as root")
 	}

pkg/cli/cmds/init_default.go

@@ -0,0 +1,7 @@
+// +build !linux !cgo
+
+package cmds
+
+func HandleInit() error {
+	return nil
+}

pkg/cli/cmds/init_linux.go

@@ -0,0 +1,83 @@
+// +build linux,cgo
+
+package cmds
+
+import (
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/erikdubbelboer/gspt"
+	"github.com/pkg/errors"
+	"github.com/rancher/k3s/pkg/version"
+	"github.com/rootless-containers/rootlesskit/pkg/parent/cgrouputil"
+)
+
+// HandleInit takes care of things that need to be done when running as process 1, usually in a
+// Docker container. This includes evacuating the root cgroup and reaping child pids.
+func HandleInit() error {
+	if os.Getpid() != 1 {
+		return nil
+	}
+
+	// The root cgroup has to be empty to enable subtree_control, so evacuate it by placing
+	// ourselves in the init cgroup.
+	if err := cgrouputil.EvacuateCgroup2("init"); err != nil {
+		return errors.Wrap(err, "failed to evacuate root cgroup")
+	}
+
+	pwd, err := os.Getwd()
+	if err != nil {
+		return errors.Wrap(err, "failed to get working directory for init process")
+	}
+
+	go reapChildren()
+
+	// fork the main process to do work so that this init process can handle reaping pids
+	// without interfering with any other exec's that the rest of the codebase may do.
+	var wstatus syscall.WaitStatus
+	pattrs := &syscall.ProcAttr{
+		Dir: pwd,
+		Env: os.Environ(),
+		Sys: &syscall.SysProcAttr{Setsid: true},
+		Files: []uintptr{
+			uintptr(syscall.Stdin),
+			uintptr(syscall.Stdout),
+			uintptr(syscall.Stderr),
+		},
+	}
+	pid, err := syscall.ForkExec(os.Args[0], os.Args, pattrs)
+	if err != nil {
+		return errors.Wrap(err, "failed to fork/exec "+version.Program)
+	}
+
+	gspt.SetProcTitle(os.Args[0] + " init")
+	// wait for main process to exit, and return its status when it does
+	_, err = syscall.Wait4(pid, &wstatus, 0, nil)
+	for err == syscall.EINTR {
+		_, err = syscall.Wait4(pid, &wstatus, 0, nil)
+	}
+	os.Exit(wstatus.ExitStatus())
+	return nil
+}
+
+//reapChildren calls Wait4 whenever SIGCHLD is received
+func reapChildren() {
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGCHLD)
+	for {
+		select {
+		case <-sigs:
+		}
+		for {
+			var wstatus syscall.WaitStatus
+			_, err := syscall.Wait4(-1, &wstatus, 0, nil)
+			for err == syscall.EINTR {
+				_, err = syscall.Wait4(-1, &wstatus, 0, nil)
+			}
+			if err == nil || err == syscall.ECHILD {
+				break
+			}
+		}
+	}
+}

pkg/cli/server/server.go

@@ -38,16 +38,10 @@ import (
 )
 
 func Run(app *cli.Context) error {
-	if err := cmds.InitLogging(); err != nil {
-		return err
-	}
 	return run(app, &cmds.ServerConfig, server.CustomControllers{}, server.CustomControllers{})
 }
 
 func RunWithControllers(app *cli.Context, leaderControllers server.CustomControllers, controllers server.CustomControllers) error {
-	if err := cmds.InitLogging(); err != nil {
-		return err
-	}
 	return run(app, &cmds.ServerConfig, leaderControllers, controllers)
 }
 
@@ -60,6 +54,16 @@ func run(app *cli.Context, cfg *cmds.Server, leaderControllers server.CustomCont
 	// database credentials or other secrets.
 	gspt.SetProcTitle(os.Args[0] + " server")
 
+	// Do init stuff if pid 1.
+	// This must be done before InitLogging as that may reexec in order to capture log output
+	if err := cmds.HandleInit(); err != nil {
+		return err
+	}
+
+	if err := cmds.InitLogging(); err != nil {
+		return err
+	}
+
 	if !cfg.DisableAgent && os.Getuid() != 0 && !cfg.Rootless {
 		return fmt.Errorf("must run as root unless --disable-agent is specified")
 	}