PYTHON-1332 - Check current user owns session

Author: ajdavis

pymongo/bulk.py

@@ -310,6 +310,8 @@ def execute_command(self, sock_info, generator, write_concern, session):
         listeners = self.collection.database.client._event_listeners
 
         with self.collection.database.client._tmp_session(session) as s:
+            # sock_info.command checks auth, but we use sock_info.write_command.
+            sock_info.check_session_auth_matches(s)
             for run in generator:
                 cmd = SON([(_COMMANDS[run.op_type], self.collection.name),
                            ('ordered', self.ordered)])

pymongo/client_session.py

@@ -66,22 +66,13 @@ def causally_consistent_reads(self):
 
 
 class ClientSession(object):
-    """A session for ordering sequential operations.
-
-    :Parameters:
-      - `client`: A :class:`~pymongo.mongo_client.MongoClient`.
-      - `options` (optional): A :class:`SessionOptions` instance.
-    """
-    def __init__(self, client, options=None):
+    """A session for ordering sequential operations."""
+    def __init__(self, client, server_session, options, authset):
+        # A MongoClient, a _ServerSession, a SessionOptions, and a set.
         self._client = client
-
-        if options is not None:
-            self._options = options
-        else:
-            self._options = SessionOptions()
-
-        # Raises ConfigurationError if sessions are not supported.
-        self._server_session = client._get_server_session()
+        self._server_session = server_session
+        self._options = options
+        self._authset = authset
 
     def end_session(self):
         """Finish this session.

pymongo/message.py

@@ -274,6 +274,8 @@ def use_command(self, sock_info, exhaust):
                 'Specifying a collation is unsupported with a max wire '
                 'version of %d.' % (sock_info.max_wire_version,))
 
+        sock_info.check_session_auth_matches(self.session)
+
         return use_find_cmd
 
     def as_command(self):
@@ -345,6 +347,7 @@ def __init__(self, db, coll, ntoreturn, cursor_id, codec_options, session,
         self.session = session
 
     def use_command(self, sock_info, exhaust):
+        sock_info.check_session_auth_matches(self.session)
         return sock_info.max_wire_version >= 4 and not exhaust
 
     def as_command(self):

pymongo/mongo_client.py

@@ -1219,12 +1219,15 @@ def start_session(self, **kwargs):
         # Driver Sessions Spec: "If startSession is called when multiple users
         # are authenticated drivers MUST raise an error with the error message
         # 'Cannot call startSession when multiple users are authenticated.'"
-        if len(self.__all_credentials) > 1:
+        authset = set(self.__all_credentials.values())
+        if len(authset) > 1:
             raise InvalidOperation("Cannot call start_session when"
                                    " multiple users are authenticated")
 
+        # Raises ConfigurationError if sessions are not supported.
+        server_session = self._get_server_session()
         opts = client_session.SessionOptions(**kwargs)
-        return client_session.ClientSession(self, opts)
+        return client_session.ClientSession(self, server_session, opts, authset)
 
     def _get_server_session(self):
         """Internal: start or resume a _ServerSession."""

pymongo/pool.py

@@ -37,6 +37,7 @@ class SSLError(socket.error):
 from pymongo.errors import (AutoReconnect,
                             ConnectionFailure,
                             ConfigurationError,
+                            InvalidOperation,
                             DocumentTooLarge,
                             NetworkTimeout,
                             NotMasterError,
@@ -437,7 +438,7 @@ def command(self, dbname, spec, slave_ok=False,
                 parse_write_concern_error=False,
                 collation=None,
                 session=None):
-        """Execute a command or raise ConnectionFailure or OperationFailure.
+        """Execute a command or raise an error.
 
         :Parameters:
           - `dbname`: name of the database on which to run the command
@@ -455,6 +456,7 @@ def command(self, dbname, spec, slave_ok=False,
           - `collation`: The collation for this command.
           - `session`: optional ClientSession instance.
         """
+        self.check_session_auth_matches(session)
         if self.max_wire_version < 4 and not read_concern.ok_for_legacy:
             raise ConfigurationError(
                 'read concern level of %s is not valid '
@@ -583,6 +585,12 @@ def authenticate(self, credentials):
         auth.authenticate(credentials, self)
         self.authset.add(credentials)
 
+    def check_session_auth_matches(self, session):
+        """Raise error if a ClientSession is logged in as a different user."""
+        if session and session._authset != self.authset:
+            raise InvalidOperation('start_session was called while'
+                                   ' authenticated with different credentials')
+
     def close(self):
         self.closed = True
         # Avoid exceptions on interpreter shutdown.

test/test_session.py

@@ -45,6 +45,12 @@ def failed(self, event):
         if not event.command_name.startswith('sasl'):
             super(SessionTestListener, self).failed(event)
 
+    def first_command_started(self):
+        assert len(self.results['started']) >= 1, (
+            "No command-started events")
+
+        return self.results['started'][0]
+
 
 def session_ids(client):
     return [s.session_id for s in client._topology._session_pool]
@@ -122,16 +128,63 @@ def _test_ops(self, client, *ops, **kwargs):
     @client_context.require_auth
     @ignore_deprecations
     def test_session_authenticate_multiple(self):
+        listener = SessionTestListener()
         # Logged in as root.
-        client = rs_or_single_client()
-        client.pymongo_test.add_user('second-user', 'pass')
-        self.addCleanup(client.pymongo_test.remove_user, 'second-user')
-
-        client.pymongo_test.authenticate('second-user', 'pass')
+        client = rs_or_single_client(event_listeners=[listener])
+        db = client.pymongo_test
+        db.add_user('second-user', 'pass')
+        self.addCleanup(db.remove_user, 'second-user')
+        db.authenticate('second-user', 'pass')
 
         with self.assertRaises(InvalidOperation):
             client.start_session()
 
+        # No implicit sessions.
+        listener.results.clear()
+        db.collection.find_one()
+        event = listener.first_command_started()
+        self.assertNotIn(
+            'lsid', event.command,
+            "find_one with multi-auth shouldn't have sent lsid with %s" % (
+                event.command_name))
+
+        # Changing auth invalidates the session. Start as root.
+        client = rs_or_single_client(event_listeners=[listener])
+        db = client.pymongo_test
+        db.collection.insert_many([{} for _ in range(10)])
+        self.addCleanup(db.collection.drop)
+        with client.start_session() as s:
+            listener.results.clear()
+            cursor = db.collection.find(session=s).batch_size(2)
+            next(cursor)
+            event = listener.first_command_started()
+            self.assertEqual(event.command_name, 'find')
+            self.assertEqual(
+                s.session_id, event.command.get('lsid'),
+                "find sent wrong lsid with %s" % (event.command_name,))
+
+            client.admin.logout()
+            db.authenticate('second-user', 'pass')
+
+            err = 'start_session was called while authenticated with' \
+                  ' different credentials'
+
+            with self.assertRaisesRegex(InvalidOperation, err):
+                # Auth has changed between find and getMore.
+                list(cursor)
+
+            with self.assertRaisesRegex(InvalidOperation, err):
+                db.collection.bulk_write([InsertOne({})], session=s)
+
+            with self.assertRaisesRegex(InvalidOperation, err):
+                db.collection_names(session=s)
+
+            with self.assertRaisesRegex(InvalidOperation, err):
+                db.collection.find_one(session=s)
+
+            with self.assertRaisesRegex(InvalidOperation, err):
+                list(db.collection.aggregate([], session=s))
+
     def test_pool_lifo(self):
         # "Pool is LIFO" test from Driver Sessions Spec.
         a = self.client.start_session()
@@ -380,8 +433,7 @@ def test_cursor(self):
         for name, f in ops:
             listener.results.clear()
             f(session=None)
-            self.assertGreaterEqual(len(listener.results['started']), 1)
-            event0 = listener.results['started'][0]
+            event0 = listener.first_command_started()
             self.assertTrue(
                 'lsid' in event0.command,
                 "%s sent no lsid with %s" % (
@@ -573,8 +625,7 @@ def test_aggregate_error(self):
         with self.assertRaises(OperationFailure):
             coll.aggregate([{'$badOperation': {'bar': 1}}])
 
-        self.assertEqual(len(listener.results['started']), 1)
-        event = listener.results['started'][0]
+        event = listener.first_command_started()
         self.assertEqual(event.command_name, 'aggregate')
         lsid = event.command['lsid']
         # Session was returned to pool despite error.