Add support for env variables

README.md

@@ -14,4 +14,4 @@ OSX implements native isolation via the command `sandbox-exec`. The command line
 
 ### Linux
 
-On Linux, the functionality is implemented with the default command `systemd-run`, which should be available on most systems and allow a vast fine-grained sandbox configuration via SecComp and EBPF
+On Linux, the functionality is implemented with the default command `systemd-run`, which should be available on most systems and allow a vast fine-grained sandbox configuration via SecComp and EBPF

env.go

@@ -0,0 +1,19 @@
+package gozero
+
+import "fmt"
+
+type Variable struct {
+	Name  string
+	Value string
+}
+
+func (v *Variable) String() string {
+	return fmt.Sprintf("%s=%s", v.Name, v.Value)
+}
+
+func extendWithVars(env []string, vars ...Variable) []string {
+	for _, v := range vars {
+		env = append(env, v.String())
+	}
+	return env
+}

go.mod

@@ -1,6 +1,6 @@
 module github.com/projectdiscovery/gozero
 
-go 1.18
+go 1.20
 
 require (
 	github.com/projectdiscovery/utils v0.0.13

gozero.go

@@ -2,17 +2,35 @@ package gozero
 
 import (
 	"context"
+	"errors"
 	"os"
 	"os/exec"
+	"time"
 
 	"github.com/projectdiscovery/gozero/command"
+	errorutil "github.com/projectdiscovery/utils/errors"
 )
 
 type Gozero struct {
 	Options *Options
 }
 
 func New(options *Options) (*Gozero, error) {
+	// attempt to locate the interpreter by executing it
+	for _, engine := range options.Engines {
+		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+		defer cancel()
+
+		cmd := exec.CommandContext(ctx, engine)
+		err := cmd.Run()
+		if err == nil || errorutil.IsAny(err, exec.ErrWaitDelay) {
+			options.engine = engine
+			break
+		}
+	}
+	if options.engine == "" {
+		return nil, errors.New("no valid engine found")
+	}
 	return &Gozero{Options: options}, nil
 }
 
@@ -24,6 +42,8 @@ func (g *Gozero) Exec(ctx context.Context, input *Source, cmd *command.Command)
 	gCmd := exec.CommandContext(ctx, cmd.Name, cmd.Args...)
 	gCmd.Stdin = input.File
 	gCmd.Stdout = output.File
+	gCmd.Env = extendWithVars(gCmd.Environ(), input.Variables...)
+
 	return output, gCmd.Run()
 }
 
@@ -32,6 +52,9 @@ func (g *Gozero) Eval(ctx context.Context, src, input *Source, args ...string) (
 	if err != nil {
 		return nil, err
 	}
+	if g.Options.EarlyCloseFileDescriptor {
+		src.File.Close()
+	}
 	switch {
 	case g.Options.PreferStartProcess:
 		err = g.runWithApi(ctx, src, input, output, args...)
@@ -45,32 +68,34 @@ func (g *Gozero) Eval(ctx context.Context, src, input *Source, args ...string) (
 }
 
 func (g *Gozero) run(ctx context.Context, src, input, output *Source, args ...string) error {
-	cmdArgs := []string{src.Filename}
+	var cmdArgs []string
+	cmdArgs = append(cmdArgs, g.Options.Args...)
+	cmdArgs = append(cmdArgs, src.Filename)
 	cmdArgs = append(cmdArgs, args...)
-	gCmd := exec.CommandContext(ctx, g.Options.Engine, cmdArgs...)
+	gCmd := exec.CommandContext(ctx, g.Options.engine, cmdArgs...)
 	gCmd.Stdin = input.File
 	gCmd.Stdout = output.File
+	gCmd.Env = extendWithVars(gCmd.Environ(), input.Variables...)
 	return gCmd.Run()
 }
 
 func (g *Gozero) runWithApi(ctx context.Context, src, input, output *Source, args ...string) error {
 	var procAttr os.ProcAttr
-	cmdArgs := []string{g.Options.Engine, src.Filename}
+	cmdArgs := []string{g.Options.engine}
+	cmdArgs = append(cmdArgs, g.Options.Args...)
+	cmdArgs = append(cmdArgs, src.Filename)
 	cmdArgs = append(cmdArgs, args...)
 	procAttr.Files = []*os.File{input.File, output.File, nil}
-	proc, err := os.StartProcess(g.Options.Engine, cmdArgs, &procAttr)
+	procAttr.Env = extendWithVars(procAttr.Env, input.Variables...)
+	proc, err := os.StartProcess(g.Options.engine, cmdArgs, &procAttr)
 	if err != nil {
 		return err
 	}
 
 	go func() {
-		for {
-			select {
-			case <-ctx.Done():
-				proc.Kill()
-				return
-			default:
-			}
+		for range ctx.Done() {
+			proc.Kill()
+			return
 		}
 	}()
 

options.go

@@ -1,8 +1,10 @@
 package gozero
 
 type Options struct {
-	Command            string
-	Engine             string
-	PreferStartProcess bool
-	Sandbox            bool
+	Engines                  []string
+	Args                     []string
+	engine                   string
+	PreferStartProcess       bool
+	Sandbox                  bool
+	EarlyCloseFileDescriptor bool
 }

source.go

@@ -11,14 +11,15 @@ import (
 )
 
 type Source struct {
+	Variables       []Variable
 	Temporary       bool
 	CloseAfterWrite bool
 	Filename        string
 	File            *os.File
 }
 
 func NewSource() (*Source, error) {
-	return NewSourceWithString("")
+	return NewSourceWithString("", "")
 }
 
 func NewSourceWithFile(src string) (*Source, error) {
@@ -32,16 +33,16 @@ func NewSourceWithFile(src string) (*Source, error) {
 	return nil, errors.New("file does not exist")
 }
 
-func NewSourceWithBytes(src []byte) (*Source, error) {
-	return NewSourceWithReader(bytes.NewReader(src))
+func NewSourceWithBytes(src []byte, wantedPattern string) (*Source, error) {
+	return NewSourceWithReader(bytes.NewReader(src), wantedPattern)
 }
 
-func NewSourceWithString(src string) (*Source, error) {
-	return NewSourceWithReader(strings.NewReader(src))
+func NewSourceWithString(src, wantedPattern string) (*Source, error) {
+	return NewSourceWithReader(strings.NewReader(src), wantedPattern)
 }
 
-func NewSourceWithReader(src io.Reader) (*Source, error) {
-	srcFile, err := os.CreateTemp("", "")
+func NewSourceWithReader(src io.Reader, wantedPattern string) (*Source, error) {
+	srcFile, err := os.CreateTemp("", wantedPattern)
 	if err != nil {
 		return nil, err
 	}
@@ -85,3 +86,7 @@ func (s *Source) Cleanup() error {
 func (s *Source) ReadAll() ([]byte, error) {
 	return os.ReadFile(s.Filename)
 }
+
+func (s *Source) AddVariable(vars ...Variable) {
+	s.Variables = append(s.Variables, vars...)
+}