1 files changed, 12 insertions, 2 deletions
diff --git a/pkg/handler/auth/login.go b/pkg/handler/auth/login.go
index 7014548..89fd87b 100644
--- a/pkg/handler/auth/login.go
+++ b/pkg/handler/auth/login.go
@@ -45,10 +45,19 @@ func (g *LoginHandler) Logout(w http.ResponseWriter, r *http.Request) error {
 }
 
 func (g *LoginHandler) Login(w http.ResponseWriter, r *http.Request) error {
+	referer := r.URL.Query().Get("referer")
+
+	// if query value is empty tries to get from header
+	if referer == "" {
+		referer = r.Header.Get("Referer")
+	}
+
 	if r.Method == "GET" {
 		ext.SetHTML(w)
 
-		login := &templates.LoginPage{}
+		login := &templates.LoginPage{
+			Referer: referer,
+		}
 		templates.WritePageTemplate(w, login, r.Context())
 	} else if r.Method == "POST" {
 
@@ -57,6 +66,7 @@ func (g *LoginHandler) Login(w http.ResponseWriter, r *http.Request) error {
 
 		if !g.auth.CheckAuth(username, password) {
 			login := &templates.LoginPage{
+				Referer:      referer,
 				ErrorMessage: "Invalid login",
 			}
 			templates.WritePageTemplate(w, login, r.Context())
@@ -78,7 +88,7 @@ func (g *LoginHandler) Login(w http.ResponseWriter, r *http.Request) error {
 			}
 
 			http.SetCookie(w, cookie)
-			ext.Redirect(w, "/")
+			ext.Redirect(w, referer)
 		}
 
 	}