From 227da5c7b9169a63f4d4adedd4a8d56c8286d53e Mon Sep 17 00:00:00 2001
From: Jordan Wright <jmwright798@gmail.com>
Date: Sun, 10 Dec 2017 18:11:32 -0600
Subject: [PATCH] Change failed login status code to 401. Fixes #833

---
 controllers/route.go      | 25 +++++++++--------
 controllers/route_test.go | 59 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 71 insertions(+), 13 deletions(-)

diff --git a/controllers/route.go b/controllers/route.go
index 300c973f..4a9435d1 100644
--- a/controllers/route.go
+++ b/controllers/route.go
@@ -108,10 +108,7 @@ func Register(w http.ResponseWriter, r *http.Request) {
 		succ, err := auth.Register(r)
 		//If we've registered, redirect to the login page
 		if succ {
-			session.AddFlash(models.Flash{
-				Type:    "success",
-				Message: "Registration successful!.",
-			})
+			Flash(w, r, "success", "Registration successful!")
 			session.Save(r, w)
 			http.Redirect(w, r, "/login", 302)
 			return
@@ -119,10 +116,7 @@ func Register(w http.ResponseWriter, r *http.Request) {
 		// Check the error
 		m := err.Error()
 		Logger.Println(err)
-		session.AddFlash(models.Flash{
-			Type:    "danger",
-			Message: m,
-		})
+		Flash(w, r, "danger", m)
 		session.Save(r, w)
 		http.Redirect(w, r, "/register", 302)
 		return
@@ -276,18 +270,26 @@ func Login(w http.ResponseWriter, r *http.Request) {
 			http.Redirect(w, r, "/", 302)
 		} else {
 			Flash(w, r, "danger", "Invalid Username/Password")
-			http.Redirect(w, r, "/login", 302)
+			params.Flashes = session.Flashes()
+			session.Save(r, w)
+			templates := template.New("template")
+			_, err := templates.ParseFiles("templates/login.html", "templates/flashes.html")
+			if err != nil {
+				Logger.Println(err)
+			}
+			w.Header().Set("Content-Type", "text/html; charset=utf-8")
+			w.WriteHeader(http.StatusUnauthorized)
+			template.Must(templates, err).ExecuteTemplate(w, "base", params)
 		}
 	}
 }
 
 // Logout destroys the current user session
 func Logout(w http.ResponseWriter, r *http.Request) {
-	// If it is a post request, attempt to register the account
-	// Now that we are all registered, we can log the user in
 	session := ctx.Get(r, "session").(*sessions.Session)
 	delete(session.Values, "id")
 	Flash(w, r, "success", "You have successfully logged out")
+	session.Save(r, w)
 	http.Redirect(w, r, "/login", 302)
 }
 
@@ -329,5 +331,4 @@ func Flash(w http.ResponseWriter, r *http.Request, t string, m string) {
 		Type:    t,
 		Message: m,
 	})
-	session.Save(r, w)
 }
diff --git a/controllers/route_test.go b/controllers/route_test.go
index 754b59df..36c7ba42 100644
--- a/controllers/route_test.go
+++ b/controllers/route_test.go
@@ -4,6 +4,9 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"strings"
+
+	"github.com/PuerkitoBio/goquery"
 )
 
 func (s *ControllersSuite) TestLoginCSRF() {
@@ -13,6 +16,60 @@ func (s *ControllersSuite) TestLoginCSRF() {
 			"password": {"gophish"},
 		})
 
-	s.Equal(resp.StatusCode, 403)
+	s.Equal(resp.StatusCode, http.StatusForbidden)
 	fmt.Println(err)
 }
+
+func (s *ControllersSuite) TestInvalidCredentials() {
+	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{}
+	req, err := http.NewRequest("POST", fmt.Sprintf("%s/login", as.URL), strings.NewReader(url.Values{
+		"username":   {"admin"},
+		"password":   {"invalid"},
+		"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.StatusUnauthorized)
+}
+
+func (s *ControllersSuite) TestSuccessfulLogin() {
+	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{}
+	req, err := http.NewRequest("POST", fmt.Sprintf("%s/login", as.URL), 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.StatusOK)
+}