Track login next URL as part of OIDC state

mozilla_django_oidc/middleware.py

@@ -156,9 +156,7 @@ def process_request(self, request):
                 'nonce': nonce
             })
 
-        add_state_and_nonce_to_session(request, state, params)
-
-        request.session['oidc_login_next'] = request.get_full_path()
+        add_state_and_nonce_to_session(request, state, params, next=request.get_full_path())
 
         query = urlencode(params, quote_via=quote)
         redirect_url = '{url}?{query}'.format(url=auth_url, query=query)

mozilla_django_oidc/utils.py

@@ -54,7 +54,7 @@ def is_authenticated(user):
     return user.is_authenticated
 
 
-def add_state_and_nonce_to_session(request, state, params):
+def add_state_and_nonce_to_session(request, state, params, next=None):
     """
     Stores the `state` and `nonce` parameters in a session dictionary including the time when it
     was added. The dictionary can contain multiple state/nonce combinations to allow parallel
@@ -97,3 +97,6 @@ def add_state_and_nonce_to_session(request, state, params):
         'nonce': nonce,
         'added_on': time.time(),
     }
+
+    if next is not None:
+        request.session['oidc_states'][state]['next'] = next

mozilla_django_oidc/views.py

@@ -39,8 +39,7 @@ def failure_url(self):
     def success_url(self):
         # Pull the next url from the session or settings--we don't need to
         # sanitize here because it should already have been sanitized.
-        next_url = self.request.session.get('oidc_login_next', None)
-        return next_url or self.get_settings('LOGIN_REDIRECT_URL', '/')
+        return self.next_url or self.get_settings('LOGIN_REDIRECT_URL', '/')
 
     def login_failure(self):
         return HttpResponseRedirect(self.failure_url)
@@ -89,6 +88,9 @@ def get(self, request):
                 msg = 'OIDC callback state not found in session `oidc_states`!'
                 raise SuspiciousOperation(msg)
 
+            # Get the next url from the state before it is deleted
+            self.next_url = request.session['oidc_states'][state].get('next', None)
+
             # Get the nonce from the dictionary for further processing and delete the entry to
             # prevent replay attacks.
             nonce = request.session['oidc_states'][state]['nonce']
@@ -185,9 +187,12 @@ def get(self, request):
                 'nonce': nonce
             })
 
-        add_state_and_nonce_to_session(request, state, params)
-
-        request.session['oidc_login_next'] = get_next_url(request, redirect_field_name)
+        add_state_and_nonce_to_session(
+                request,
+                state,
+                params,
+                next=get_next_url(request, redirect_field_name)
+        )
 
         query = urlencode(params)
         redirect_url = '{url}?{query}'.format(url=self.OIDC_OP_AUTH_ENDPOINT, query=query)

tests/test_views.py

@@ -68,9 +68,8 @@ def test_get_auth_success_next_url(self):
         client = Client()
         request.session = client.session
         request.session['oidc_states'] = {
-            'example_state': {'nonce': None, 'added_on': time.time()},
+            'example_state': {'next': '/foobar', 'nonce': None, 'added_on': time.time()},
         }
-        request.session['oidc_login_next'] = '/foobar'
         callback_view = views.OIDCAuthenticationCallbackView.as_view()
 
         with patch('mozilla_django_oidc.views.auth.authenticate') as mock_auth:
@@ -538,22 +537,9 @@ def test_next_url(self):
         request.session = dict()
         login_view = views.OIDCAuthenticationRequestView.as_view()
         login_view(request)
-        self.assertTrue('oidc_login_next' in request.session)
-        self.assertEqual(request.session['oidc_login_next'], '/foo')
-
-    @override_settings(OIDC_OP_AUTHORIZATION_ENDPOINT='https://server.example.com/auth')
-    @override_settings(OIDC_RP_CLIENT_ID='example_id')
-    def test_missing_next_url(self):
-        """Test that `next` url gets invalidated in user session."""
-        url = reverse('oidc_authentication_init')
-        request = self.factory.get(url)
-        request.session = {
-            'oidc_login_next': 'foobar'
-        }
-        login_view = views.OIDCAuthenticationRequestView.as_view()
-        login_view(request)
-        self.assertTrue('oidc_login_next' in request.session)
-        self.assertTrue(request.session['oidc_login_next'] is None)
+        generated_state = list(request.session['oidc_states'].values())[0]
+        self.assertTrue('next' in generated_state)
+        self.assertEqual(generated_state['next'], '/foo')
 
 
 class OIDCLogoutViewTestCase(TestCase):