From 93fcb0c4926f36a82ccf1f09b5dbe9b3ae77a2d8 Mon Sep 17 00:00:00 2001
From: Jed Cunningham <jedcunningham@apache.org>
Date: Tue, 16 Aug 2022 21:01:41 -0600
Subject: [PATCH] Rotate session id during login

---
 airflow/www/fab_security/manager.py | 15 +++++++++++++++
 tests/www/views/test_session.py     | 28 ++++++++++++++++++++++++++--
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/airflow/www/fab_security/manager.py b/airflow/www/fab_security/manager.py
index 8399f11df367e..3317fb014b489 100644
--- a/airflow/www/fab_security/manager.py
+++ b/airflow/www/fab_security/manager.py
@@ -25,6 +25,7 @@
 import logging
 import re
 from typing import Any, Dict, List, Optional, Set, Tuple, Type
+from uuid import uuid4
 
 from flask import current_app, g, session, url_for
 from flask_appbuilder import AppBuilder
@@ -70,6 +71,7 @@
 from flask_login import AnonymousUserMixin, LoginManager, current_user
 from werkzeug.security import check_password_hash, generate_password_hash
 
+from airflow.configuration import conf
 from airflow.www.fab_security.sqla.models import Action, Permission, RegisterUser, Resource, Role, User
 from airflow.www.views import ResourceModelView
 
@@ -854,6 +856,14 @@ def update_user_auth_stat(self, user, success=True):
             user.fail_login_count += 1
         self.update_user(user)
 
+    def _rotate_session_id(self):
+        """
+        Upon successful authentication when using the database session backend,
+        we need to rotate the session id
+        """
+        if conf.get('webserver', 'SESSION_BACKEND') == 'database':
+            session.sid = str(uuid4())
+
     def auth_user_db(self, username, password):
         """
         Method for authenticating user, auth db style
@@ -878,6 +888,7 @@ def auth_user_db(self, username, password):
             log.info(LOGMSG_WAR_SEC_LOGIN_FAILED.format(username))
             return None
         elif check_password_hash(user.password, password):
+            self._rotate_session_id()
             self.update_user_auth_stat(user, True)
             return user
         else:
@@ -1174,6 +1185,7 @@ def auth_user_ldap(self, username, password):
 
             # LOGIN SUCCESS (only if user is now registered)
             if user:
+                self._rotate_session_id()
                 self.update_user_auth_stat(user)
                 return user
             else:
@@ -1201,6 +1213,7 @@ def auth_user_oid(self, email):
             log.info(LOGMSG_WAR_SEC_LOGIN_FAILED.format(email))
             return None
         else:
+            self._rotate_session_id()
             self.update_user_auth_stat(user)
             return user
 
@@ -1230,6 +1243,7 @@ def auth_user_remote_user(self, username):
             log.info(LOGMSG_WAR_SEC_LOGIN_FAILED.format(username))
             return None
 
+        self._rotate_session_id()
         self.update_user_auth_stat(user)
         return user
 
@@ -1315,6 +1329,7 @@ def auth_user_oauth(self, userinfo):
 
         # LOGIN SUCCESS (only if user is now registered)
         if user:
+            self._rotate_session_id()
             self.update_user_auth_stat(user)
             return user
         else:
diff --git a/tests/www/views/test_session.py b/tests/www/views/test_session.py
index 9fb6f364695f7..0753f2069a293 100644
--- a/tests/www/views/test_session.py
+++ b/tests/www/views/test_session.py
@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from unittest import mock
+
 import pytest
 
 from airflow.exceptions import AirflowConfigException
@@ -23,12 +25,16 @@
 from tests.test_utils.decorators import dont_initialize_flask_app_submodules
 
 
+def get_session_cookie(client):
+    return next((cookie for cookie in client.cookie_jar if cookie.name == 'session'), None)
+
+
 def test_session_cookie_created_on_login(user_client):
-    assert any(cookie.name == 'session' for cookie in user_client.cookie_jar)
+    assert get_session_cookie(user_client) is not None
 
 
 def test_session_inaccessible_after_logout(user_client):
-    session_cookie = next((cookie for cookie in user_client.cookie_jar if cookie.name == 'session'), None)
+    session_cookie = get_session_cookie(user_client)
     assert session_cookie is not None
 
     resp = user_client.get('/logout/')
@@ -63,3 +69,21 @@ def poorly_configured_app_factory():
     )
     with pytest.raises(AirflowConfigException, match=expected_exc_regex):
         poorly_configured_app_factory()
+
+
+def test_session_id_rotates(app, user_client):
+    old_session_cookie = get_session_cookie(user_client)
+    assert old_session_cookie is not None
+
+    resp = user_client.get('/logout/')
+    assert resp.status_code == 302
+
+    patch_path = "airflow.www.fab_security.manager.check_password_hash"
+    with mock.patch(patch_path) as check_password_hash:
+        check_password_hash.return_value = True
+        resp = user_client.post("/login/", data={"username": "test_user", "password": "test_user"})
+    assert resp.status_code == 302
+
+    new_session_cookie = get_session_cookie(user_client)
+    assert new_session_cookie is not None
+    assert old_session_cookie.value != new_session_cookie.value