From 9214b293e3f69965faa70ad024394cfb0f62e36b Mon Sep 17 00:00:00 2001
From: Mihai Parparita <mihai@tailscale.com>
Date: Wed, 28 Sep 2022 16:37:31 -0700
Subject: [PATCH] tstime: add ParseDuration helper function

More expressive than time.ParseDuration, also accepting d (days) and
w (weeks) literals.

Signed-off-by: Mihai Parparita <mihai@tailscale.com>
---
 tstime/tstime.go      | 24 ++++++++++++++++++++++++
 tstime/tstime_test.go | 23 +++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/tstime/tstime.go b/tstime/tstime.go
index 61aab5ca4..669393455 100644
--- a/tstime/tstime.go
+++ b/tstime/tstime.go
@@ -8,6 +8,8 @@ package tstime
 import (
 	"errors"
 	"fmt"
+	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -141,3 +143,25 @@ func Parse3339(s string) (time.Time, error) {
 func Parse3339B(b []byte) (time.Time, error) {
 	return parse3339m(mem.B(b))
 }
+
+// ParseDuration is more expressive than time.ParseDuration, also accepting d
+// (days) and w (weeks) literals.
+func ParseDuration(s string) (time.Duration, error) {
+	for {
+		end := strings.IndexAny(s, "dw")
+		if end < 0 {
+			break
+		}
+		start := end - (len(s[:end]) - len(strings.TrimRight(s[:end], "012345789")))
+		n, err := strconv.Atoi(s[start:end])
+		if err != nil {
+			return 0, err
+		}
+		hours := 24
+		if s[end] == 'w' {
+			hours *= 7
+		}
+		s = s[:start] + s[end+1:] + strconv.Itoa(n*hours) + "h"
+	}
+	return time.ParseDuration(s)
+}
diff --git a/tstime/tstime_test.go b/tstime/tstime_test.go
index a2f76a58d..c458e01be 100644
--- a/tstime/tstime_test.go
+++ b/tstime/tstime_test.go
@@ -102,6 +102,29 @@ func TestParseInt(t *testing.T) {
 	}
 }
 
+func TestParseDuration(t *testing.T) {
+	tests := []struct {
+		in   string
+		want time.Duration
+	}{
+		{"1h", time.Hour},
+		{"1d", 24 * time.Hour},
+		{"1d1d", 48 * time.Hour},
+		{"1h1d", 25 * time.Hour},
+		{"1d1h", 25 * time.Hour},
+		{"1w", 7 * 24 * time.Hour},
+		{"1w1d1h", 8*24*time.Hour + time.Hour},
+		{"1w1d1h", 8*24*time.Hour + time.Hour},
+		{"1y", 0},
+		{"", 0},
+	}
+	for _, tt := range tests {
+		if got, _ := ParseDuration(tt.in); got != tt.want {
+			t.Errorf("ParseDuration(%q) = %d; want %d", tt.in, got, tt.want)
+		}
+	}
+}
+
 func BenchmarkGoParse3339(b *testing.B) {
 	run := func(in string) func(*testing.B) {
 		return func(b *testing.B) {