From 8492239cba18c66c2b6dc76eaf3314f11b38b10b Mon Sep 17 00:00:00 2001
From: Ophestra Umiker <cat@ophivana.moe>
Date: Wed, 25 Sep 2024 14:00:30 +0900
Subject: [PATCH] helper/args: simplify argument parsing and eliminate excess
 memory copies

Signed-off-by: Ophestra Umiker <cat@ophivana.moe>
---
 dbus/dbus.go     |  7 ++--
 helper/args.go   | 97 ++++++++++++++----------------------------------
 helper/helper.go |  4 +-
 3 files changed, 33 insertions(+), 75 deletions(-)

diff --git a/dbus/dbus.go b/dbus/dbus.go
index d68dcfa..70151cf 100644
--- a/dbus/dbus.go
+++ b/dbus/dbus.go
@@ -54,8 +54,6 @@ func (p *Proxy) Seal(session, system *Config) error {
 		return errors.New("no configuration to seal")
 	}
 
-	seal := helper.NewArgs()
-
 	var args []string
 	if session != nil {
 		args = append(args, session.Args(p.session)...)
@@ -63,11 +61,12 @@ func (p *Proxy) Seal(session, system *Config) error {
 	if system != nil {
 		args = append(args, system.Args(p.system)...)
 	}
-	if err := seal.Seal(args); err != nil {
+	if seal, err := helper.NewCheckedArgs(args); err != nil {
 		return err
+	} else {
+		p.seal = seal
 	}
 
-	p.seal = seal
 	return nil
 }
 
diff --git a/helper/args.go b/helper/args.go
index 8ef07f4..fabbbf2 100644
--- a/helper/args.go
+++ b/helper/args.go
@@ -1,100 +1,59 @@
 package helper
 
 import (
-	"bytes"
 	"errors"
-	"fmt"
 	"io"
 	"strings"
-	"sync"
 )
 
 var (
 	ErrContainsNull = errors.New("argument contains null character")
 )
 
-// Args is sealed with a slice of arguments for writing to the helper args FD.
-// The sealing args is checked to not contain null characters.
-// Attempting to seal an instance twice will cause a panic.
-type Args interface {
-	Seal(args []string) error
-	io.WriterTo
-	fmt.Stringer
-}
+type argsFD []string
 
-// argsFD implements Args for helpers expecting null terminated arguments to a file descriptor.
-// argsFD must not be copied after first use.
-type argsFD struct {
-	seal []byte
-	sync.RWMutex
-}
-
-func (a *argsFD) Seal(args []string) error {
-	a.Lock()
-	defer a.Unlock()
-
-	if a.seal != nil {
-		panic("args sealed twice")
-	}
-
-	seal := bytes.Buffer{}
-
-	n := 0
-	for _, arg := range args {
-		// reject argument strings containing null
-		if hasNull(arg) {
-			return ErrContainsNull
+// checks whether any element contains the null character
+// must be called before args use and args must not be modified after call
+func (a argsFD) check() error {
+	for _, arg := range a {
+		for _, b := range arg {
+			if b == '\x00' {
+				return ErrContainsNull
+			}
 		}
-
-		// accumulate buffer size
-		n += len(arg) + 1
-	}
-	seal.Grow(n)
-
-	// write null terminated arguments
-	for _, arg := range args {
-		seal.WriteString(arg)
-		seal.WriteByte('\x00')
 	}
 
-	a.seal = seal.Bytes()
 	return nil
 }
 
-func (a *argsFD) WriteTo(w io.Writer) (int64, error) {
-	a.RLock()
-	defer a.RUnlock()
+func (a argsFD) WriteTo(w io.Writer) (int64, error) {
+	// assuming already checked
 
-	if a.seal == nil {
-		panic("attempted to activate unsealed args")
+	nt := 0
+	// write null terminated arguments
+	for _, arg := range a {
+		n, err := w.Write([]byte(arg + "\x00"))
+		nt += n
+
+		if err != nil {
+			return int64(nt), err
+		}
 	}
 
-	n, err := w.Write(a.seal)
-	return int64(n), err
+	return int64(nt), nil
 }
 
-func (a *argsFD) String() string {
+func (a argsFD) String() string {
 	if a == nil {
 		return "(invalid helper args)"
 	}
 
-	if a.seal == nil {
-		return "(unsealed helper args)"
-	}
-
-	return strings.ReplaceAll(string(a.seal), "\x00", " ")
+	return strings.Join(a, " ")
 }
 
-func hasNull(s string) bool {
-	for _, b := range s {
-		if b == '\x00' {
-			return true
-		}
-	}
-	return false
-}
-
-// NewArgs returns a new instance of Args
-func NewArgs() Args {
-	return new(argsFD)
+// NewCheckedArgs returns a checked argument writer for args.
+// Callers must not retain any references to args.
+func NewCheckedArgs(args []string) (io.WriterTo, error) {
+	a := argsFD(args)
+	return a, a.check()
 }
diff --git a/helper/helper.go b/helper/helper.go
index 6f52922..5488ec5 100644
--- a/helper/helper.go
+++ b/helper/helper.go
@@ -19,7 +19,6 @@ var (
 // Helper wraps *exec.Cmd and manages status and args fd.
 // Args is always 3 and status if set is always 4.
 type Helper struct {
-	lock sync.RWMutex
 	args io.WriterTo
 
 	statP [2]*os.File
@@ -32,6 +31,7 @@ type Helper struct {
 	// standard error. If non-nil, entry i becomes file descriptor 5+i.
 	ExtraFiles []*os.File
 
+	lock sync.RWMutex
 	*exec.Cmd
 }
 
@@ -177,7 +177,7 @@ func (h *Helper) Start() error {
 
 func New(wt io.WriterTo, name string, arg ...string) *Helper {
 	if wt == nil {
-		panic("attempted to create helper with nil argument writer")
+		panic("attempted to create helper with invalid argument writer")
 	}
 
 	return &Helper{args: wt, Cmd: exec.Command(name, arg...)}