fmsg: support temporarily withholding output

Trying to print to a shared stdout is a terrible idea. This change makes it possible to withhold output for the lifetime of the sandbox.

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

config.go

@@ -65,7 +65,7 @@ func tryTemplate() {
 		} else {
 			fmt.Println(string(s))
 		}
-		os.Exit(0)
+		fmsg.Exit(0)
 	}
 }
 

error.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"errors"
-	"fmt"
 
 	"git.ophivana.moe/security/fortify/internal/app"
 	"git.ophivana.moe/security/fortify/internal/fmsg"
@@ -51,6 +50,6 @@ func logBaseError(err error, message string) {
 	if fmsg.AsBaseError(err, &e) {
 		fmsg.Print(e.Message())
 	} else {
-		fmt.Println(message, err)
+		fmsg.Println(message, err)
 	}
 }

helper/stub.go

@@ -11,6 +11,7 @@ import (
 	"testing"
 
 	"git.ophivana.moe/security/fortify/helper/bwrap"
+	"git.ophivana.moe/security/fortify/internal/fmsg"
 )
 
 // InternalChildStub is an internal function but exported because it is cross-package;
@@ -33,7 +34,7 @@ func InternalChildStub() {
 		genericStub(argsFD, statFD)
 	}
 
-	os.Exit(0)
+	fmsg.Exit(0)
 }
 
 // InternalReplaceExecCommand is an internal function but exported because it is cross-package;

internal/fmsg/defer.go

@@ -0,0 +1,70 @@
+package fmsg
+
+import (
+	"os"
+	"sync"
+	"sync/atomic"
+)
+
+var (
+	wstate   atomic.Bool
+	withhold = make(chan struct{}, 1)
+	msgbuf   = make(chan dOp, 64) // these ops are tiny so a large buffer is allocated for withholding output
+
+	dequeueOnce sync.Once
+	queueSync   sync.WaitGroup
+)
+
+func dequeue() {
+	go func() {
+		for {
+			select {
+			case op := <-msgbuf:
+				op.Do()
+				queueSync.Done()
+			case <-withhold:
+				<-withhold
+			}
+		}
+	}()
+}
+
+type dOp interface{ Do() }
+
+func Exit(code int) {
+	queueSync.Wait()
+	os.Exit(code)
+}
+
+func Withhold() {
+	if wstate.CompareAndSwap(false, true) {
+		withhold <- struct{}{}
+	}
+}
+
+func Resume() {
+	if wstate.CompareAndSwap(true, false) {
+		withhold <- struct{}{}
+	}
+}
+
+type dPrint []any
+
+func (v dPrint) Do() {
+	std.Print(v...)
+}
+
+type dPrintf struct {
+	format string
+	v      []any
+}
+
+func (d *dPrintf) Do() {
+	std.Printf(d.format, d.v...)
+}
+
+type dPrintln []any
+
+func (v dPrintln) Do() {
+	std.Println(v...)
+}

internal/fmsg/fmsg.go

@@ -4,38 +4,40 @@ package fmsg
 import (
 	"log"
 	"os"
-	"sync/atomic"
 )
 
-var (
+var std = log.New(os.Stderr, "fortify: ", 0)
-	std  = log.New(os.Stdout, "fortify: ", 0)
-	warn = log.New(os.Stderr, "fortify: ", 0)
-
-	verbose = new(atomic.Bool)
-)
 
 func SetPrefix(prefix string) {
 	prefix += ": "
 	std.SetPrefix(prefix)
-	warn.SetPrefix(prefix)
+	std.SetPrefix(prefix)
 }
 
 func Print(v ...any) {
-	warn.Print(v...)
+	dequeueOnce.Do(dequeue)
+	queueSync.Add(1)
+	msgbuf <- dPrint(v)
 }
 
 func Printf(format string, v ...any) {
-	warn.Printf(format, v...)
+	dequeueOnce.Do(dequeue)
+	queueSync.Add(1)
+	msgbuf <- &dPrintf{format, v}
 }
 
 func Println(v ...any) {
-	warn.Println(v...)
+	dequeueOnce.Do(dequeue)
+	queueSync.Add(1)
+	msgbuf <- dPrintln(v)
 }
 
 func Fatal(v ...any) {
-	warn.Fatal(v...)
+	Print(v...)
+	Exit(1)
 }
 
 func Fatalf(format string, v ...any) {
-	warn.Fatalf(format, v...)
+	Printf(format, v...)
+	Exit(1)
 }

internal/fmsg/verbose.go

@@ -1,5 +1,9 @@
 package fmsg
 
+import "sync/atomic"
+
+var verbose = new(atomic.Bool)
+
 func Verbose() bool {
 	return verbose.Load()
 }
@@ -10,12 +14,12 @@ func SetVerbose(v bool) {
 
 func VPrintf(format string, v ...any) {
 	if verbose.Load() {
-		std.Printf(format, v...)
+		Printf(format, v...)
 	}
 }
 
 func VPrintln(v ...any) {
 	if verbose.Load() {
-		std.Println(v...)
+		Println(v...)
 	}
 }

internal/init/main.go

@@ -129,7 +129,7 @@ func doInit(fd uintptr) {
 		select {
 		case s := <-sig:
 			fmsg.VPrintln("received", s.String())
-			os.Exit(0)
+			fmsg.Exit(0)
 		case w := <-info:
 			if w.wpid == cmd.Process.Pid {
 				switch {
@@ -147,10 +147,10 @@ func doInit(fd uintptr) {
 				}()
 			}
 		case <-done:
-			os.Exit(r)
+			fmsg.Exit(r)
 		case <-timeout:
 			fmsg.Println("timeout exceeded waiting for lingering processes")
-			os.Exit(r)
+			fmsg.Exit(r)
 		}
 	}
 }

internal/shim/main.go

@@ -134,9 +134,9 @@ func doShim(socket string) {
 			fmsg.VPrintln("wait:", err)
 		}
 		if b.Unwrap().ProcessState != nil {
-			os.Exit(b.Unwrap().ProcessState.ExitCode())
+			fmsg.Exit(b.Unwrap().ProcessState.ExitCode())
 		} else {
-			os.Exit(127)
+			fmsg.Exit(127)
 		}
 	}
 }

internal/state/print.go

@@ -21,8 +21,7 @@ func MustPrintLauncherStateSimpleGlobal(w **tabwriter.Writer, runDir string) {
 
 	// read runtime directory to get all UIDs
 	if dirs, err := os.ReadDir(path.Join(runDir, "state")); err != nil && !errors.Is(err, os.ErrNotExist) {
-		fmsg.Println("cannot read runtime directory:", err)
+		fmsg.Fatal("cannot read runtime directory:", err)
-		os.Exit(1)
 	} else {
 		for _, e := range dirs {
 			// skip non-directories
@@ -112,13 +111,11 @@ func (s *simpleStore) mustPrintLauncherState(w **tabwriter.Writer, now time.Time
 	}); err != nil {
 		fmsg.Printf("cannot perform action on store %q: %s", path.Join(s.path...), err)
 		if !ok {
-			fmsg.Println("store faulted before printing")
+			fmsg.Fatal("store faulted before printing")
-			os.Exit(1)
 		}
 	}
 
 	if innerErr != nil {
-		fmsg.Printf("cannot print launcher state for store %q: %s", path.Join(s.path...), innerErr)
+		fmsg.Fatalf("cannot print launcher state for store %q: %s", path.Join(s.path...), innerErr)
-		os.Exit(1)
 	}
 }

internal/system.go

@@ -109,7 +109,7 @@ func (s *Std) Open(name string) (fs.File, error) {
 	return os.Open(name)
 }
 func (s *Std) Exit(code int) {
-	os.Exit(code)
+	fmsg.Exit(code)
 }
 
 const xdgRuntimeDir = "XDG_RUNTIME_DIR"

state.go

@@ -30,6 +30,6 @@ func tryState() {
 			fmt.Println("No information available")
 		}
 
-		os.Exit(0)
+		fmsg.Exit(0)
 	}
 }