From 18cef9c60da2a41d18d8419cec998d740e5c7b5b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=92=9F=E9=80=8A?= <zhongxun@xsyxsc.com>
Date: Wed, 4 Sep 2019 18:39:28 +0800
Subject: [PATCH 1/4] extended custom_callback decorator, use a parameter to
 control whether or not to use custom state. I believe there are two values
 for statefield: `custom`, `destination`, correct me if I'm wrong.

---
 flask_oidc/__init__.py | 34 +++++++++++++++++++---------------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/flask_oidc/__init__.py b/flask_oidc/__init__.py
index 5ca54a4..a248ad7 100644
--- a/flask_oidc/__init__.py
+++ b/flask_oidc/__init__.py
@@ -666,21 +666,25 @@ def _is_id_token_valid(self, id_token):
 
     WRONG_GOOGLE_APPS_DOMAIN = 'WRONG_GOOGLE_APPS_DOMAIN'
 
-    def custom_callback(self, view_func):
-        """
-        Wrapper function to use a custom callback.
-        The custom OIDC callback will get the custom state field passed in with
-        redirect_to_auth_server.
-        """
-        @wraps(view_func)
-        def decorated(*args, **kwargs):
-            plainreturn, data = self._process_callback('custom')
-            if plainreturn:
-                return data
-            else:
-                return view_func(data, *args, **kwargs)
-        self._custom_callback = decorated
-        return decorated
+    def custom_callback(self, statefield='custom'):
+        def _custom_callback(view_func):
+            """
+            Wrapper function to use a custom callback.
+            The custom OIDC callback will get the custom state field passed in with
+            redirect_to_auth_server.
+            """
+
+            @wraps(view_func)
+            def decorated(*args, **kwargs):
+                plainreturn, data = self._process_callback(statefield)
+                if plainreturn:
+                    return data
+                else:
+                    return view_func(data, *args, **kwargs)
+
+            self._custom_callback = decorated
+            return decorated
+        return _custom_callback
 
     def _oidc_callback(self):
         plainreturn, data = self._process_callback('destination')

From 8f37821f2f4caac1d29870d32e89f71a4c6baef0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=92=9F=E9=80=8A?= <zhongxun@xsyxsc.com>
Date: Wed, 4 Sep 2019 18:49:57 +0800
Subject: [PATCH 2/4] update documents

---
 docs/index.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/index.rst b/docs/index.rst
index f8f8692..77c8a4e 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -99,7 +99,7 @@ Example::
         return oidc.redirect_to_auth_server(None, flask.request.values)
 
     @app.route('/custom_callback')
-    @oidc.custom_callback
+    @oidc.custom_callback('custom')
     def callback(data):
         return 'Hello. You submitted %s' % data
 

From 99a20da98952e9855de30d1c9d4b2890c29809cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=92=9F=E9=80=8A?= <zhongxun@xsyxsc.com>
Date: Wed, 4 Sep 2019 18:52:00 +0800
Subject: [PATCH 3/4] remove redundant blank lines

---
 flask_oidc/__init__.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/flask_oidc/__init__.py b/flask_oidc/__init__.py
index a248ad7..46be703 100644
--- a/flask_oidc/__init__.py
+++ b/flask_oidc/__init__.py
@@ -673,7 +673,6 @@ def _custom_callback(view_func):
             The custom OIDC callback will get the custom state field passed in with
             redirect_to_auth_server.
             """
-
             @wraps(view_func)
             def decorated(*args, **kwargs):
                 plainreturn, data = self._process_callback(statefield)
@@ -681,7 +680,6 @@ def decorated(*args, **kwargs):
                     return data
                 else:
                     return view_func(data, *args, **kwargs)
-
             self._custom_callback = decorated
             return decorated
         return _custom_callback

From 956e125544638884abd0926d062a73fe3fd75e92 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=92=9F=E9=80=8A?= <zhongxun@xsyxsc.com>
Date: Wed, 4 Sep 2019 18:56:56 +0800
Subject: [PATCH 4/4] add explanation of custom_callback parameter.

---
 docs/index.rst | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/docs/index.rst b/docs/index.rst
index 77c8a4e..953c226 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -91,6 +91,8 @@ To do this, add the decorator `oidc.custom_callback` to your callback function.
 This will get the (json-serializable) custom state that you passed in as
 `customstate` to `oidc.redirect_to_auth_server`.
 Note that to use this, you will need to set `OVERWRITE_REDIRECT_URI`.
+Use parameter `custom` if you want to use custom state, `destination` if you
+want to keep default state settings.
 
 Example::