From 8ecbb1157822360f5bdb24231fd50f25a6247620 Mon Sep 17 00:00:00 2001
From: Miguel Grinberg <miguel.grinberg@gmail.com>
Date: Sun, 26 Apr 2020 18:51:41 +0100
Subject: [PATCH] Support an optional=True argument in login_required decorator

---
 docs/index.rst      |  8 ++++++++
 flask_httpauth.py   | 10 ++++++----
 tests/test_token.py |  9 +++++++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/docs/index.rst b/docs/index.rst
index 9c53aa0..c376217 100755
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -297,6 +297,14 @@ API Documentation
       def private_page():
           return "Only for admins!"
 
+    An optional ``optional`` argument can be set to ``True`` to allow the route to execute also when authentication is not included with the request, in which case ``auth.current_user()`` will be set to ``None``. Example::
+
+      @app.route('/private')
+      @auth.login_required(optional=True)
+      def private_page():
+          user = auth.current_user()
+          return "Hello {}!".format(user.name if user is not None else 'anonymous')
+
   .. method:: current_user()
 
     The user object returned by the ``verify_password`` callback on successful authentication. If no user is returned by the callback, this is set to the username passed by the client. Example::
diff --git a/flask_httpauth.py b/flask_httpauth.py
index 623b598..6678332 100644
--- a/flask_httpauth.py
+++ b/flask_httpauth.py
@@ -124,9 +124,11 @@ def authorize(self, role, user, auth):
             elif role in user_roles:
                 return True
 
-    def login_required(self, f=None, role=None):
-        if f is not None and role is not None:  # pragma: no cover
-            raise ValueError('role is the only supported argument')
+    def login_required(self, f=None, role=None, optional=None):
+        if f is not None and \
+                (role is not None or optional is not None):  # pragma: no cover
+            raise ValueError(
+                'role and optional are the only supported arguments')
 
         def login_required_internal(f):
             @wraps(f)
@@ -147,7 +149,7 @@ def decorated(*args, **kwargs):
                         status = 401
                     elif not self.authorize(role, user, auth):
                         status = 403
-                    if status:
+                    if not optional and status:
                         # Clear TCP receive buffer of any pending data
                         request.data
                         try:
diff --git a/tests/test_token.py b/tests/test_token.py
index 767d456..1b72762 100644
--- a/tests/test_token.py
+++ b/tests/test_token.py
@@ -36,6 +36,11 @@ def index():
         def token_auth_route():
             return 'token_auth:' + token_auth.current_user()
 
+        @app.route('/protected-optional')
+        @token_auth.login_required(optional=True)
+        def token_auth_optional_route():
+            return 'token_auth:' + str(token_auth.current_user())
+
         @app.route('/protected2')
         @token_auth2.login_required
         def token_auth_route2():
@@ -74,6 +79,10 @@ def test_token_auth_login_valid_different_case(self):
                                    'mytoken this-is-the-token!'})
         self.assertEqual(response.data.decode('utf-8'), 'token_auth:user')
 
+    def test_token_auth_login_optional(self):
+        response = self.client.get('/protected-optional')
+        self.assertEqual(response.data.decode('utf-8'), 'token_auth:None')
+
     def test_token_auth_login_invalid_token(self):
         response = self.client.get(
             '/protected', headers={'Authorization':