From ccce3f8e0a081f893402be0cb1d0142a663aa02c Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Thu, 2 Jun 2022 19:22:51 -0700
Subject: [PATCH] cmd/golink: expand, add tests for URL expansions

Change-Id: Ic91358499c8f5eecfdd70aaf9d67ef58b9a104a8
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 golink.go      | 40 ++++++++++++++++++++++++++++++-----
 golink_test.go | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 5 deletions(-)
 create mode 100644 golink_test.go

diff --git a/golink.go b/golink.go
index 16f7d48..6857375 100644
--- a/golink.go
+++ b/golink.go
@@ -14,8 +14,8 @@ import (
 	"net/http"
 	"net/url"
 	"os"
-	"path"
 	"path/filepath"
+	"regexp"
 	"sort"
 	"strings"
 	"text/template"
@@ -177,14 +177,44 @@ func serveGo(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	dest, err := url.Parse(dl.Long)
+	target, err := expandLink(dl.Long, remainder, expandEnv{Now: time.Now().UTC()})
 	if err != nil {
-		log.Printf("error parsing URL %q: %v", dl.Long, err)
+		log.Printf("expanding %q: %v", dl.Long, err)
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
-	dest.Path = path.Join(dest.RequestURI(), remainder)
-	http.Redirect(w, r, dest.String(), http.StatusFound)
+	http.Redirect(w, r, target, http.StatusFound)
+}
+
+var reVarExpand = regexp.MustCompile(`\$\{\w+\}`)
+
+type expandEnv struct {
+	Now time.Time
+}
+
+func expandLink(long, remainder string, env expandEnv) (string, error) {
+	if strings.HasPrefix(long, "$") {
+		long = reVarExpand.ReplaceAllStringFunc(long[1:], func(m string) string {
+			switch m {
+			case "${YYYY}":
+				return env.Now.Format("2006")
+			case "${MM}":
+				return env.Now.Format("01")
+			case "${DD}":
+				return env.Now.Format("02")
+			default:
+				return m
+			}
+		})
+	}
+	_, err := url.Parse(long)
+	if err != nil {
+		return "", err
+	}
+	if remainder != "" {
+		return strings.TrimSuffix(long, "/") + "/" + remainder, nil
+	}
+	return long, nil
 }
 
 func devMode() bool { return *dev != "" }
diff --git a/golink_test.go b/golink_test.go
new file mode 100644
index 0000000..17eb9d8
--- /dev/null
+++ b/golink_test.go
@@ -0,0 +1,57 @@
+package main
+
+import (
+	"testing"
+	"time"
+)
+
+func TestExpandLink(t *testing.T) {
+	tests := []struct {
+		name      string
+		long      string
+		remainder string
+		now       time.Time
+		want      string
+	}{
+		{
+			name: "dont-mangle-escapes",
+			long: "http://host.com/foo%2f/bar",
+			want: "http://host.com/foo%2f/bar",
+		},
+		{
+			name:      "dont-mangle-escapes-and-remainder",
+			long:      "http://host.com/foo%2f/bar",
+			remainder: "extra",
+			want:      "http://host.com/foo%2f/bar/extra",
+		},
+		{
+			name:      "remainder-insert-slash",
+			long:      "http://host.com/foo",
+			remainder: "extra",
+			want:      "http://host.com/foo/extra",
+		},
+		{
+			name:      "remainder-long-as-trailing-slash",
+			long:      "http://host.com/foo/",
+			remainder: "extra",
+			want:      "http://host.com/foo/extra",
+		},
+		{
+			name: "var-expansions-time",
+			long: "$https://roamresearch.com/#/app/ts-corp/page/${MM}-${DD}-${YYYY}",
+			want: "https://roamresearch.com/#/app/ts-corp/page/06-02-2022",
+			now:  time.Date(2022, 06, 02, 1, 2, 3, 4, time.UTC),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := expandLink(tt.long, tt.remainder, expandEnv{Now: tt.now})
+			if err != nil {
+				t.Fatalf("expandLink(%q): %v", tt.long, err)
+			}
+			if got != tt.want {
+				t.Errorf("expandLink(%q) = %q; want %q", tt.long, got, tt.want)
+			}
+		})
+	}
+}