internal/frontend: set cookie on alternative mod redirect

When a request to /<path> is redirected because it is an alternative
module path, a cookie is now set with the originating path. This will be
used to set the value of a banner in a future CL.

For golang/go#43693

Change-Id: I92bf674c927ecd3bd00148f893c3cb12984dce6c
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/283622
Trust: Julie Qiu <julie@golang.org>
Run-TryBot: Julie Qiu <julie@golang.org>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
diff --git a/internal/frontend/404.go b/internal/frontend/404.go
index 09ca6df..5fcf653 100644
--- a/internal/frontend/404.go
+++ b/internal/frontend/404.go
@@ -14,6 +14,7 @@
 	"github.com/google/safehtml/template"
 	"github.com/google/safehtml/template/uncheckedconversions"
 	"golang.org/x/pkgsite/internal"
+	"golang.org/x/pkgsite/internal/cookie"
 	"golang.org/x/pkgsite/internal/derrors"
 	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/log"
@@ -71,6 +72,8 @@
 	}
 	switch fr.status {
 	case http.StatusFound, derrors.ToStatus(derrors.AlternativeModule):
+		u := constructUnitURL(fr.goModPath, fr.goModPath, internal.LatestVersion)
+		cookie.Set(w, cookie.AlternativeModuleFlash, fullPath, u)
 		http.Redirect(w, r, constructUnitURL(fr.goModPath, fr.goModPath, internal.LatestVersion), http.StatusFound)
 		return
 	case http.StatusInternalServerError:
diff --git a/internal/frontend/server_test.go b/internal/frontend/server_test.go
index 06b6df6..8aa8306 100644
--- a/internal/frontend/server_test.go
+++ b/internal/frontend/server_test.go
@@ -9,6 +9,7 @@
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"os"
 	"regexp"
 	"strings"
@@ -19,6 +20,7 @@
 	"github.com/jba/templatecheck"
 	"golang.org/x/net/html"
 	"golang.org/x/pkgsite/internal"
+	"golang.org/x/pkgsite/internal/cookie"
 	"golang.org/x/pkgsite/internal/derrors"
 	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/middleware"
@@ -1230,14 +1232,14 @@
 	_, handlerNoExp, _ := newTestServer(t, nil)
 
 	for _, test := range []struct {
-		name, path string
-		wantCode   int
+		name, path, flash string
+		wantCode          int
 	}{
-		{"not found", "/invalid-page", http.StatusNotFound},
-		{"bad request", "/gocloud.dev/@latest/blob", http.StatusBadRequest},
-		{"alternative module", "/" + alternativeModule.ModulePath, http.StatusFound},
-		{"module not in v1", "/" + v1modpath, http.StatusFound},
-		{"import path not in v1", "/" + v1path, http.StatusFound},
+		{"not found", "/invalid-page", "", http.StatusNotFound},
+		{"bad request", "/gocloud.dev/@latest/blob", "", http.StatusBadRequest},
+		{"alternative module", "/" + alternativeModule.ModulePath, "module.path/alternative", http.StatusFound},
+		{"module not in v1", "/" + v1modpath, "notinv1.mod", http.StatusFound},
+		{"import path not in v1", "/" + v1path, "notinv1.mod/foo", http.StatusFound},
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			w := httptest.NewRecorder()
@@ -1246,10 +1248,23 @@
 				t.Errorf("%q: got status code = %d, want %d", test.path, w.Code, test.wantCode)
 			}
 
-			handlerNoExp.ServeHTTP(w, httptest.NewRequest("GET", test.path, nil))
+			w2 := httptest.NewRecorder()
+			handlerNoExp.ServeHTTP(w2, httptest.NewRequest("GET", test.path, nil))
 			if test.path == v1modpath || test.path == v1path {
 				test.wantCode = http.StatusNotFound
 			}
+
+			r := &http.Request{
+				Header: http.Header{"Cookie": w.Header()["Set-Cookie"]},
+				URL:    &url.URL{Path: test.path},
+			}
+			got, err := cookie.Extract(w, r, cookie.AlternativeModuleFlash)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if got != test.flash {
+				t.Fatalf("got %q; want %q", got, test.flash)
+			}
 			if w.Code != test.wantCode {
 				t.Errorf("%q: got status code = %d, want %d", test.path, w.Code, test.wantCode)
 			}