3 files changed, 239 insertions, 80 deletions

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