The verify view was not passing along the next param to the two factor

view.

Now if it is passed and it is a url on the same domain that request
originates from then it is used.

app/main/views/sign_in.py

@@ -4,7 +4,8 @@ from flask import (
     url_for,
     session,
     abort,
-    flash
+    flash,
+    request
 )
 
 from flask.ext.login import (current_user, login_fresh, confirm_login)
@@ -41,7 +42,10 @@ def sign_in():
                 return redirect(url_for('.verify'))
             elif user.is_active():
                 users_dao.send_verify_code(user.id, 'sms', user.mobile_number)
-                return redirect(url_for('.two_factor'))
+                if request.args.get('next'):
+                    return redirect(url_for('.two_factor', next=request.args.get('next')))
+                else:
+                    return redirect(url_for('.two_factor'))
         # Vague error message for login in case of user not known, locked, inactive or password not verified
         flash('Username or password is incorrect')
 

app/main/views/two_factor.py

@@ -3,7 +3,8 @@ from flask import (
     render_template,
     redirect,
     session,
-    url_for
+    url_for,
+    request
 )
 
 from flask_login import login_user
@@ -37,9 +38,23 @@ def two_factor():
             login_user(user, remember=True)
         finally:
             del session['user_details']
+
+        next_url = request.args.get('next')
+        if next_url and _is_safe_redirect_url(next_url):
+            return redirect(next_url)
+
         if len(services) == 1:
             return redirect(url_for('main.service_dashboard', service_id=services[0]['id']))
         else:
             return redirect(url_for('main.choose_service'))
 
     return render_template('views/two-factor.html', form=form)
+
+
+# see http://flask.pocoo.org/snippets/62/
+def _is_safe_redirect_url(target):
+    from urllib.parse import urlparse, urljoin
+    host_url = urlparse(request.host_url)
+    redirect_url = urlparse(urljoin(request.host_url, target))
+    return redirect_url.scheme in ('http', 'https') and \
+        host_url.netloc == redirect_url.netloc

tests/app/main/views/test_two_factor.py

@@ -40,6 +40,50 @@ def test_should_login_user_and_redirect_to_service_dashboard(app_,
             )
 
 
+def test_should_login_user_and_should_redirect_to_next_url(app_,
+                                                           api_user_active,
+                                                           mock_get_user,
+                                                           mock_get_user_by_email,
+                                                           mock_check_verify_code,
+                                                           mock_get_services):
+    with app_.test_request_context():
+        with app_.test_client() as client:
+            with client.session_transaction() as session:
+                session['user_details'] = {
+                    'id': api_user_active.id,
+                    'email': api_user_active.email_address}
+            response = client.post(url_for('main.two_factor', next='/services/{}/dashboard'.format(SERVICE_ONE_ID)),
+                                   data={'sms_code': '12345'})
+            assert response.status_code == 302
+            assert response.location == url_for(
+                'main.service_dashboard',
+                service_id=SERVICE_ONE_ID,
+                _external=True
+            )
+
+
+def test_should_login_user_and_not_redirect_to_external_url(app_,
+                                                            api_user_active,
+                                                            mock_get_user,
+                                                            mock_get_user_by_email,
+                                                            mock_check_verify_code,
+                                                            mock_get_services_with_one_service):
+    with app_.test_request_context():
+        with app_.test_client() as client:
+            with client.session_transaction() as session:
+                session['user_details'] = {
+                    'id': api_user_active.id,
+                    'email': api_user_active.email_address}
+            response = client.post(url_for('main.two_factor', next='http://www.google.com'),
+                                   data={'sms_code': '12345'})
+            assert response.status_code == 302
+            assert response.location == url_for(
+                'main.service_dashboard',
+                service_id=SERVICE_ONE_ID,
+                _external=True
+            )
+
+
 def test_should_login_user_and_redirect_to_choose_services(app_,
                                                            api_user_active,
                                                            mock_get_user,