From 96d539898afda8f07f2024acb8b7c45898c23a80 Mon Sep 17 00:00:00 2001
From: Alexander Scheel <alex.scheel@hashicorp.com>
Date: Tue, 1 Nov 2022 09:42:31 -0400
Subject: [PATCH] feat(parseutil): Add ParseDurationSecondWithUnits

This allows ParseDurationSecond to specify custom additional units (such
as days, months, and years) which must stand independently of other
units (e.g., "5d6mo7y" is not allowed, but "5d", "6mo", and "7y" are),
mapping them to time.Duration values.

Signed-off-by: Alexander Scheel <alex.scheel@hashicorp.com>
---
 parseutil/parseutil.go      | 42 ++++++++++++++++++----
 parseutil/parseutil_test.go | 69 +++++++++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+), 7 deletions(-)

diff --git a/parseutil/parseutil.go b/parseutil/parseutil.go
index e469499..ced5b58 100644
--- a/parseutil/parseutil.go
+++ b/parseutil/parseutil.go
@@ -97,6 +97,22 @@ func ParseCapacityString(in interface{}) (uint64, error) {
 // a time.Duration; when units are missing (such as when a numeric type is
 // provided), the duration is assumed to be in seconds.
 func ParseDurationSecond(in interface{}) (time.Duration, error) {
+	return ParseDurationSecondWithUnits(in, map[string]time.Duration{
+		"d": 24 * time.Hour,
+	})
+}
+
+// Parse a duration from an arbitrary value (a string or numeric value) into
+// a time.Duration; when units are missing (such as when a numeric type is
+// provided), the duration is assumed to be in seconds.
+//
+// Unlike ParseDurationSecond, this allows the caller to augment with the
+// desired additional units via a suffix map.
+//
+// This behaves like the existing implementation, wherein the augmented units
+// must be a suffix and must be strictly by an (optionally signed) integer.
+// Mixing custom units with built-in units is not supported.
+func ParseDurationSecondWithUnits(in interface{}, units map[string]time.Duration) (time.Duration, error) {
 	var dur time.Duration
 	jsonIn, ok := in.(json.Number)
 	if ok {
@@ -114,15 +130,11 @@ func ParseDurationSecond(in interface{}) (time.Duration, error) {
 			return time.Duration(v) * time.Second, nil
 		}
 
-		if strings.HasSuffix(inp, "d") {
-			v, err := strconv.ParseInt(inp[:len(inp)-1], 10, 64)
-			if err != nil {
-				return dur, err
-			}
-			return time.Duration(v) * 24 * time.Hour, nil
+		found, parsed, err := parseDurationStringWithUnits(inp, units)
+		if found {
+			return parsed, err
 		}
 
-		var err error
 		if dur, err = time.ParseDuration(inp); err != nil {
 			return dur, err
 		}
@@ -151,6 +163,22 @@ func ParseDurationSecond(in interface{}) (time.Duration, error) {
 	return dur, nil
 }
 
+func parseDurationStringWithUnits(inp string, units map[string]time.Duration) (bool, time.Duration, error) {
+	for suffix, value := range units {
+		if strings.HasSuffix(inp, suffix) {
+			v, err := strconv.ParseInt(inp[:len(inp)-len(suffix)], 10, 64)
+			if err != nil {
+				return true, 0 * time.Second, err
+			}
+
+			return true, time.Duration(v) * value, nil
+		}
+	}
+
+	return false, 0 * time.Second, nil
+
+}
+
 // Parse an absolute timestamp from the provided arbitrary value (string or
 // numeric value). When an untyped numeric value is provided, it is assumed
 // to be seconds from the Unix Epoch.
diff --git a/parseutil/parseutil_test.go b/parseutil/parseutil_test.go
index 75b62eb..b3bba5e 100644
--- a/parseutil/parseutil_test.go
+++ b/parseutil/parseutil_test.go
@@ -231,6 +231,75 @@ func Test_ParseDurationSecond(t *testing.T) {
 	}
 }
 
+func Test_ParseDurationSecondWithUnits(t *testing.T) {
+	type Test struct {
+		in      interface{}
+		units   map[string]time.Duration
+		out     time.Duration
+		invalid bool
+	}
+
+	daysOnly := map[string]time.Duration{
+		"d": 24 * time.Hour,
+	}
+	daysMonthsYears := map[string]time.Duration{
+		"d":  24 * time.Hour,
+		"mo": 30 * 24 * time.Hour,
+		"y":  365 * 24 * time.Hour,
+	}
+
+	tests := []Test{
+		// String inputs
+		{in: "50ms", units: daysOnly, out: 50 * time.Millisecond},
+		{in: "50s", units: daysOnly, out: 50 * time.Second},
+		{in: "5m", units: daysMonthsYears, out: 5 * time.Minute},
+		{in: "6h", units: daysMonthsYears, out: 6 * time.Hour},
+		{in: "5d", units: daysOnly, out: 5 * 24 * time.Hour},
+		{in: "-5d", units: daysOnly, out: -5 * 24 * time.Hour},
+		{in: "05d", units: daysOnly, out: 5 * 24 * time.Hour},
+		{in: "500d", units: daysOnly, out: 500 * 24 * time.Hour},
+		{in: "5mo", units: daysMonthsYears, out: 5 * 30 * 24 * time.Hour},
+		{in: "5y", units: daysMonthsYears, out: 5 * 365 * 24 * time.Hour},
+	}
+
+	// Invalid inputs
+	for _, s := range []string{
+		"d",
+		"mo",
+		"yr",
+		"m",
+		"5d6mo12y",
+		"5d6mo",
+		"1d2y",
+		"2mo3y",
+		"3y2mo",
+	} {
+		tests = append(tests, Test{
+			in:      s,
+			units:   daysMonthsYears,
+			invalid: true,
+		})
+	}
+
+	for _, test := range tests {
+		out, err := ParseDurationSecondWithUnits(test.in, test.units)
+		if test.invalid {
+			if err == nil {
+				t.Fatalf("%q: expected error, got nil", test.in)
+			}
+			continue
+		}
+
+		if err != nil {
+			t.Fatalf("unexpected error parsing %v: %v", test.in, err)
+		}
+
+		if out != test.out {
+			t.Fatalf("%q: expected: %q, got: %q", test.in, test.out, out)
+		}
+	}
+}
+
 func Test_ParseAbsoluteTime(t *testing.T) {
 	testCases := []struct {
 		inp      interface{}