internal/cookie: add logic for flashCookie
Logic for getting and setting a flash cookie is added. This will be used
for the redirected from banner.
For golang/go#43693
Change-Id: I7d07eb1ee4beeefe33e9b20a10ce3581504ade57
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/283619
Trust: Julie Qiu <julie@golang.org>
Run-TryBot: Julie Qiu <julie@golang.org>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
diff --git a/internal/cookie/flash.go b/internal/cookie/flash.go
new file mode 100644
index 0000000..b45c64c
--- /dev/null
+++ b/internal/cookie/flash.go
@@ -0,0 +1,47 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package cookie is used to get and set HTTP cookies.
+package cookie
+
+import (
+ "encoding/base64"
+ "fmt"
+ "net/http"
+ "time"
+
+ "golang.org/x/pkgsite/internal/derrors"
+)
+
+// AlternativeModuleFlash indicates the alternative module path that
+// a request was redirected from.
+const AlternativeModuleFlash = "tmp-redirected-from-alternative-module"
+
+// Extract returns the value of the cookie at name and deletes the cookie.
+func Extract(w http.ResponseWriter, r *http.Request, name string) (_ string, err error) {
+ defer derrors.Wrap(&err, "Extract")
+ c, err := r.Cookie(name)
+ if err != nil && err != http.ErrNoCookie {
+ return "", fmt.Errorf("r.Cookie(%q): %v", name, err)
+ }
+ if c == nil {
+ return "", nil
+ }
+ val, err := base64.URLEncoding.DecodeString(c.Value)
+ if err != nil {
+ return "", nil
+ }
+ http.SetCookie(w, &http.Cookie{
+ Name: name,
+ Path: r.URL.Path,
+ Expires: time.Unix(0, 0),
+ })
+ return string(val), nil
+}
+
+// Set sets a cookie at the urlPath with name and val.
+func Set(w http.ResponseWriter, name, val, urlPath string) {
+ value := base64.URLEncoding.EncodeToString([]byte(val))
+ http.SetCookie(w, &http.Cookie{Name: name, Value: value, Path: urlPath})
+}
diff --git a/internal/cookie/flash_test.go b/internal/cookie/flash_test.go
new file mode 100644
index 0000000..61ced7a
--- /dev/null
+++ b/internal/cookie/flash_test.go
@@ -0,0 +1,35 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package cookie
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "net/url"
+ "testing"
+)
+
+const (
+ testName = "testname"
+ testVal = "testvalue"
+)
+
+func TestFlashMessages(t *testing.T) {
+ w := httptest.NewRecorder()
+
+ Set(w, testName, testVal, "/foo")
+ r := &http.Request{
+ Header: http.Header{"Cookie": w.Header()["Set-Cookie"]},
+ URL: &url.URL{Path: "/foo"},
+ }
+
+ got, err := Extract(w, r, testName)
+ if err != nil {
+ t.Fatal(err)
+ }
+ if got != testVal {
+ t.Errorf("got %q, want %q", got, testVal)
+ }
+}
