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>

3 files changed, 91 insertions, 15 deletions

diff --git a/internal/fmsg/defer.go b/internal/fmsg/defer.go
new file mode 100644
index 00000000..cee3c737
--- /dev/null
+++ b/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...)
+}
diff --git a/internal/fmsg/fmsg.go b/internal/fmsg/fmsg.go
index c72913f8..dbe590bf 100644
--- a/internal/fmsg/fmsg.go
+++ b/internal/fmsg/fmsg.go
@@ -4,38 +4,40 @@ package fmsg
 import (
 	"log"
 	"os"
-	"sync/atomic"
 )
 
-var (
-	std  = log.New(os.Stdout, "fortify: ", 0)
-	warn = log.New(os.Stderr, "fortify: ", 0)
-
-	verbose = new(atomic.Bool)
-)
+var std = log.New(os.Stderr, "fortify: ", 0)
 
 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)
 }
diff --git a/internal/fmsg/verbose.go b/internal/fmsg/verbose.go
index 36faadeb..72a92a69 100644
--- a/internal/fmsg/verbose.go
+++ b/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...)
 	}
 }