From cafed5f23468daf2765655325ebe7d879acdde24 Mon Sep 17 00:00:00 2001
From: Ophestra Umiker <cat@ophivana.moe>
Date: Mon, 21 Oct 2024 21:23:56 +0900
Subject: [PATCH] shim: abort setup on failed start and process exit

Shim setup listens on a socket in the process share, if shim setup hasn't happened on exit revert will fail. This change makes sure shim setup is aborted on a doomed launch.

Signed-off-by: Ophestra Umiker <cat@ophivana.moe>
---
 internal/app/app.go     |  2 ++
 internal/app/start.go   |  7 ++++++-
 internal/shim/parent.go | 35 +++++++++++++++++++++++++++++------
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/internal/app/app.go b/internal/app/app.go
index 5c5e345..0055dd7 100644
--- a/internal/app/app.go
+++ b/internal/app/app.go
@@ -24,6 +24,8 @@ type app struct {
 	id *ID
 	// underlying user switcher process
 	cmd *exec.Cmd
+	// shim setup abort reason and completion
+	abort chan error
 	// child process related information
 	seal *appSeal
 	// error returned waiting for process
diff --git a/internal/app/start.go b/internal/app/start.go
index b6c9404..464d23d 100644
--- a/internal/app/start.go
+++ b/internal/app/start.go
@@ -63,7 +63,8 @@ func (a *app) Start() error {
 	a.cmd.Stdin, a.cmd.Stdout, a.cmd.Stderr = os.Stdin, os.Stdout, os.Stderr
 	a.cmd.Dir = a.seal.RunDirPath
 
-	if err := shim.ServeConfig(confSockPath, a.seal.sys.UID(), &shim.Payload{
+	a.abort = make(chan error)
+	if err := shim.ServeConfig(confSockPath, a.abort, a.seal.sys.UID(), &shim.Payload{
 		Argv:  a.seal.command,
 		Exec:  shimExec,
 		Bwrap: a.seal.sys.bwrap,
@@ -71,6 +72,8 @@ func (a *app) Start() error {
 
 		Verbose: fmsg.Verbose(),
 	}, a.seal.wl); err != nil {
+		a.abort <- err
+		<-a.abort
 		return fmsg.WrapErrorSuffix(err,
 			"cannot serve shim setup:")
 	}
@@ -232,6 +235,8 @@ func (a *app) Wait() (int, error) {
 				}
 			}
 
+			a.abort <- errors.New("shim exited")
+			<-a.abort
 			if err := a.seal.sys.Revert(ec); err != nil {
 				return err.(RevertCompoundError)
 			}
diff --git a/internal/shim/parent.go b/internal/shim/parent.go
index 6c393c5..c6209fe 100644
--- a/internal/shim/parent.go
+++ b/internal/shim/parent.go
@@ -13,7 +13,7 @@ import (
 
 // called in the parent process
 
-func ServeConfig(socket string, uid int, payload *Payload, wl *Wayland) error {
+func ServeConfig(socket string, abort chan error, 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
@@ -23,9 +23,25 @@ func ServeConfig(socket string, uid int, payload *Payload, wl *Wayland) error {
 		}
 	}
 
+	// setup success state accessed by abort
+	var success bool
+
 	if c, err := net.ListenUnix("unix", &net.UnixAddr{Name: socket, Net: "unix"}); err != nil {
 		return err
 	} else {
+		c.SetUnlinkOnClose(true)
+
+		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)
+		}()
+
 		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)
@@ -34,7 +50,11 @@ func ServeConfig(socket string, uid int, payload *Payload, wl *Wayland) error {
 		go func() {
 			var conn *net.UnixConn
 			if conn, err = c.AcceptUnix(); err != nil {
-				fmsg.Println("cannot accept connection from shim:", err)
+				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)
@@ -70,11 +90,14 @@ func ServeConfig(socket string, uid int, payload *Payload, wl *Wayland) error {
 					_ = conn.Close()
 				}
 			}
+
+			success = true
 			if err = c.Close(); err != nil {
-				fmsg.Println("cannot close shim socket:", err)
-			}
-			if err = os.Remove(socket); err != nil && !errors.Is(err, os.ErrNotExist) {
-				fmsg.Println("cannot remove dangling shim socket:", err)
+				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