helper/args: simplify argument parsing and eliminate excess memory copies

Signed-off-by: Ophestra Umiker <cat@ophivana.moe>
This commit is contained in:
Ophestra Umiker 2024-09-25 14:00:30 +09:00
parent a8b4b3634b
commit 8492239cba
Signed by: cat
SSH Key Fingerprint: SHA256:gQ67O0enBZ7UdZypgtspB2FDM1g3GVw8nX0XSdcFw8Q
3 changed files with 33 additions and 75 deletions

View File

@ -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
}

View File

@ -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) {
// 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()
}

View File

@ -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...)}