use StringSlice

projectdiscovery · Jul 21, 2023 · 9831d71 · 9831d71
1 parent a96e9a9
commit 9831d71
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 11 deletions.
diff --git a/ratelimit_var.go b/ratelimit_var.go
@@ -69,28 +69,31 @@ func (rateLimitMap RateLimitMap) String() string {
 
 	var items string
 	for k, v := range rateLimitMap.kv {
-		items += fmt.Sprintf("\"%s\"=\"%s\"%s", k, v.Duration.String(), kvSep)
+		items += fmt.Sprintf("\"%s\":\"%d/%s\",", k, v.MaxCount, v.Duration.String())
 	}
-	defaultBuilder.WriteString(stringsutil.TrimSuffixAny(items, ",", "="))
+	defaultBuilder.WriteString(stringsutil.TrimSuffixAny(items, ",", ":"))
 	defaultBuilder.WriteString("}")
 	return defaultBuilder.String()
 }
 
 // RateLimitMapVar adds a ratelimit flag with a longname
-func (flagSet *FlagSet) RateLimitMapVar(field *RateLimitMap, long string, defaultValue []string, usage string) *FlagData {
-	return flagSet.RateLimitMapVarP(field, long, "", defaultValue, usage)
+func (flagSet *FlagSet) RateLimitMapVar(field *RateLimitMap, long string, defaultValue []string, usage string, options Options) *FlagData {
+	return flagSet.RateLimitMapVarP(field, long, "", defaultValue, usage, options)
 }
 
 // RateLimitMapVarP adds a ratelimit flag with a short name and long name.
 // It is equivalent to RateLimitMapVar, and also allows specifying ratelimits in days (e.g., "hackertarget=2/d" 2 requests per day, which is equivalent to 24h).
-func (flagSet *FlagSet) RateLimitMapVarP(field *RateLimitMap, long, short string, defaultValue []string, usage string) *FlagData {
+func (flagSet *FlagSet) RateLimitMapVarP(field *RateLimitMap, long, short string, defaultValue StringSlice, usage string, options Options) *FlagData {
 	if field == nil {
 		panic(fmt.Errorf("field cannot be nil for flag -%v", long))
 	}
 
-	for _, item := range defaultValue {
-		if err := field.Set(item); err != nil {
-			panic(fmt.Errorf("failed to set default value for flag -%v: %v", long, err))
+	for _, defaultItem := range defaultValue {
+		values, _ := ToStringSlice(defaultItem, options)
+		for _, value := range values {
+			if err := field.Set(value); err != nil {
+				panic(fmt.Errorf("failed to set default value for flag -%v: %v", long, err))
+			}
 		}
 	}
 

diff --git a/ratelimit_var_test.go b/ratelimit_var_test.go
@@ -14,7 +14,7 @@ func TestRateLimitMapVar(t *testing.T) {
 		var rateLimitMap RateLimitMap
 		flagSet := NewFlagSet()
 		flagSet.CreateGroup("Config", "Config",
-			flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", []string{"hackertarget=1/ms"}, "rate limits"),
+			flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", []string{"hackertarget=1/ms"}, "rate limits", CommaSeparatedStringSliceOptions),
 		)
 		os.Args = []string{
 			os.Args[0],
@@ -25,11 +25,27 @@ func TestRateLimitMapVar(t *testing.T) {
 		tearDown(t.Name())
 	})
 
+	t.Run("multiple-default-value", func(t *testing.T) {
+		var rateLimitMap RateLimitMap
+		flagSet := NewFlagSet()
+		flagSet.CreateGroup("Config", "Config",
+			flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", []string{"hackertarget=1/s,github=1/ms"}, "rate limits", CommaSeparatedStringSliceOptions),
+		)
+		os.Args = []string{
+			os.Args[0],
+		}
+		err := flagSet.Parse()
+		assert.Nil(t, err)
+		assert.Equal(t, RateLimit{MaxCount: 1, Duration: time.Second}, rateLimitMap.AsMap()["hackertarget"])
+		assert.Equal(t, RateLimit{MaxCount: 1, Duration: time.Millisecond}, rateLimitMap.AsMap()["github"])
+		tearDown(t.Name())
+	})
+
 	t.Run("valid-rate-limit", func(t *testing.T) {
 		var rateLimitMap RateLimitMap
 		flagSet := NewFlagSet()
 		flagSet.CreateGroup("Config", "Config",
-			flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", nil, "rate limits"),
+			flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", nil, "rate limits", CommaSeparatedStringSliceOptions),
 		)
 		os.Args = []string{
 			os.Args[0],
@@ -45,7 +61,7 @@ func TestRateLimitMapVar(t *testing.T) {
 		var rateLimitMap RateLimitMap
 		flagSet := NewFlagSet()
 		flagSet.CreateGroup("Config", "Config",
-			flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", nil, "rate limits"),
+			flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", nil, "rate limits", CommaSeparatedStringSliceOptions),
 		)
 		os.Args = []string{
 			os.Args[0],