From 3d3980eb64918f660e743c12ebb7bfd83fa3ac93 Mon Sep 17 00:00:00 2001
From: Tim Gross <tgross@hashicorp.com>
Date: Fri, 26 Mar 2021 14:35:49 -0400
Subject: [PATCH] redirect from HTTP root to UI should include query params

The OTT feature relies on having a query parameter for a one-time token which
gets handled by the UI. We need to make sure that query param is preserved
when redirecting from the root URL to the `/ui/` URI.
---
 command/agent/http.go      |  6 +++++-
 command/agent/http_test.go | 17 +++++++++++++----
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/command/agent/http.go b/command/agent/http.go
index 8853a0e81..b0a3b0895 100644
--- a/command/agent/http.go
+++ b/command/agent/http.go
@@ -407,7 +407,11 @@ func (s *HTTPServer) handleUI(h http.Handler) http.Handler {
 func (s *HTTPServer) handleRootFallthrough() http.Handler {
 	return s.auditHTTPHandler(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		if req.URL.Path == "/" {
-			http.Redirect(w, req, "/ui/", 307)
+			url := "/ui/"
+			if req.URL.RawQuery != "" {
+				url = url + "?" + req.URL.RawQuery
+			}
+			http.Redirect(w, req, url, 307)
 		} else {
 			w.WriteHeader(http.StatusNotFound)
 		}
diff --git a/command/agent/http_test.go b/command/agent/http_test.go
index ee252ab97..c8ea3a732 100644
--- a/command/agent/http_test.go
+++ b/command/agent/http_test.go
@@ -74,10 +74,11 @@ func TestRootFallthrough(t *testing.T) {
 	t.Parallel()
 
 	cases := []struct {
-		desc         string
-		path         string
-		expectedPath string
-		expectedCode int
+		desc             string
+		path             string
+		expectedPath     string
+		expectedRawQuery string
+		expectedCode     int
 	}{
 		{
 			desc:         "unknown endpoint 404s",
@@ -90,6 +91,13 @@ func TestRootFallthrough(t *testing.T) {
 			expectedPath: "/ui/",
 			expectedCode: 307,
 		},
+		{
+			desc:             "root path with one-time token redirects to ui",
+			path:             "/?ott=whatever",
+			expectedPath:     "/ui/",
+			expectedRawQuery: "ott=whatever",
+			expectedCode:     307,
+		},
 	}
 
 	s := makeHTTPServer(t, nil)
@@ -115,6 +123,7 @@ func TestRootFallthrough(t *testing.T) {
 				loc, err := resp.Location()
 				require.NoError(t, err)
 				require.Equal(t, tc.expectedPath, loc.Path)
+				require.Equal(t, tc.expectedRawQuery, loc.RawQuery)
 			}
 		})
 	}