diff --git a/src/job/utils.go b/src/job/utils.go index 3d6ad5ce..0c519554 100644 --- a/src/job/utils.go +++ b/src/job/utils.go @@ -26,6 +26,7 @@ import ( "context" "encoding/base64" "fmt" + "strconv" "time" "github.com/mitchellh/mapstructure" @@ -61,13 +62,26 @@ func logJob(ctx context.Context, args config.Args, globalConfig *GlobalConfig, a func setVarJob(ctx context.Context, args config.Args, globalConfig *GlobalConfig, a *metrics.Accumulator, logger *zap.Logger) (data any, err error) { var jobConfig struct { Value string + Type string } if err := mapstructure.Decode(args, &jobConfig); err != nil { return nil, fmt.Errorf("error parsing job config: %w", err) } - return templates.ParseAndExecute(logger, jobConfig.Value, ctx), nil + switch jobConfig.Type { + case "int": + return strconv.Atoi(templates.ParseAndExecute(logger, jobConfig.Value, ctx)) + case "uint": + val, err := strconv.ParseUint(templates.ParseAndExecute(logger, jobConfig.Value, ctx), 10, 32) + return uint(val), err + case "int64": + return strconv.ParseInt(templates.ParseAndExecute(logger, jobConfig.Value, ctx), 10, 64) + case "uint64": + return strconv.ParseUint(templates.ParseAndExecute(logger, jobConfig.Value, ctx), 10, 64) + default: + return templates.ParseAndExecute(logger, jobConfig.Value, ctx), nil + } } // "check" in config diff --git a/src/utils/templates/math.go b/src/utils/templates/math.go new file mode 100644 index 00000000..dc6c5ce5 --- /dev/null +++ b/src/utils/templates/math.go @@ -0,0 +1,49 @@ +package templates + +func mod(lhs, rhs int) int { + return lhs % rhs +} + +func add(lhs, rhs int) int { + return lhs + rhs +} + +func sub(lhs, rhs int) int { + return lhs - rhs +} + +func umod(lhs, rhs uint) uint { + return lhs % rhs +} + +func uadd(lhs, rhs uint) uint { + return lhs + rhs +} + +func usub(lhs, rhs uint) uint { + return lhs - rhs +} + +func mod64(lhs, rhs int64) int64 { + return lhs % rhs +} + +func add64(lhs, rhs int64) int64 { + return lhs + rhs +} + +func sub64(lhs, rhs int64) int64 { + return lhs - rhs +} + +func umod64(lhs, rhs uint64) uint64 { + return lhs % rhs +} + +func uadd64(lhs, rhs uint64) uint64 { + return lhs + rhs +} + +func usub64(lhs, rhs uint64) uint64 { + return lhs - rhs +} diff --git a/src/utils/templates/templates.go b/src/utils/templates/templates.go index bfaf8a1b..1e24e0d5 100644 --- a/src/utils/templates/templates.go +++ b/src/utils/templates/templates.go @@ -91,14 +91,6 @@ func randomAplhaNum(n int) string { return randomString(n, "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789") } -func mod(lhs, rhs int) int { - return lhs % rhs -} - -func add(lhs, rhs int) int { - return lhs + rhs -} - // ContextKey used to work with context and not trigger linter type ContextKey string @@ -144,6 +136,16 @@ func Parse(input string) (*template.Template, error) { "get_url": getURLContent, "mod": mod, "add": add, + "sub": sub, + "umod": umod, + "uadd": uadd, + "usub": usub, + "mod64": mod64, + "add64": add64, + "sub64": sub64, + "umod64": umod64, + "uadd64": uadd64, + "usub64": usub64, "ctx_key": ctxKey, "cookie_string": cookieString, }).Parse(input)