From aa8c770e7303af495e75d47fa0076460760acafc Mon Sep 17 00:00:00 2001
From: Jordan Wright <jmwright798@gmail.com>
Date: Sun, 10 Dec 2017 21:40:46 -0600
Subject: [PATCH] Adding "next" parameter to support redirecting after
 successful login.

---
 controllers/route.go      | 11 ++++++++++-
 controllers/route_test.go | 35 +++++++++++++++++++++++++++++++++++
 middleware/middleware.go  |  4 +++-
 templates/login.html      |  6 +++---
 4 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/controllers/route.go b/controllers/route.go
index 4a9435d1..b7a3bd98 100644
--- a/controllers/route.go
+++ b/controllers/route.go
@@ -5,6 +5,7 @@ import (
 	"html/template"
 	"log"
 	"net/http"
+	"net/url"
 	"os"
 
 	"github.com/gophish/gophish/auth"
@@ -267,7 +268,15 @@ func Login(w http.ResponseWriter, r *http.Request) {
 		if succ {
 			session.Values["id"] = u.Id
 			session.Save(r, w)
-			http.Redirect(w, r, "/", 302)
+			next := "/"
+			url, err := url.Parse(r.FormValue("next"))
+			if err == nil {
+				path := url.Path
+				if path != "" {
+					next = path
+				}
+			}
+			http.Redirect(w, r, next, 302)
 		} else {
 			Flash(w, r, "danger", "Invalid Username/Password")
 			params.Flashes = session.Flashes()
diff --git a/controllers/route_test.go b/controllers/route_test.go
index 36c7ba42..09d7ec45 100644
--- a/controllers/route_test.go
+++ b/controllers/route_test.go
@@ -73,3 +73,38 @@ func (s *ControllersSuite) TestSuccessfulLogin() {
 	s.Equal(err, nil)
 	s.Equal(resp.StatusCode, http.StatusOK)
 }
+
+func (s *ControllersSuite) TestSuccessfulRedirect() {
+	next := "/campaigns"
+	resp, err := http.Get(fmt.Sprintf("%s/login", as.URL))
+	s.Equal(err, nil)
+	s.Equal(resp.StatusCode, http.StatusOK)
+
+	doc, err := goquery.NewDocumentFromResponse(resp)
+	s.Equal(err, nil)
+	elem := doc.Find("input[name='csrf_token']").First()
+	token, ok := elem.Attr("value")
+	s.Equal(ok, true)
+
+	client := &http.Client{
+		CheckRedirect: func(req *http.Request, via []*http.Request) error {
+			return http.ErrUseLastResponse
+		},
+	}
+	req, err := http.NewRequest("POST", fmt.Sprintf("%s/login?next=%s", as.URL, next), strings.NewReader(url.Values{
+		"username":   {"admin"},
+		"password":   {"gophish"},
+		"csrf_token": {token},
+	}.Encode()))
+	s.Equal(err, nil)
+
+	req.Header.Set("Cookie", resp.Header.Get("Set-Cookie"))
+	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
+
+	resp, err = client.Do(req)
+	s.Equal(err, nil)
+	s.Equal(resp.StatusCode, http.StatusFound)
+	url, err := resp.Location()
+	s.Equal(err, nil)
+	s.Equal(url.Path, next)
+}
diff --git a/middleware/middleware.go b/middleware/middleware.go
index 5756caa1..ef7a0d99 100644
--- a/middleware/middleware.go
+++ b/middleware/middleware.go
@@ -94,7 +94,9 @@ func RequireLogin(handler http.Handler) http.HandlerFunc {
 		if u := ctx.Get(r, "user"); u != nil {
 			handler.ServeHTTP(w, r)
 		} else {
-			http.Redirect(w, r, "/login", 302)
+			q := r.URL.Query()
+			q.Set("next", r.URL.Path)
+			http.Redirect(w, r, fmt.Sprintf("/login?%s", q.Encode()), 302)
 		}
 	}
 }
diff --git a/templates/login.html b/templates/login.html
index 4dca3ce7..a4cb58b2 100644
--- a/templates/login.html
+++ b/templates/login.html
@@ -41,13 +41,13 @@
         </div>
     </div>
     <div class="container">
-        <form class="form-signin" action="/login" method="POST">
+        <form class="form-signin" action="" method="POST">
             <img id="logo" src="/images/logo_purple.png" />
             <h2 class="form-signin-heading">Please sign in</h2>
             {{template "flashes" .Flashes}}
             <input type="text" name="username" class="form-control top-input" placeholder="Username" required autofocus>
             <input type="password" name="password" class="form-control bottom-input" placeholder="Password" autocomplete="off" required>
-            <input type="hidden" name="csrf_token" value="{{.Token}}"/>
+            <input type="hidden" name="csrf_token" value="{{.Token}}" />
             <button class="btn btn-lg btn-primary btn-block" type="submit">Sign in</button>
         </form>
     </div>
@@ -56,4 +56,4 @@
 </body>
 
 </html>
-{{ end }}
+{{ end }}
\ No newline at end of file