fix bug in command evaluation

dagu-org · Jan 7, 2025 · e566c84 · e566c84
1 parent 310e35c
commit e566c84
Show file tree

Hide file tree

Showing 5 changed files with 89 additions and 31 deletions.
diff --git a/internal/cmdutil/cmdutil.go b/internal/cmdutil/cmdutil.go
@@ -192,16 +192,15 @@ func SubstituteCommands(input string) (string, error) {
 		// Execute the command and replace the command with the output.
 		command := matches[i]
 
-		parser := shellwords.NewParser()
-		parser.ParseBacktick = true
-		parser.ParseEnv = false
-
-		res, err := parser.Parse(escapeReplacer.Replace(command))
+		sh := GetShellCommand("")
+		cmd := exec.Command(sh, "-c", command[1:len(command)-1])
+		cmd.Env = os.Environ()
+		out, err := cmd.Output()
 		if err != nil {
-			return "", fmt.Errorf("failed to substitute command: %w", err)
+			return "", fmt.Errorf("failed to execute command: %w", err)
 		}
 
-		ret = strings.ReplaceAll(ret, command, strings.Join(res, " "))
+		ret = strings.ReplaceAll(ret, command, strings.TrimSpace(string(out)))
 	}
 
 	return ret, nil
@@ -227,7 +226,16 @@ func GetShellCommand(configuredShell string) string {
 }
 
 type EvalOptions struct {
-	Variables []map[string]string
+	ExpandEnv  bool
+	Substitute bool
+	Variables  []map[string]string
+}
+
+func newEvalOptions() *EvalOptions {
+	return &EvalOptions{
+		ExpandEnv:  true,
+		Substitute: true,
+	}
 }
 
 type EvalOption func(*EvalOptions)
@@ -238,6 +246,25 @@ func WithVariables(vars map[string]string) EvalOption {
 	}
 }
 
+func WithoutExpandEnv() EvalOption {
+	return func(opts *EvalOptions) {
+		opts.ExpandEnv = false
+	}
+}
+
+func WithoutSubstitute() EvalOption {
+	return func(opts *EvalOptions) {
+		opts.Substitute = false
+	}
+}
+
+func OnlyReplaceVars() EvalOption {
+	return func(opts *EvalOptions) {
+		opts.ExpandEnv = false
+		opts.Substitute = false
+	}
+}
+
 var regEscapedKeyValue = regexp.MustCompile(`^[^\s=]+="[^"]+"$`)
 
 // BuildCommandEscapedString constructs a single shell-ready string from a command and its arguments.
@@ -279,33 +306,40 @@ func BuildCommandEscapedString(command string, args []string) string {
 
 // EvalString substitutes environment variables and commands in the input string
 func EvalString(input string, opts ...EvalOption) (string, error) {
-	options := &EvalOptions{}
+	options := newEvalOptions()
 	for _, opt := range opts {
 		opt(options)
 	}
 	value := input
 	for _, vars := range options.Variables {
 		value = replaceVars(value, vars)
 	}
-	value = os.ExpandEnv(value)
-	value, err := SubstituteCommands(value)
-	if err != nil {
-		return "", fmt.Errorf("failed to substitute string in %q: %w", input, err)
+	if options.Substitute {
+		var err error
+		value, err = SubstituteCommands(value)
+		if err != nil {
+			return "", fmt.Errorf("failed to substitute string in %q: %w", input, err)
+		}
+	}
+	if options.ExpandEnv {
+		value = os.ExpandEnv(value)
 	}
 	return value, nil
 }
 
 // EvalIntString substitutes environment variables and commands in the input string
 func EvalIntString(input string, opts ...EvalOption) (int, error) {
-	options := &EvalOptions{}
+	options := newEvalOptions()
 	for _, opt := range opts {
 		opt(options)
 	}
 	value := input
 	for _, vars := range options.Variables {
 		value = replaceVars(value, vars)
 	}
-	value = os.ExpandEnv(value)
+	if options.ExpandEnv {
+		value = os.ExpandEnv(value)
+	}
 	value, err := SubstituteCommands(value)
 	if err != nil {
 		return 0, err
@@ -321,7 +355,7 @@ func EvalIntString(input string, opts ...EvalOption) (int, error) {
 // variables and substituting command outputs. It takes a struct value and returns a new
 // modified struct value.
 func EvalStringFields[T any](obj T, opts ...EvalOption) (T, error) {
-	options := &EvalOptions{}
+	options := newEvalOptions()
 	for _, opt := range opts {
 		opt(options)
 	}
@@ -356,12 +390,20 @@ func processStructFields(v reflect.Value, opts *EvalOptions) error {
 			for _, vars := range opts.Variables {
 				value = replaceVars(value, vars)
 			}
-			value = os.ExpandEnv(value)
-			processed, err := SubstituteCommands(value)
-			if err != nil {
-				return fmt.Errorf("field %q: %w", t.Field(i).Name, err)
+
+			if opts.Substitute {
+				var err error
+				value, err = SubstituteCommands(value)
+				if err != nil {
+					return fmt.Errorf("field %q: %w", t.Field(i).Name, err)
+				}
 			}
-			field.SetString(processed)
+
+			if opts.ExpandEnv {
+				value = os.ExpandEnv(value)
+			}
+
+			field.SetString(value)
 
 		case reflect.Struct:
 			if err := processStructFields(field, opts); err != nil {

diff --git a/internal/digraph/condition.go b/internal/digraph/condition.go
@@ -56,19 +56,19 @@ func (c Condition) eval(ctx context.Context) (bool, error) {
 func (c Condition) evalCommand(ctx context.Context) (bool, error) {
 	var commandToRun string
 	if IsStepContext(ctx) {
-		command, err := GetStepContext(ctx).EvalString(c.Command)
+		command, err := GetStepContext(ctx).EvalString(c.Command, cmdutil.OnlyReplaceVars())
 		if err != nil {
 			return false, err
 		}
 		commandToRun = command
 	} else if IsContext(ctx) {
-		command, err := GetContext(ctx).EvalString(c.Command)
+		command, err := GetContext(ctx).EvalString(c.Command, cmdutil.OnlyReplaceVars())
 		if err != nil {
 			return false, err
 		}
 		commandToRun = command
 	} else {
-		command, err := cmdutil.EvalString(c.Command)
+		command, err := cmdutil.EvalString(c.Command, cmdutil.OnlyReplaceVars())
 		if err != nil {
 			return false, err
 		}

diff --git a/internal/digraph/condition_test.go b/internal/digraph/condition_test.go
@@ -66,6 +66,22 @@ func TestCondition_Eval(t *testing.T) {
 			},
 			wantErr: true,
 		},
+		{
+			name: "ComplexCommand",
+			condition: []Condition{
+				{
+					Command: "test 1 -eq 1",
+				},
+			},
+		},
+		{
+			name: "EvenMoreComplexCommand",
+			condition: []Condition{
+				{
+					Command: "df / | awk 'NR==2 {exit $4 > 5000 ? 0 : 1}'",
+				},
+			},
+		},
 		{
 			name: "CommandResultTest",
 			condition: []Condition{

diff --git a/internal/digraph/context.go b/internal/digraph/context.go
@@ -45,8 +45,9 @@ func (c Context) WithEnv(key, value string) Context {
 	return c
 }
 
-func (c Context) EvalString(s string) (string, error) {
-	return cmdutil.EvalString(s, cmdutil.WithVariables(c.envs))
+func (c Context) EvalString(s string, opts ...cmdutil.EvalOption) (string, error) {
+	opts = append(opts, cmdutil.WithVariables(c.envs))
+	return cmdutil.EvalString(s, opts...)
 }
 
 func NewContext(ctx context.Context, dag *DAG, client DBClient, requestID, logFile string) context.Context {

diff --git a/internal/digraph/step_context.go b/internal/digraph/step_context.go
@@ -61,11 +61,10 @@ func (c StepContext) MailerConfig() (mailer.Config, error) {
 	})
 }
 
-func (c StepContext) EvalString(s string) (string, error) {
-	return cmdutil.EvalString(s,
-		cmdutil.WithVariables(c.envs),
-		cmdutil.WithVariables(c.outputVariables.Variables()),
-	)
+func (c StepContext) EvalString(s string, opts ...cmdutil.EvalOption) (string, error) {
+	opts = append(opts, cmdutil.WithVariables(c.envs))
+	opts = append(opts, cmdutil.WithVariables(c.outputVariables.Variables()))
+	return cmdutil.EvalString(s, opts...)
 }
 
 func (c StepContext) EvalBool(value any) (bool, error) {