shim: user switcher process management struct

This change moves all user switcher and shim management to the shim package and withholds output while shim is alive. This also eliminated all exit scenarios where revert is skipped.

Signed-off-by: Ophestra Umiker <cat@ophivana.moe>

6 files changed, 306 insertions, 159 deletions

diff --git a/internal/app/app.go b/internal/app/app.go
index de2c461f..264975c5 100644
--- a/internal/app/app.go
+++ b/internal/app/app.go
@@ -1,10 +1,10 @@
 package app
 
 import (
-	"os/exec"
 	"sync"
 
 	"git.ophivana.moe/security/fortify/internal"
+	"git.ophivana.moe/security/fortify/internal/shim"
 )
 
 type App interface {
@@ -26,10 +26,8 @@ type app struct {
 	id *ID
 	// operating system interface
 	os internal.System
-	// underlying user switcher process
-	cmd *exec.Cmd
-	// shim setup abort reason and completion
-	abort chan error
+	// shim process manager
+	shim *shim.Shim
 	// child process related information
 	seal *appSeal
 	// error returned waiting for process
@@ -50,8 +48,8 @@ func (a *app) String() string {
 	a.lock.RLock()
 	defer a.lock.RUnlock()
 
-	if a.cmd != nil {
-		return a.cmd.String()
+	if a.shim != nil {
+		return a.shim.String()
 	}
 
 	if a.seal != nil {
diff --git a/internal/app/start.go b/internal/app/start.go
index 49a4678a..ea216018 100644
--- a/internal/app/start.go
+++ b/internal/app/start.go
@@ -3,12 +3,10 @@ package app
 import (
 	"errors"
 	"fmt"
-	"os"
 	"os/exec"
 	"path"
 	"path/filepath"
 	"strings"
-	"time"
 
 	"git.ophivana.moe/security/fortify/helper"
 	"git.ophivana.moe/security/fortify/internal/fmsg"
@@ -17,7 +15,8 @@ import (
 	"git.ophivana.moe/security/fortify/internal/system"
 )
 
-// Start starts the fortified child
+// Start selects a user switcher and starts shim.
+// Note that Wait must be called regardless of error returned by Start.
 func (a *app) Start() error {
 	a.lock.Lock()
 	defer a.lock.Unlock()
@@ -41,12 +40,8 @@ func (a *app) Start() error {
 		}
 	}
 
-	if err := a.seal.sys.Commit(); err != nil {
-		return err
-	}
-
 	// select command builder
-	var commandBuilder func(shimEnv string) (args []string)
+	var commandBuilder shim.CommandBuilder
 	switch a.seal.launchOption {
 	case LaunchMethodSudo:
 		commandBuilder = a.commandBuilderSudo
@@ -56,60 +51,45 @@ func (a *app) Start() error {
 		panic("unreachable")
 	}
 
-	// configure child process
-	confSockPath := path.Join(a.seal.share, "shim")
-	a.cmd = exec.Command(a.seal.toolPath, commandBuilder(shim.EnvShim+"="+confSockPath)...)
-	a.cmd.Env = []string{}
-	a.cmd.Stdin, a.cmd.Stdout, a.cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
-	a.cmd.Dir = a.seal.RunDirPath
+	// construct shim manager
+	a.shim = shim.New(a.seal.toolPath, uint32(a.seal.sys.UID()), path.Join(a.seal.share, "shim"), a.seal.wl,
+		&shim.Payload{
+			Argv:  a.seal.command,
+			Exec:  shimExec,
+			Bwrap: a.seal.sys.bwrap,
+			WL:    a.seal.wl != nil,
 
-	a.abort = make(chan error)
-	procReady := make(chan struct{})
-	if err := shim.ServeConfig(confSockPath, a.abort, func() {
-		<-procReady
-		if err := a.cmd.Process.Signal(os.Interrupt); err != nil {
-			fmsg.Println("cannot kill shim on faulted setup:", err)
-		}
-		fmt.Print("\r")
-	}, a.seal.sys.UID(), &shim.Payload{
-		Argv:  a.seal.command,
-		Exec:  shimExec,
-		Bwrap: a.seal.sys.bwrap,
-		WL:    a.seal.wl != nil,
+			Verbose: fmsg.Verbose(),
+		},
+	)
 
-		Verbose: fmsg.Verbose(),
-	}, a.seal.wl); err != nil {
-		a.abort <- err
-		<-a.abort
-		return fmsg.WrapErrorSuffix(err,
-			"cannot serve shim setup:")
+	// startup will go ahead, commit system setup
+	if err := a.seal.sys.Commit(); err != nil {
+		return err
 	}
+	a.seal.sys.needRevert = true
 
-	// start shim
-	fmsg.VPrintln("starting shim as target user:", a.cmd)
-	if err := a.cmd.Start(); err != nil {
-		return fmsg.WrapErrorSuffix(err,
-			"cannot start process:")
-	}
-	startTime := time.Now().UTC()
-	close(procReady)
+	if startTime, err := a.shim.Start(commandBuilder); err != nil {
+		return err
+	} else {
+		// shim start and setup success, create process state
+		sd := state.State{
+			PID:        a.shim.Unwrap().Process.Pid,
+			Command:    a.seal.command,
+			Capability: a.seal.et,
+			Method:     method[a.seal.launchOption],
+			Argv:       a.shim.Unwrap().Args,
+			Time:       *startTime,
+		}
 
-	// create process state
-	sd := state.State{
-		PID:        a.cmd.Process.Pid,
-		Command:    a.seal.command,
-		Capability: a.seal.et,
-		Method:     method[a.seal.launchOption],
-		Argv:       a.cmd.Args,
-		Time:       startTime,
+		// register process state
+		var err0 = new(StateStoreError)
+		err0.Inner, err0.DoErr = a.seal.store.Do(func(b state.Backend) {
+			err0.InnerErr = b.Save(&sd)
+		})
+		a.seal.sys.saveState = true
+		return err0.equiv("cannot save process state:")
 	}
-
-	// register process state
-	var err = new(StateStoreError)
-	err.Inner, err.DoErr = a.seal.store.Do(func(b state.Backend) {
-		err.InnerErr = b.Save(&sd)
-	})
-	return err.equiv("cannot save process state:")
 }
 
 // StateStoreError is returned for a failed state save
@@ -173,21 +153,28 @@ func (a *app) Wait() (int, error) {
 
 	var r int
 
-	// wait for process and resolve exit code
-	if err := a.cmd.Wait(); err != nil {
-		var exitError *exec.ExitError
-		if !errors.As(err, &exitError) {
-			// should be unreachable
-			a.waitErr = err
-		}
-
-		// store non-zero return code
-		r = exitError.ExitCode()
+	if cmd := a.shim.Unwrap(); cmd == nil {
+		// failure prior to process start
+		r = 255
 	} else {
-		r = a.cmd.ProcessState.ExitCode()
+		// wait for process and resolve exit code
+		if err := cmd.Wait(); err != nil {
+			var exitError *exec.ExitError
+			if !errors.As(err, &exitError) {
+				// should be unreachable
+				a.waitErr = err
+			}
+
+			// store non-zero return code
+			r = exitError.ExitCode()
+		} else {
+			r = cmd.ProcessState.ExitCode()
+		}
+		fmsg.VPrintf("process %d exited with exit code %d", cmd.Process.Pid, r)
 	}
 
-	fmsg.VPrintf("process %d exited with exit code %d", a.cmd.Process.Pid, r)
+	// child process exited, resume output
+	fmsg.Resume()
 
 	// close wayland connection
 	if a.seal.wl != nil {
@@ -201,8 +188,10 @@ func (a *app) Wait() (int, error) {
 	e.Inner, e.DoErr = a.seal.store.Do(func(b state.Backend) {
 		e.InnerErr = func() error {
 			// destroy defunct state entry
-			if err := b.Destroy(a.cmd.Process.Pid); err != nil {
-				return err
+			if cmd := a.shim.Unwrap(); cmd != nil && a.seal.sys.saveState {
+				if err := b.Destroy(cmd.Process.Pid); err != nil {
+					return err
+				}
 			}
 
 			// enablements of remaining launchers
@@ -243,8 +232,7 @@ func (a *app) Wait() (int, error) {
 				}
 			}
 
-			a.abort <- errors.New("shim exited")
-			<-a.abort
+			a.shim.AbortWait(errors.New("shim exited"))
 			if err := a.seal.sys.Revert(ec); err != nil {
 				return err.(RevertCompoundError)
 			}
diff --git a/internal/app/system.go b/internal/app/system.go
index 7ccc2eeb..74b8e611 100644
--- a/internal/app/system.go
+++ b/internal/app/system.go
@@ -22,6 +22,8 @@ type appSealSys struct {
 	// target user sealed from config
 	user *user.User
 
+	needRevert bool
+	saveState  bool
 	*system.I
 
 	// protected by upstream mutex
diff --git a/internal/shim/parent.go b/internal/shim/parent.go
index cf6b125d..2ac9d6ae 100644
--- a/internal/shim/parent.go
+++ b/internal/shim/parent.go
@@ -1,106 +1,202 @@
 package shim
 
 import (
-	"encoding/gob"
 	"errors"
 	"net"
+	"os"
+	"os/exec"
+	"sync"
+	"sync/atomic"
 	"syscall"
+	"time"
 
 	"git.ophivana.moe/security/fortify/acl"
 	"git.ophivana.moe/security/fortify/internal/fmsg"
 )
 
-// called in the parent process
+// used by the parent process
 
-func ServeConfig(socket string, abort chan error, killShim func(), uid int, payload *Payload, wl *Wayland) error {
-	if payload.WL {
-		if f, err := net.DialUnix("unix", nil, &net.UnixAddr{Name: wl.Path, Net: "unix"}); err != nil {
-			return err
-		} else {
-			fmsg.VPrintf("connected to wayland at %q", wl.Path)
-			wl.UnixConn = f
-		}
+type Shim struct {
+	// user switcher process
+	cmd *exec.Cmd
+	// uid of shim target user
+	uid uint32
+	// whether to check shim pid
+	checkPid bool
+	// user switcher executable path
+	executable string
+	// path to setup socket
+	socket string
+	// shim setup abort reason and completion
+	abort     chan error
+	abortErr  atomic.Pointer[error]
+	abortOnce sync.Once
+	// wayland mediation, nil if disabled
+	wl *Wayland
+	// shim setup payload
+	payload *Payload
+}
+
+func New(executable string, uid uint32, socket string, wl *Wayland, payload *Payload) *Shim {
+	// checkPid is impossible at the moment since there is no way to obtain shim's pid
+	// this feature is disabled here until sudo is replaced by fortify suid wrapper
+	return &Shim{uid: uid, executable: executable, socket: socket, wl: wl, payload: payload}
+}
+
+func (s *Shim) String() string {
+	if s.cmd == nil {
+		return "(unused shim manager)"
 	}
+	return s.cmd.String()
+}
+
+func (s *Shim) Unwrap() *exec.Cmd {
+	return s.cmd
+}
+
+func (s *Shim) Abort(err error) {
+	s.abortOnce.Do(func() {
+		s.abortErr.Store(&err)
+		// s.abort is buffered so this will never block
+		s.abort <- err
+	})
+}
+
+func (s *Shim) AbortWait(err error) {
+	s.Abort(err)
+	<-s.abort
+}
 
-	// setup success state accessed by abort
-	var success bool
+type CommandBuilder func(shimEnv string) (args []string)
 
-	if c, err := net.ListenUnix("unix", &net.UnixAddr{Name: socket, Net: "unix"}); err != nil {
-		return err
+func (s *Shim) Start(f CommandBuilder) (*time.Time, error) {
+	var (
+		cf     chan *net.UnixConn
+		accept func()
+	)
+
+	// listen on setup socket
+	if c, a, err := s.serve(); err != nil {
+		return nil, fmsg.WrapErrorSuffix(err,
+			"cannot listen on shim setup socket:")
 	} else {
-		c.SetUnlinkOnClose(true)
+		// accepts a connection after each call to accept
+		// connections are sent to the channel cf
+		cf, accept = c, a
+	}
 
-		go func() {
-			err1 := <-abort
-			if !success {
-				fmsg.VPrintln("aborting shim setup, reason:", err1)
-				if err1 = c.Close(); err1 != nil {
-					fmsg.Println("cannot abort shim setup:", err1)
-				}
-			}
-			close(abort)
-		}()
+	// start user switcher process and save time
+	s.cmd = exec.Command(s.executable, f(EnvShim+"="+s.socket)...)
+	s.cmd.Env = []string{}
+	s.cmd.Stdin, s.cmd.Stdout, s.cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
+	s.cmd.Dir = "/"
+	fmsg.VPrintln("starting shim via user switcher:", s.cmd)
+	fmsg.Withhold() // withhold messages to stderr
+	if err := s.cmd.Start(); err != nil {
+		return nil, fmsg.WrapErrorSuffix(err,
+			"cannot start user switcher:")
+	}
+	startTime := time.Now().UTC()
+
+	// kill shim if something goes wrong and an error is returned
+	killShim := func() {
+		if err := s.cmd.Process.Signal(os.Interrupt); err != nil {
+			fmsg.Println("cannot terminate shim on faulted setup:", err)
+		}
+	}
+	defer func() { killShim() }()
+
+	accept()
+	conn := <-cf
+	if conn == nil {
+		return &startTime, fmsg.WrapErrorSuffix(*s.abortErr.Load(), "cannot accept call on setup socket:")
+	}
 
-		fmsg.VPrintf("configuring shim on socket %q", socket)
-		if err = acl.UpdatePerm(socket, uid, acl.Read, acl.Write, acl.Execute); err != nil {
-			fmsg.Println("cannot change permissions of shim setup socket:", err)
+	// authenticate against called provided uid and shim pid
+	if cred, err := peerCred(conn); err != nil {
+		return &startTime, fmsg.WrapErrorSuffix(*s.abortErr.Load(), "cannot retrieve shim credentials:")
+	} else if cred.Uid != s.uid {
+		fmsg.Printf("process %d owned by user %d tried to connect, expecting %d",
+			cred.Pid, cred.Uid, s.uid)
+		err = errors.New("compromised fortify build")
+		s.Abort(err)
+		return &startTime, err
+	} else if s.checkPid && cred.Pid != int32(s.cmd.Process.Pid) {
+		fmsg.Printf("process %d tried to connect to shim setup socket, expecting shim %d",
+			cred.Pid, s.cmd.Process.Pid)
+		err = errors.New("compromised target user")
+		s.Abort(err)
+		return &startTime, err
+	}
+
+	// serve payload and wayland fd if enabled
+	// this also closes the connection
+	err := s.payload.serve(conn, s.wl)
+	if err == nil {
+		killShim = func() {}
+	}
+	s.Abort(err) // aborting with nil indicates success
+	return &startTime, err
+}
+
+func (s *Shim) serve() (chan *net.UnixConn, func(), error) {
+	if s.abort != nil {
+		panic("attempted to serve shim setup twice")
+	}
+	s.abort = make(chan error, 1)
+
+	cf := make(chan *net.UnixConn)
+	accept := make(chan struct{}, 1)
+
+	if l, err := net.ListenUnix("unix", &net.UnixAddr{Name: s.socket, Net: "unix"}); err != nil {
+		return nil, nil, err
+	} else {
+		l.SetUnlinkOnClose(true)
+
+		fmsg.VPrintf("listening on shim setup socket %q", s.socket)
+		if err = acl.UpdatePerm(s.socket, int(s.uid), acl.Read, acl.Write, acl.Execute); err != nil {
+			fmsg.Println("cannot append ACL entry to shim setup socket:", err)
+			s.Abort(err) // ensures setup socket cleanup
 		}
 
 		go func() {
-			var conn *net.UnixConn
-			if conn, err = c.AcceptUnix(); err != nil {
-				if errors.Is(err, net.ErrClosed) {
-					fmsg.VPrintln("accept failed due to shim setup abort")
-				} else {
-					fmsg.Println("cannot accept connection from shim:", err)
-				}
-			} else {
-				if err = gob.NewEncoder(conn).Encode(*payload); err != nil {
-					fmsg.Println("cannot stream shim payload:", err)
-					killShim()
+			for {
+				select {
+				case err = <-s.abort:
+					if err != nil {
+						fmsg.VPrintln("aborting shim setup, reason:", err)
+					}
+					if err = l.Close(); err != nil {
+						fmsg.Println("cannot close setup socket:", err)
+					}
+					close(s.abort)
+					close(cf)
 					return
-				}
-
-				if payload.WL {
-					// get raw connection
-					var rc syscall.RawConn
-					if rc, err = wl.SyscallConn(); err != nil {
-						fmsg.Println("cannot obtain raw wayland connection:", err)
-						killShim()
-						return
+				case <-accept:
+					if conn, err0 := l.AcceptUnix(); err0 != nil {
+						s.Abort(err0) // does not block, breaks loop
+						cf <- nil     // receiver sees nil value and loads err0 stored during abort
 					} else {
-						go func() {
-							// pass wayland socket fd
-							if err = rc.Control(func(fd uintptr) {
-								if _, _, err = conn.WriteMsgUnix(nil, syscall.UnixRights(int(fd)), nil); err != nil {
-									fmsg.Println("cannot pass wayland connection to shim:", err)
-									killShim()
-									return
-								}
-								_ = conn.Close()
-
-								// block until shim exits
-								<-wl.done
-								fmsg.VPrintln("releasing wayland connection")
-							}); err != nil {
-								fmsg.Println("cannot obtain wayland connection fd:", err)
-							}
-						}()
+						cf <- conn
 					}
-				} else {
-					_ = conn.Close()
-				}
-			}
-
-			success = true
-			if err = c.Close(); err != nil {
-				if errors.Is(err, net.ErrClosed) {
-					fmsg.VPrintln("close failed due to shim setup abort")
-				} else {
-					fmsg.Println("cannot close shim socket:", err)
 				}
 			}
 		}()
-		return nil
 	}
+
+	return cf, func() { accept <- struct{}{} }, nil
+}
+
+// peerCred fetches peer credentials of conn
+func peerCred(conn *net.UnixConn) (ucred *syscall.Ucred, err error) {
+	var raw syscall.RawConn
+	if raw, err = conn.SyscallConn(); err != nil {
+		return
+	}
+
+	err0 := raw.Control(func(fd uintptr) {
+		ucred, err = syscall.GetsockoptUcred(int(fd), syscall.SOL_SOCKET, syscall.SO_PEERCRED)
+	})
+	err = errors.Join(err, err0)
+	return
 }
diff --git a/internal/shim/payload.go b/internal/shim/payload.go
index 5d9552d4..ba39ec04 100644
--- a/internal/shim/payload.go
+++ b/internal/shim/payload.go
@@ -1,6 +1,13 @@
 package shim
 
-import "git.ophivana.moe/security/fortify/helper/bwrap"
+import (
+	"encoding/gob"
+	"errors"
+	"net"
+
+	"git.ophivana.moe/security/fortify/helper/bwrap"
+	"git.ophivana.moe/security/fortify/internal/fmsg"
+)
 
 const EnvShim = "FORTIFY_SHIM"
 
@@ -17,3 +24,19 @@ type Payload struct {
 	// verbosity pass through
 	Verbose bool
 }
+
+func (p *Payload) serve(conn *net.UnixConn, wl *Wayland) error {
+	if err := gob.NewEncoder(conn).Encode(*p); err != nil {
+		return fmsg.WrapErrorSuffix(err,
+			"cannot stream shim payload:")
+	}
+
+	if wl != nil {
+		if err := wl.WriteUnix(conn); err != nil {
+			return errors.Join(err, conn.Close())
+		}
+	}
+
+	return fmsg.WrapErrorSuffix(conn.Close(),
+		"cannot close setup connection:")
+}
diff --git a/internal/shim/wayland.go b/internal/shim/wayland.go
index 05f04278..3bac55bc 100644
--- a/internal/shim/wayland.go
+++ b/internal/shim/wayland.go
@@ -1,8 +1,12 @@
 package shim
 
 import (
+	"fmt"
 	"net"
 	"sync"
+	"syscall"
+
+	"git.ophivana.moe/security/fortify/internal/fmsg"
 )
 
 // Wayland implements wayland mediation.
@@ -11,7 +15,7 @@ type Wayland struct {
 	Path string
 
 	// wayland connection
-	*net.UnixConn
+	conn *net.UnixConn
 
 	connErr error
 	sync.Once
@@ -19,10 +23,46 @@ type Wayland struct {
 	done chan struct{}
 }
 
+func (wl *Wayland) WriteUnix(conn *net.UnixConn) error {
+	// connect to host wayland socket
+	if f, err := net.DialUnix("unix", nil, &net.UnixAddr{Name: wl.Path, Net: "unix"}); err != nil {
+		return fmsg.WrapErrorSuffix(err,
+			fmt.Sprintf("cannot connect to wayland at %q:", wl.Path))
+	} else {
+		fmsg.VPrintf("connected to wayland at %q", wl.Path)
+		wl.conn = f
+	}
+
+	// set up for passing wayland socket
+	if rc, err := wl.conn.SyscallConn(); err != nil {
+		return fmsg.WrapErrorSuffix(err, "cannot obtain raw wayland connection:")
+	} else {
+		ec := make(chan error)
+		go func() {
+			// pass wayland connection fd
+			if err = rc.Control(func(fd uintptr) {
+				if _, _, err = conn.WriteMsgUnix(nil, syscall.UnixRights(int(fd)), nil); err != nil {
+					ec <- fmsg.WrapErrorSuffix(err, "cannot pass wayland connection to shim:")
+					return
+				}
+				ec <- nil
+
+				// block until shim exits
+				<-wl.done
+				fmsg.VPrintln("releasing wayland connection")
+			}); err != nil {
+				ec <- fmsg.WrapErrorSuffix(err, "cannot obtain wayland connection fd:")
+				return
+			}
+		}()
+		return <-ec
+	}
+}
+
 func (wl *Wayland) Close() error {
 	wl.Do(func() {
 		close(wl.done)
-		wl.connErr = wl.UnixConn.Close()
+		wl.connErr = wl.conn.Close()
 	})
 
 	return wl.connErr