From f78966682abb1aa05106ea1f60a4c20a83d496af Mon Sep 17 00:00:00 2001
From: Rasmus Kaj <kaj@kth.se>
Date: Thu, 3 Aug 2017 12:50:01 +0200
Subject: [PATCH] Support "next" parameter in login form.

---
 Cargo.toml                  |  1 +
 src/main.rs                 | 56 ++++++++++++++++++++++++++++++++-----
 templates/login.rs.html     |  8 ++++--
 templates/not_found.rs.html | 10 ++++---
 4 files changed, 62 insertions(+), 13 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index b907d43..fba5741 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -41,3 +41,4 @@ memcached-rs = "^0.1"
 flate2 = "*"
 brotli2 = "*"
 mime = "0.2.6"
+regex = "*"
diff --git a/src/main.rs b/src/main.rs
index b32d30a..ab164f0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,14 +18,16 @@ extern crate diesel;
 extern crate r2d2_diesel;
 extern crate dotenv;
 extern crate memcached;
+extern crate regex;
 
 use chrono::Datelike;
 use diesel::pg::PgConnection;
 use diesel::prelude::*;
 use dotenv::dotenv;
 use hyper::header::ContentType;
-use nickel::{Action, Continue, FormBody, Halt, HttpRouter, MediaType, MiddlewareResult, Nickel,
-             NickelError, Request, Response};
+use nickel::{Action, Continue, FormBody, Halt, HttpRouter, MediaType,
+             MiddlewareResult, Nickel, NickelError, QueryString, Request,
+             Response};
 use nickel::extensions::response::Redirect;
 use nickel::status::StatusCode::NotFound;
 use nickel_diesel::{DieselMiddleware, DieselRequestExtensions};
@@ -118,7 +120,7 @@ fn main() {
 fn custom_errors(err: &mut NickelError, req: &mut Request) -> Action {
     if let Some(ref mut res) = err.stream {
         if res.status() == NotFound {
-            templates::not_found(res, req.authorized_user()).unwrap();
+            templates::not_found(res, req).unwrap();
             return Halt(());
         }
     }
@@ -126,11 +128,12 @@ fn custom_errors(err: &mut NickelError, req: &mut Request) -> Action {
     Continue(())
 }
 
-fn login<'mw>(_req: &mut Request,
+fn login<'mw>(req: &mut Request,
               mut res: Response<'mw>)
               -> MiddlewareResult<'mw> {
     res.clear_jwt();
-    res.ok(|o| templates::login(o))
+    let next = sanitize_next(req.query().get("next"));
+    res.ok(|o| templates::login(o, next))
 }
 
 fn do_login<'mw>(req: &mut Request,
@@ -138,6 +141,7 @@ fn do_login<'mw>(req: &mut Request,
                  -> MiddlewareResult<'mw> {
     let c: &PgConnection = &req.db_conn();
     let form_data = try_with!(res, req.form_body());
+    let next = sanitize_next(form_data.get("next"));
     if let (Some(user), Some(pw)) = (form_data.get("user"),
                                      form_data.get("password")) {
         use rphotos::schema::users::dsl::*;
@@ -148,14 +152,52 @@ fn do_login<'mw>(req: &mut Request,
             if djangohashers::check_password_tolerant(pw, &hash) {
                 info!("User {} logged in", user);
                 res.set_jwt_user(user);
-                return res.redirect("/");
+                return res.redirect(next.unwrap_or("/"));
             }
             debug!("Password verification failed");
         } else {
             debug!("No hash found for {}", user);
         }
     }
-    res.ok(|o| templates::login(o))
+    res.ok(|o| templates::login(o, next))
+}
+
+fn sanitize_next(next: Option<&str>) -> Option<&str> {
+    if let Some(next) = next {
+        use regex::Regex;
+        let re = Regex::new(r"^/([a-z0-9.-]+/?)*$").unwrap();
+        if re.is_match(next) {
+            return Some(next)
+        }
+    }
+    None
+}
+
+#[test]
+fn test_sanitize_bad_1() {
+    assert_eq!(None, sanitize_next(Some("https://evil.org/")))
+}
+
+#[test]
+fn test_sanitize_bad_2() {
+    assert_eq!(None, sanitize_next(Some("//evil.org/")))
+}
+#[test]
+fn test_sanitize_bad_3() {
+    assert_eq!(None, sanitize_next(Some("/evil\"hack")))
+}
+#[test]
+fn test_sanitize_bad_4() {
+    assert_eq!(None, sanitize_next(Some("/evil'hack")))
+}
+
+#[test]
+fn test_sanitize_good_1() {
+    assert_eq!(Some("/foo/"), sanitize_next(Some("/foo/")))
+}
+#[test]
+fn test_sanitize_good_2() {
+    assert_eq!(Some("/2017/7/15"), sanitize_next(Some("/2017/7/15")))
 }
 
 fn logout<'mw>(_req: &mut Request,
diff --git a/templates/login.rs.html b/templates/login.rs.html
index 4c286f8..6844211 100644
--- a/templates/login.rs.html
+++ b/templates/login.rs.html
@@ -1,6 +1,6 @@
 @use templates::page_base;
 
-@()
+@(next: Option<&str>)
 
 @:page_base("login", &[], &None, {
     <form action="/login" method="post">
@@ -8,6 +8,10 @@
 	<input id="user" name="user"></p>
       <p><label for="password">Password:</label>
 	<input id="password" name="password" type="password"></p>
-      <p><span></span><input type="submit" value="Log in"></p>
+      <p><span>@if let Some(next) = next {
+	  <input type="hidden" name="next" value="@next">
+	  }</span>
+	<input type="submit" value="Log in">
+      </p>
     </form>
 })
diff --git a/templates/not_found.rs.html b/templates/not_found.rs.html
index 16078ce..db3ef21 100644
--- a/templates/not_found.rs.html
+++ b/templates/not_found.rs.html
@@ -1,11 +1,13 @@
+@use nickel::Request;
 @use templates::page_base;
+@use nickel_jwt_session::SessionRequestExtensions;
 
-@(user: Option<String>)
+@(req: &Request)
 
-@:page_base("Not found", &[], &user, {
+@:page_base("Not found", &[], &req.authorized_user(), {
   <p>No page or photo match that url.</p>
-  @if user.is_none() {
+  @if req.authorized_user().is_none() {
     <p>At least nothing publicly visible, you might try
-    <a href="/login">logging in</a>.</p>
+    <a href="/login@if let Some(p) = req.path_without_query() {?next=@p}">logging in</a>.</p>
   }
 })