diff --git a/internal/runner/notify.go b/internal/runner/notify.go new file mode 100644 index 0000000..2538e94 --- /dev/null +++ b/internal/runner/notify.go @@ -0,0 +1,46 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +package runner + +// notifications keeps the service manager's sender in this wrapper while +// its engine runs in a separate scope. Only the child gets the private socket. +type notifications struct { + env []string + failed <-chan error + stop func() error +} + +func (n *notifications) close() error { + if n == nil { + return nil + } + return n.stop() +} + +func (n *notifications) watch(done <-chan struct{}, fail func(error)) <-chan error { + result := make(chan error, 1) + if n == nil { + result <- nil + return result + } + go func() { + select { + case err := <-n.failed: + select { + case <-done: + default: + fail(err) + } + result <- err + case <-done: + select { + case err := <-n.failed: + result <- err + default: + result <- nil + } + } + }() + return result +} diff --git a/internal/runner/notify_linux.go b/internal/runner/notify_linux.go new file mode 100644 index 0000000..fd298a0 --- /dev/null +++ b/internal/runner/notify_linux.go @@ -0,0 +1,139 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +package runner + +import ( + "errors" + "fmt" + "net" + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + + "golang.org/x/sys/unix" +) + +func startNotifications() (*notifications, error) { + original := os.Getenv("NOTIFY_SOCKET") + if original == "" { + return nil, nil + } + if original[0] != '/' && original[0] != '@' { + return nil, errors.New("unsupported NOTIFY_SOCKET address") + } + target, err := net.DialUnix("unixgram", nil, &net.UnixAddr{Name: original, Net: "unixgram"}) + if err != nil { + return nil, fmt.Errorf("connect notification socket: %w", err) + } + dir, err := os.MkdirTemp(os.Getenv("XDG_RUNTIME_DIR"), "berth-notify-") + if err != nil { + target.Close() + return nil, err + } + path := filepath.Join(dir, "notify") + source, err := net.ListenUnixgram("unixgram", &net.UnixAddr{Name: path, Net: "unixgram"}) + if err != nil { + target.Close() + return nil, errors.Join(err, os.RemoveAll(dir)) + } + failed := make(chan error, 1) + finished := make(chan struct{}) + var closing atomic.Bool + go func() { + defer close(finished) + defer source.Close() + for { + if err := forwardNotification(source, target); err != nil { + if !closing.Load() { + failed <- err + } + return + } + } + }() + var env []string + for _, entry := range os.Environ() { + if !strings.HasPrefix(entry, "NOTIFY_SOCKET=") { + env = append(env, entry) + } + } + env = append(env, "NOTIFY_SOCKET="+path) + return ¬ifications{env: env, failed: failed, stop: sync.OnceValue(func() error { + closing.Store(true) + source.Close() + target.Close() + <-finished + return os.RemoveAll(dir) + })}, nil +} + +func forwardNotification(source, target *net.UnixConn) error { + // systemd's receiver uses NOTIFY_BUFFER_MAX=PIPE_BUF and NOTIFY_FD_MAX=768. + data, control := make([]byte, 4096), make([]byte, unix.CmsgSpace(768*4)) + raw, err := source.SyscallConn() + if err != nil { + return err + } + var n, oobn, flags int + var receiveErr error + err = raw.Read(func(fd uintptr) bool { + for { + n, oobn, flags, _, receiveErr = unix.Recvmsg(int(fd), data, control, unix.MSG_CMSG_CLOEXEC) + if receiveErr != unix.EINTR { + return receiveErr != unix.EAGAIN + } + } + }) + if err != nil { + return err + } + if receiveErr != nil { + return receiveErr + } + messages, err := unix.ParseSocketControlMessage(control[:oobn]) + if err != nil { + return err + } + var fds []int + defer func() { + for _, fd := range fds { + _ = unix.Close(fd) + } + }() + for _, message := range messages { + if message.Header.Level == unix.SOL_SOCKET && message.Header.Type == unix.SCM_RIGHTS { + rights, err := unix.ParseUnixRights(&message) + if err != nil { + return err + } + fds = append(fds, rights...) + } + } + if flags&(unix.MSG_TRUNC|unix.MSG_CTRUNC) != 0 { + return errors.New("truncated notification datagram") + } + var ancillary []byte + if len(fds) > 0 { + ancillary = unix.UnixRights(fds...) + } + raw, err = target.SyscallConn() + if err != nil { + return err + } + var sendErr error + err = raw.Write(func(fd uintptr) bool { + for { + _, sendErr = unix.SendmsgN(int(fd), data[:n], ancillary, nil, 0) + if sendErr != unix.EINTR { + return sendErr != unix.EAGAIN + } + } + }) + if err != nil { + return err + } + return sendErr +} diff --git a/internal/runner/notify_linux_test.go b/internal/runner/notify_linux_test.go new file mode 100644 index 0000000..8fc4fe2 --- /dev/null +++ b/internal/runner/notify_linux_test.go @@ -0,0 +1,427 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +//go:build linux + +package runner + +import ( + "crypto/rand" + "io" + "net" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "golang.org/x/sys/unix" +) + +func TestRunNotificationSenderIsWrapper(t *testing.T) { + for _, leased := range []bool{false, true} { + name := "unleased" + if leased { + name = "leased" + } + t.Run(name, func(t *testing.T) { testRunNotification(t, leased, false) }) + } +} + +func TestRunNotifyCleanupFailure(t *testing.T) { + if unix.Geteuid() == 0 { + t.Skip("directory permission refusal requires an unprivileged user") + } + testRunNotification(t, false, true) +} + +func testRunNotification(t *testing.T, leased, denyCleanup bool) { + managerPath := filepath.Join(t.TempDir(), "notify") + manager, err := net.ListenUnixgram("unixgram", &net.UnixAddr{Name: managerPath, Net: "unixgram"}) + if err != nil { + t.Fatal(err) + } + defer manager.Close() + raw, err := manager.SyscallConn() + if err != nil { + t.Fatal(err) + } + var optionErr error + if err := raw.Control(func(fd uintptr) { + optionErr = unix.SetsockoptInt(int(fd), unix.SOL_SOCKET, unix.SO_PASSCRED, 1) + }); err != nil || optionErr != nil { + t.Fatalf("receive sender credentials: %v, %v", err, optionErr) + } + t.Setenv("NOTIFY_SOCKET", managerPath) + t.Setenv("BERTH_NOTIFY_CHILD", "1") + t.Setenv("XDG_RUNTIME_DIR", t.TempDir()) + if denyCleanup { + t.Setenv("BERTH_NOTIFY_DENY_CLEANUP", "1") + runtimeDir := os.Getenv("XDG_RUNTIME_DIR") + t.Cleanup(func() { + dirs, err := filepath.Glob(filepath.Join(runtimeDir, "berth-notify-*")) + if err != nil { + t.Error(err) + } + for _, dir := range dirs { + if err := os.Chmod(dir, 0700); err != nil { + t.Error(err) + } + if err := os.RemoveAll(dir); err != nil { + t.Error(err) + } + } + }) + } + executable, err := os.Executable() + if err != nil { + t.Fatal(err) + } + exited := make(chan int, 1) + socket := filepath.Join(t.TempDir(), "absent") + if leased { + socket, _ = parkDaemon(t, &parkCard{total: 24 << 30}, 16<<30) + } + go func() { + exited <- Run(Options{Socket: socket, Engine: "notify-fixture", Priority: "interactive", + Argv: []string{executable, "-test.run=^TestNotifyChild$"}, Stderr: io.Discard}) + }() + if err := manager.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal(err) + } + var sender int32 + for _, want := range []string{"READY=1\nSTATUS=child ready", "BARRIER=1"} { + data, control := make([]byte, 4096), make([]byte, 4096) + n, oobn, _, _, err := manager.ReadMsgUnix(data, control) + if err != nil { + t.Fatal(err) + } + messages, err := unix.ParseSocketControlMessage(control[:oobn]) + if err != nil { + t.Fatal(err) + } + for _, message := range messages { + switch message.Header.Type { + case unix.SCM_CREDENTIALS: + credentials, err := unix.ParseUnixCredentials(&message) + if err != nil { + t.Fatal(err) + } + sender = credentials.Pid + case unix.SCM_RIGHTS: + fds, err := unix.ParseUnixRights(&message) + if err != nil { + t.Fatal(err) + } + for _, fd := range fds { + _ = unix.Close(fd) + } + } + } + if string(data[:n]) != want { + t.Errorf("notification = %q, want %q", data[:n], want) + } + } + select { + case code := <-exited: + want := 0 + if denyCleanup { + want = 1 + } + if code != want { + t.Errorf("runner exit = %d, want %d", code, want) + } + case <-time.After(5 * time.Second): + t.Fatal("child barrier did not finish") + } + if sender != int32(os.Getpid()) { + t.Fatalf("notification sender = %d, want wrapper %d; a scoped child cannot ready the wrapper's service", sender, os.Getpid()) + } +} + +func closeNotify(t *testing.T, n *notifications) { + t.Helper() + if err := n.close(); err != nil { + t.Error(err) + } +} + +func notifySocket(n *notifications) string { + for _, entry := range n.env { + if value, ok := strings.CutPrefix(entry, "NOTIFY_SOCKET="); ok { + return value + } + } + return "" +} + +func notifyManager(t *testing.T, abstract bool) (*net.UnixConn, string) { + t.Helper() + path := filepath.Join(t.TempDir(), "manager") + if abstract { + path = "@berth-notify-test-" + rand.Text() + } + manager, err := net.ListenUnixgram("unixgram", &net.UnixAddr{Name: path, Net: "unixgram"}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { manager.Close() }) + t.Setenv("NOTIFY_SOCKET", path) + t.Setenv("XDG_RUNTIME_DIR", t.TempDir()) + return manager, path +} + +func TestNotifySetupAndCleanup(t *testing.T) { + for _, abstract := range []bool{false, true} { + name := "filesystem" + if abstract { + name = "abstract" + } + t.Run(name, func(t *testing.T) { + _, original := notifyManager(t, abstract) + n, err := startNotifications() + if err != nil { + t.Fatal(err) + } + defer closeNotify(t, n) + private := notifySocket(n) + if private == "" || private == original || os.Getenv("NOTIFY_SOCKET") != original { + t.Fatal("only the child must receive the private notification socket") + } + info, err := os.Stat(filepath.Dir(private)) + if err != nil || info.Mode().Perm() != 0700 { + t.Fatalf("private directory permissions: %v, %v", info, err) + } + closeNotify(t, n) + closeNotify(t, n) + if _, err := os.Stat(filepath.Dir(private)); !os.IsNotExist(err) { + t.Fatalf("private notification directory remains: %v", err) + } + select { + case err := <-n.failed: + t.Fatalf("intentional shutdown reported as failure: %v", err) + default: + } + }) + } +} + +func TestNotifyMissingSocketIsNoop(t *testing.T) { + t.Setenv("NOTIFY_SOCKET", "") + n, err := startNotifications() + if err != nil || n != nil { + t.Fatalf("without a notification socket: %v, %v", n, err) + } + closeNotify(t, n) +} + +func TestNotifyStartupRefusesInvalidSocket(t *testing.T) { + for _, address := range []string{"relative", "/nonexistent-berth-notify-test/socket"} { + t.Run(filepath.Base(address), func(t *testing.T) { + t.Setenv("NOTIFY_SOCKET", address) + marker := filepath.Join(t.TempDir(), "started") + code := Run(Options{Socket: filepath.Join(t.TempDir(), "absent"), + Argv: []string{"touch", marker}, Stderr: io.Discard}) + if code != 127 { + t.Fatalf("startup exit = %d, want 127", code) + } + if _, err := os.Stat(marker); !os.IsNotExist(err) { + t.Fatal("child ran despite notification setup failure") + } + }) + } +} + +func TestNotifyFailedExecCleansSocket(t *testing.T) { + notifyManager(t, false) + code := Run(Options{Socket: filepath.Join(t.TempDir(), "absent"), + Argv: []string{"/nonexistent-berth-notify-test/executable"}, Stderr: io.Discard}) + if code != 127 { + t.Fatalf("missing executable exit = %d", code) + } + entries, err := os.ReadDir(os.Getenv("XDG_RUNTIME_DIR")) + if err != nil || len(entries) != 0 { + t.Fatalf("failed exec left notification state: %v, %v", entries, err) + } +} + +func TestNotifyFailureClosesBarrier(t *testing.T) { + for _, oversized := range []bool{false, true} { + name := "target_closed" + if oversized { + name = "oversized" + } + t.Run(name, func(t *testing.T) { + manager, _ := notifyManager(t, false) + n, err := startNotifications() + if err != nil { + t.Fatal(err) + } + defer closeNotify(t, n) + if !oversized { + manager.Close() + } + conn, err := net.DialUnix("unixgram", nil, &net.UnixAddr{Name: notifySocket(n), Net: "unixgram"}) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + read, write, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer read.Close() + data := "BARRIER=1" + if oversized { + data = strings.Repeat("X", 4097) + } + writeNotify(t, conn, data, int(write.Fd())) + write.Close() + select { + case err := <-n.failed: + if err == nil { + t.Fatal("notification transport failure was not surfaced") + } + case <-time.After(5 * time.Second): + t.Fatal("notification transport failure was lost") + } + if err := read.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal(err) + } + if _, err := io.ReadAll(read); err != nil { + t.Fatalf("failed forwarding leaked a barrier descriptor: %v", err) + } + }) + } +} + +func writeNotify(t *testing.T, conn *net.UnixConn, data string, fds ...int) { + t.Helper() + raw, err := conn.SyscallConn() + if err != nil { + t.Fatal(err) + } + var sendErr error + err = raw.Write(func(fd uintptr) bool { + _, sendErr = unix.SendmsgN(int(fd), []byte(data), unix.UnixRights(fds...), nil, 0) + return sendErr != unix.EAGAIN + }) + if err != nil || sendErr != nil { + t.Fatalf("send notification: %v, %v", err, sendErr) + } +} + +func TestNotifyExitKeepsPendingError(t *testing.T) { + failed := make(chan error, 1) + failed <- io.ErrClosedPipe + done := make(chan struct{}) + close(done) + n := ¬ifications{failed: failed} + drained := false + got := <-n.watch(done, func(error) { drained = true }) + if got != io.ErrClosedPipe || drained { + t.Fatalf("completed child: error=%v, drained=%v; keep the failure without stopping an exited child", got, drained) + } +} + +func TestNotifyFileStore(t *testing.T) { + manager, _ := notifyManager(t, false) + n, err := startNotifications() + if err != nil { + t.Fatal(err) + } + defer closeNotify(t, n) + conn, err := net.DialUnix("unixgram", nil, &net.UnixAddr{Name: notifySocket(n), Net: "unixgram"}) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + state, err := os.CreateTemp(t.TempDir(), "state") + if err != nil { + t.Fatal(err) + } + defer state.Close() + if _, err := state.WriteString("stored engine state"); err != nil { + t.Fatal(err) + } + if _, err := state.Seek(0, io.SeekStart); err != nil { + t.Fatal(err) + } + payload := "FDSTORE=1\nFDNAME=engine-state\nX_BERTH_TEST=opaque" + writeNotify(t, conn, payload, int(state.Fd())) + if err := manager.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal(err) + } + data, control := make([]byte, 4096), make([]byte, 4096) + count, oobn, _, _, err := manager.ReadMsgUnix(data, control) + if err != nil { + t.Fatal(err) + } + messages, err := unix.ParseSocketControlMessage(control[:oobn]) + if err != nil { + t.Fatal(err) + } + var received []int + for _, message := range messages { + fds, err := unix.ParseUnixRights(&message) + if err != nil { + t.Fatal(err) + } + received = append(received, fds...) + } + if string(data[:count]) != payload || len(received) != 1 { + for _, fd := range received { + _ = unix.Close(fd) + } + t.Fatalf("file-store notification changed: %q, fds=%v", data[:count], received) + } + stored := os.NewFile(uintptr(received[0]), "received-state") + defer stored.Close() + content, err := io.ReadAll(stored) + if err != nil || string(content) != "stored engine state" { + t.Fatalf("forwarded descriptor: %q, %v", content, err) + } +} + +func TestNotifyChild(t *testing.T) { + if os.Getenv("BERTH_NOTIFY_CHILD") != "1" { + return + } + conn, err := net.DialUnix("unixgram", nil, &net.UnixAddr{Name: os.Getenv("NOTIFY_SOCKET"), Net: "unixgram"}) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + if _, err := conn.Write([]byte("READY=1\nSTATUS=child ready")); err != nil { + t.Fatal(err) + } + read, write, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer read.Close() + raw, err := conn.SyscallConn() + if err != nil { + t.Fatal(err) + } + var sendErr error + err = raw.Write(func(fd uintptr) bool { + _, sendErr = unix.SendmsgN(int(fd), []byte("BARRIER=1"), unix.UnixRights(int(write.Fd())), nil, 0) + return sendErr != unix.EAGAIN + }) + _ = write.Close() + if sendErr != nil { + t.Fatal(sendErr) + } + if err != nil { + t.Fatal(err) + } + if _, err := io.ReadAll(read); err != nil { + t.Fatal(err) + } + if os.Getenv("BERTH_NOTIFY_DENY_CLEANUP") == "1" { + if err := os.Chmod(filepath.Dir(os.Getenv("NOTIFY_SOCKET")), 0500); err != nil { + t.Fatal(err) + } + } +} diff --git a/internal/runner/notify_other.go b/internal/runner/notify_other.go new file mode 100644 index 0000000..e729f72 --- /dev/null +++ b/internal/runner/notify_other.go @@ -0,0 +1,8 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +//go:build !linux + +package runner + +func startNotifications() (*notifications, error) { return nil, nil } diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 879f597..75a3acd 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -130,7 +130,7 @@ func (o Options) pollEvery() time.Duration { return o.Retry } -func Run(o Options) int { +func Run(o Options) (code int) { say := func(format string, a ...any) { fmt.Fprintf(o.Stderr, "berth: "+format+"\n", a...) } if o.APIKeyEnv != "" || o.APIKeyFile != "" { key, err := (config.Engine{APIKeyEnv: o.APIKeyEnv, APIKeyFile: o.APIKeyFile}).APIKey() @@ -178,7 +178,17 @@ func Run(o Options) int { } } - cmd := command(o, o.Output) + cmd, notify, err := command(o, o.Output) + if err != nil { + say("notification socket: %v", err) + return 127 + } + defer func() { + if err := notify.close(); err != nil { + say("notification cleanup: %v", err) + code = 1 + } + }() session := comfySession{granted: time.Now()} if err := cmd.Start(); err != nil { say("start %s: %v", o.Argv[0], err) @@ -227,6 +237,17 @@ func Run(o Options) int { drain(opts, cmd, line, done, reason, say) }() } + notifyResult := notify.watch(done, func(error) { + select { + case <-done: + return + case <-drainTicket: + } + defer func() { drainTicket <- struct{}{} }() + opts := o + opts.Park, opts.Hook = nil, false + drain(opts, cmd, line, done, "notification forwarding failed", say) + }) if o.ExitComfyUIIdle { go func() { complete := false @@ -289,6 +310,10 @@ func Run(o Options) int { err = cmd.Wait() close(done) + if notifyErr := <-notifyResult; notifyErr != nil { + say("notification forwarding: %v", notifyErr) + return 1 + } if errors.Is(err, exec.ErrWaitDelay) { err = nil // exited cleanly; a grandchild still held the output copy } @@ -576,8 +601,19 @@ func slotsBusy(url, key string) (int, error) { return n, nil } -func runUnleased(o Options, reason string) int { - cmd := command(o, o.Output) +func runUnleased(o Options, reason string) (code int) { + say := func(format string, a ...any) { fmt.Fprintf(o.Stderr, "berth: "+format+"\n", a...) } + cmd, notify, err := command(o, o.Output) + if err != nil { + say("notification socket: %v", err) + return 127 + } + defer func() { + if err := notify.close(); err != nil { + say("notification cleanup: %v", err) + code = 1 + } + }() if err := cmd.Start(); err != nil { fmt.Fprintf(o.Stderr, "berth: start %s: %v\n", o.Argv[0], err) return 127 @@ -586,8 +622,17 @@ func runUnleased(o Options, reason string) int { defer record(o, cmd, reason)() done := make(chan struct{}) go forwardSignals(cmd, done) - err := cmd.Wait() + notifyResult := notify.watch(done, func(error) { + opts := o + opts.Park, opts.Hook = nil, false + drain(opts, cmd, nil, done, "notification forwarding failed", say) + }) + err = cmd.Wait() close(done) + if notifyErr := <-notifyResult; notifyErr != nil { + say("notification forwarding: %v", notifyErr) + return 1 + } if errors.Is(err, exec.ErrWaitDelay) { err = nil } @@ -608,9 +653,16 @@ func record(o Options, cmd *exec.Cmd, reason string) func() { return remove } -func command(o Options, out io.Writer) *exec.Cmd { +func command(o Options, out io.Writer) (*exec.Cmd, *notifications, error) { argv := scopeToChild(o.Engine, o.Argv, o.Log) + notify, err := startNotifications() + if err != nil { + return nil, nil, err + } cmd := exec.Command(argv[0], argv[1:]...) + if notify != nil { + cmd.Env = notify.env + } cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr if out != nil { cmd.Stdout, cmd.Stderr = io.MultiWriter(os.Stdout, out), io.MultiWriter(os.Stderr, out) @@ -619,7 +671,7 @@ func command(o Options, out io.Writer) *exec.Cmd { cmd.WaitDelay = 2 * time.Second } isolate(cmd) - return cmd + return cmd, notify, nil } // forwardSignals passes stop signals (llama-swap stopping a model, Ctrl-C)