gddo-server: add UsePkgGoDev field to gddoEvent

This boolean field will reflect whether the godoc.org request was
redirected to pkg.go.dev.

Change-Id: I521c98e7342f5880021bd215640adb570935cb06
Reviewed-on: https://go-review.googlesource.com/c/gddo/+/235957
Reviewed-by: Jonathan Amsterdam <jba@google.com>
diff --git a/gddo-server/main.go b/gddo-server/main.go
index f02c539..d316344 100644
--- a/gddo-server/main.go
+++ b/gddo-server/main.go
@@ -1070,6 +1070,7 @@
 	RedirectHost string
 	Latency      time.Duration
 	IsRobot      bool
+	UsePkgGoDev  bool
 }
 
 func newGDDOEvent(r *http.Request, latency time.Duration, isRobot bool) *gddoEvent {
@@ -1082,6 +1083,7 @@
 		RedirectHost: pkgGoDevURL.String(),
 		Latency:      latency,
 		IsRobot:      isRobot,
+		UsePkgGoDev:  shouldRedirectToPkgGoDev(r),
 	}
 }
 
@@ -1114,6 +1116,19 @@
 	return req.FormValue("utm_source") == "backtogodoc"
 }
 
+func shouldRedirectToPkgGoDev(req *http.Request) bool {
+	// API requests are not redirected.
+	if strings.HasPrefix(req.URL.Host, "api") {
+		return false
+	}
+	redirectParam := req.FormValue(pkgGoDevRedirectParam)
+	if redirectParam == pkgGoDevRedirectOn || redirectParam == pkgGoDevRedirectOff {
+		return redirectParam == pkgGoDevRedirectOn
+	}
+	cookie, err := req.Cookie(pkgGoDevRedirectCookie)
+	return (err == nil && cookie.Value == pkgGoDevRedirectOn)
+}
+
 // pkgGoDevRedirectHandler redirects requests from godoc.org to pkg.go.dev,
 // based on whether a cookie is set for pkggodev-redirect. The cookie
 // can be turned on/off using a query param.
@@ -1134,19 +1149,7 @@
 			http.SetCookie(w, cookie)
 		}
 
-		var shouldRedirect bool
-		if redirectParam == pkgGoDevRedirectOn || redirectParam == pkgGoDevRedirectOff {
-			shouldRedirect = redirectParam == pkgGoDevRedirectOn
-		} else {
-			for _, v := range r.Cookies() {
-				if v.Name == pkgGoDevRedirectCookie {
-					shouldRedirect = v.Value == pkgGoDevRedirectOn
-					break
-				}
-			}
-		}
-
-		if !shouldRedirect {
+		if !shouldRedirectToPkgGoDev(r) {
 			return f(w, r)
 		}
 
diff --git a/gddo-server/main_test.go b/gddo-server/main_test.go
index 1624306..1f29fa7 100644
--- a/gddo-server/main_test.go
+++ b/gddo-server/main_test.go
@@ -166,53 +166,139 @@
 
 func TestNewGDDOEvent(t *testing.T) {
 	for _, test := range []struct {
-		url  string
-		want *gddoEvent
+		name   string
+		url    string
+		cookie *http.Cookie
+		want   *gddoEvent
 	}{
 		{
-			url: "https://godoc.org",
+			name: "home page request",
+			url:  "https://godoc.org",
 			want: &gddoEvent{
-				Host: "godoc.org",
-				Path: "",
+				Host:        "godoc.org",
+				Path:        "",
+				UsePkgGoDev: false,
 			},
 		},
 		{
-			url: "https://godoc.org/-/about",
+			name:   "home page request with cookie on should redirect",
+			url:    "https://godoc.org",
+			cookie: &http.Cookie{Name: "pkggodev-redirect", Value: "on"},
 			want: &gddoEvent{
-				Host: "godoc.org",
-				Path: "/-/about",
+				Host:        "godoc.org",
+				Path:        "",
+				UsePkgGoDev: true,
 			},
 		},
 		{
-			url: "https://godoc.org/?q=test",
+			name: "about page request",
+			url:  "https://godoc.org/-/about",
 			want: &gddoEvent{
-				Host: "godoc.org",
-				Path: "/",
+				Host:        "godoc.org",
+				Path:        "/-/about",
+				UsePkgGoDev: false,
 			},
 		},
 		{
-			url: "https://godoc.org/net/http",
+			name: "request with search query parameter",
+			url:  "https://godoc.org/?q=test",
 			want: &gddoEvent{
-				Host: "godoc.org",
-				Path: "/net/http",
+				Host:        "godoc.org",
+				Path:        "/",
+				UsePkgGoDev: false,
 			},
 		},
 		{
-			url: "https://api.godoc.org/imports/net/http",
+			name: "package page request",
+			url:  "https://godoc.org/net/http",
 			want: &gddoEvent{
-				Host: "api.godoc.org",
-				Path: "/imports/net/http",
+				Host:        "godoc.org",
+				Path:        "/net/http",
+				UsePkgGoDev: false,
+			},
+		},
+		{
+			name:   "package page request with wrong cookie on should not redirect",
+			url:    "https://godoc.org/net/http",
+			cookie: &http.Cookie{Name: "bogus-cookie", Value: "on"},
+			want: &gddoEvent{
+				Host:        "godoc.org",
+				Path:        "/net/http",
+				UsePkgGoDev: false,
+			},
+		},
+		{
+			name:   "package page request with query parameter off should not redirect",
+			url:    "https://godoc.org/net/http?redirect=off",
+			cookie: &http.Cookie{Name: "pkggodev-redirect", Value: "on"},
+			want: &gddoEvent{
+				Host:        "godoc.org",
+				Path:        "/net/http",
+				UsePkgGoDev: false,
+			},
+		},
+		{
+			name:   "package page request with cookie on should redirect",
+			url:    "https://godoc.org/net/http",
+			cookie: &http.Cookie{Name: "pkggodev-redirect", Value: "on"},
+			want: &gddoEvent{
+				Host:        "godoc.org",
+				Path:        "/net/http",
+				UsePkgGoDev: true,
+			},
+		},
+		{
+			name:   "package page request with query parameter on should redirect",
+			url:    "https://godoc.org/net/http?redirect=on",
+			cookie: &http.Cookie{Name: "pkggodev-redirect", Value: ""},
+			want: &gddoEvent{
+				Host:        "godoc.org",
+				Path:        "/net/http",
+				UsePkgGoDev: true,
+			},
+		},
+		{
+			name: "api request",
+			url:  "https://api.godoc.org/imports/net/http",
+			want: &gddoEvent{
+				Host:        "api.godoc.org",
+				Path:        "/imports/net/http",
+				UsePkgGoDev: false,
+			},
+		},
+		{
+			name:   "api requests should never redirect, even with cookie on",
+			url:    "https://api.godoc.org/imports/net/http",
+			cookie: &http.Cookie{Name: "pkggodev-redirect", Value: "on"},
+			want: &gddoEvent{
+				Host:        "api.godoc.org",
+				Path:        "/imports/net/http",
+				UsePkgGoDev: false,
+			},
+		},
+		{
+			name:   "api requests should never redirect, even with query parameter on",
+			url:    "https://api.godoc.org/imports/net/http?redirect=on",
+			cookie: &http.Cookie{Name: "pkggodev-redirect", Value: ""},
+			want: &gddoEvent{
+				Host:        "api.godoc.org",
+				Path:        "/imports/net/http",
+				UsePkgGoDev: false,
 			},
 		},
 	} {
-		t.Run(test.url, func(t *testing.T) {
-			r := httptest.NewRequest("GET", test.url, nil)
+		t.Run(test.name, func(t *testing.T) {
 			want := test.want
 			want.Latency = 100
 			want.RedirectHost = "https://" + pkgGoDevHost
 			want.URL = test.url
 			want.Header = http.Header{}
 			want.IsRobot = true
+			r := httptest.NewRequest("GET", test.url, nil)
+			if test.cookie != nil {
+				r.AddCookie(test.cookie)
+				want.Header.Add("Cookie", test.cookie.String())
+			}
 			got := newGDDOEvent(r, want.Latency, want.IsRobot)
 			if diff := cmp.Diff(want, got); diff != "" {
 				t.Fatalf("mismatch (-want +got):\n%s", diff)