From bf260d4b7e0908ec7153cb2066976080111a8066 Mon Sep 17 00:00:00 2001
From: Matt Wright <matt@nobien.net>
Date: Tue, 28 May 2013 11:01:25 -0400
Subject: [PATCH] Add optional `next` parameter to registration endpoint. Fixes
 #117.

---
 CHANGES                    | 14 +++++++++++---
 flask_security/utils.py    | 13 ++++++++++---
 flask_security/views.py    | 19 ++++++++-----------
 tests/configured_tests.py  |  8 ++++++++
 tests/test_app/__init__.py |  4 ++++
 5 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/CHANGES b/CHANGES
index 3a4c12e..7cf2674 100644
--- a/CHANGES
+++ b/CHANGES
@@ -4,6 +4,14 @@ Flask-Security Changelog
 Here you can see the full list of changes between each Flask-Security release.
 
 
+Version 1.6.4
+-------------
+
+Not yet released
+
+- Added optional `next` parameter to registration endpoint, similar to that of login
+
+
 Version 1.6.3
 -------------
 
@@ -117,10 +125,10 @@ Version 1.2.0
 
 Released March 12th 2012
 
-- Added configuration option `SECURITY_FLASH_MESSAGES` which can be set to a 
+- Added configuration option `SECURITY_FLASH_MESSAGES` which can be set to a
   boolean value to specify if Flask-Security should flash messages or not.
-  
+
 Version 1.1.0
 -------------
 
-Initial release
\ No newline at end of file
+Initial release
diff --git a/flask_security/utils.py b/flask_security/utils.py
index d4a0334..537e4cf 100644
--- a/flask_security/utils.py
+++ b/flask_security/utils.py
@@ -144,11 +144,18 @@ def url_for_security(endpoint, **values):
     return url_for(endpoint, **values)
 
 
-def get_post_login_redirect():
-    """Returns the URL to redirect to after a user logs in successfully."""
+def get_post_action_redirect(config_key):
     return (get_url(request.args.get('next')) or
             get_url(request.form.get('next')) or
-            find_redirect('SECURITY_POST_LOGIN_VIEW'))
+            find_redirect(config_key))
+
+
+def get_post_login_redirect():
+    return get_post_action_redirect('SECURITY_POST_LOGIN_VIEW')
+
+
+def get_post_register_redirect():
+    return get_post_action_redirect('SECURITY_POST_REGISTER_VIEW')
 
 
 def find_redirect(key):
diff --git a/flask_security/views.py b/flask_security/views.py
index 1b8488d..89c5342 100644
--- a/flask_security/views.py
+++ b/flask_security/views.py
@@ -10,24 +10,23 @@
 """
 
 from flask import current_app, redirect, request, render_template, jsonify, \
-     after_this_request, Blueprint
+    after_this_request, Blueprint
 from flask_login import current_user
 from werkzeug.datastructures import MultiDict
 from werkzeug.local import LocalProxy
 
 from .confirmable import send_confirmation_instructions, \
-     confirm_user, confirm_email_token_status
+    confirm_user, confirm_email_token_status
 from .decorators import login_required, anonymous_user_required
 from .passwordless import send_login_instructions, \
-     login_token_status
+    login_token_status
 from .recoverable import reset_password_token_status, \
-     send_reset_password_instructions, update_password
+    send_reset_password_instructions, update_password
 from .changeable import change_user_password
 from .registerable import register_user
-from .utils import get_url, get_post_login_redirect, do_flash, \
-     get_message, login_user, logout_user, url_for_security as url_for, \
-     config_value
-
+from .utils import config_value, do_flash, get_url, get_post_login_redirect, \
+    get_post_register_redirect, get_message, login_user, logout_user, \
+    url_for_security as url_for
 
 # Convenient references
 _security = LocalProxy(lambda: current_app.extensions['security'])
@@ -123,9 +122,7 @@ def register():
             login_user(user)
 
         if not request.json:
-            post_register_url = get_url(_security.post_register_view)
-            post_login_url = get_url(_security.post_login_view)
-            return redirect(post_register_url or post_login_url)
+            return redirect(get_post_register_redirect())
 
     if request.json:
         return _render_json(form)
diff --git a/tests/configured_tests.py b/tests/configured_tests.py
index 26ec9e9..b8ae248 100644
--- a/tests/configured_tests.py
+++ b/tests/configured_tests.py
@@ -68,6 +68,14 @@ class ConfiguredSecurityTests(SecurityTest):
         r = self._post('/register', data=data, follow_redirects=True)
         self.assertIn('Post Register', r.data)
 
+    def test_register_with_next_querystring_argument(self):
+        data = dict(email='dude@lp.com',
+                    password='password',
+                    password_confirm='password')
+
+        r = self._post('/register?next=/page1', data=data, follow_redirects=True)
+        self.assertIn('Page 1', r.data)
+
     def test_register_json(self):
         data = '{ "email": "dude@lp.com", "password": "password", "csrf_token":"%s" }' % self.csrf_token
         r = self._post('/register', data=data, content_type='application/json')
diff --git a/tests/test_app/__init__.py b/tests/test_app/__init__.py
index ebf550c..53b3de0 100644
--- a/tests/test_app/__init__.py
+++ b/tests/test_app/__init__.py
@@ -113,6 +113,10 @@ def create_app(config):
     def invalid_role():
         return 'success' if ds.find_role('bogus') is None else 'failure'
 
+    @app.route('/page1')
+    def page_1():
+        return 'Page 1'
+
     return app