internal/short: allow relative paths from the shortend URLs

golang.org/s/key URL shortener supported only simple redirects
from key to a link. This change allows the short link to have
extra path elements, and computes the redirects by appending
the extra path elements to the resolved link. For example,
if golang.org/s/foo is configured to be resolved to example.com,
golang.org/s/foo/bar will result in a redirects to
example.com/bar.

Change-Id: I6aa9f4aab2d5a74c76fda446a29aae998fe48ad6
Reviewed-on: https://go-review.googlesource.com/c/website/+/227654
Run-TryBot: Andrew Bonventre <andybons@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Andrew Bonventre <andybons@golang.org>

internal/short/short.go

diff --git a/internal/short/short.go b/internal/short/short.go
index c04e789..0b81ef7 100644
--- a/internal/short/short.go
+++ b/internal/short/short.go
@@ -21,6 +21,7 @@
 	"net/http"
 	"net/url"
 	"regexp"
+	"strings"
 
 	"cloud.google.com/go/datastore"
 	"golang.org/x/website/internal/memcache"
@@ -61,14 +62,16 @@
 }
 
 // linkHandler services requests to short URLs.
-//   http://golang.org/s/key
+//   http://golang.org/s/key[/remaining/path]
 // It consults memcache and datastore for the Link for key.
 // It then sends a redirects or an error message.
+// If the remaining path part is not empty, the redirects
+// will be the relative path from the resolved Link.
 func (h server) linkHandler(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 
-	key := r.URL.Path[len(prefix)+1:]
-	if !validKey.MatchString(key) {
+	key, remainingPath, err := extractKey(r)
+	if err != nil { // invalid key or url
 		http.Error(w, "not found", http.StatusNotFound)
 		return
 	}
@@ -96,7 +99,28 @@
 		}
 	}
 
-	http.Redirect(w, r, link.Target, http.StatusFound)
+	target := link.Target
+	if remainingPath != "" {
+		target += remainingPath
+	}
+	http.Redirect(w, r, target, http.StatusFound)
+}
+
+func extractKey(r *http.Request) (key, remainingPath string, err error) {
+	path := r.URL.Path
+	if !strings.HasPrefix(path, prefix+"/") {
+		return "", "", errors.New("invalid path")
+	}
+
+	key, remainingPath = path[len(prefix)+1:], ""
+	if slash := strings.Index(key, "/"); slash > 0 {
+		key, remainingPath = key[:slash], key[slash:]
+	}
+
+	if !validKey.MatchString(key) {
+		return "", "", errors.New("invalid key")
+	}
+	return key, remainingPath, nil
 }
 
 var adminTemplate = template.Must(template.New("admin").Parse(templateHTML))

internal/short/short_test.go

diff --git a/internal/short/short_test.go b/internal/short/short_test.go
new file mode 100644
index 0000000..5ae4443
--- /dev/null
+++ b/internal/short/short_test.go
@@ -0,0 +1,39 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by the Apache 2.0
+// license that can be found in the LICENSE file.
+
+// +build golangorg
+
+package short
+
+import (
+	"net/http/httptest"
+	"testing"
+)
+
+func TestExtractKey(t *testing.T) {
+	testCases := []struct {
+		in                     string
+		wantKey, wantRemaining string
+		wantErr                bool
+	}{
+		{in: "/s/foo", wantKey: "foo", wantRemaining: ""},
+		{in: "/s/foo/", wantKey: "foo", wantRemaining: "/"},
+		{in: "/s/foo/bar/", wantKey: "foo", wantRemaining: "/bar/"},
+		{in: "/s/foo.bar/baz", wantKey: "foo.bar", wantRemaining: "/baz"},
+		{in: "/s/s/s/s", wantKey: "s", wantRemaining: "/s/s"},
+		{in: "/", wantErr: true},
+		{in: "/s/", wantErr: true},
+		{in: "/s", wantErr: true},
+		{in: "/t/foo", wantErr: true},
+		{in: "/s/foo*", wantErr: true},
+	}
+
+	for _, tc := range testCases {
+		req := httptest.NewRequest("GET", tc.in, nil)
+		gotKey, gotRemaining, gotErr := extractKey(req)
+		if gotKey != tc.wantKey || gotRemaining != tc.wantRemaining || (gotErr != nil) != tc.wantErr {
+			t.Errorf("extractKey(%q) = (%q, %q, %v), want (%q, %q, err=%v)", tc.in, gotKey, gotRemaining, gotErr, tc.wantKey, tc.wantRemaining, tc.wantErr)
+		}
+	}
+}