sensu-plugins-mongodb-mrtrotl 1.4.0

data/lib/pymongo/network.py

@@ -0,0 +1,311 @@
+# Copyright 2015-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Internal network layer helper methods."""
+
+import datetime
+import errno
+import socket
+import struct
+import time
+
+from bson import _decode_all_selective
+from pymongo import _csot, helpers, message, ssl_support
+from pymongo.common import MAX_MESSAGE_SIZE
+from pymongo.compression_support import _NO_COMPRESSION, decompress
+from pymongo.errors import (
+    NotPrimaryError,
+    OperationFailure,
+    ProtocolError,
+    _OperationCancelled,
+)
+from pymongo.message import _UNPACK_REPLY, _OpMsg
+from pymongo.monitoring import _is_speculative_authenticate
+from pymongo.socket_checker import _errno_from_exception
+
+_UNPACK_HEADER = struct.Struct("<iiii").unpack
+
+
+def command(
+    sock_info,
+    dbname,
+    spec,
+    is_mongos,
+    read_preference,
+    codec_options,
+    session,
+    client,
+    check=True,
+    allowable_errors=None,
+    address=None,
+    listeners=None,
+    max_bson_size=None,
+    read_concern=None,
+    parse_write_concern_error=False,
+    collation=None,
+    compression_ctx=None,
+    use_op_msg=False,
+    unacknowledged=False,
+    user_fields=None,
+    exhaust_allowed=False,
+    write_concern=None,
+):
+    """Execute a command over the socket, or raise socket.error.
+
+    :Parameters:
+      - `sock`: a raw socket instance
+      - `dbname`: name of the database on which to run the command
+      - `spec`: a command document as an ordered dict type, eg SON.
+      - `is_mongos`: are we connected to a mongos?
+      - `read_preference`: a read preference
+      - `codec_options`: a CodecOptions instance
+      - `session`: optional ClientSession instance.
+      - `client`: optional MongoClient instance for updating $clusterTime.
+      - `check`: raise OperationFailure if there are errors
+      - `allowable_errors`: errors to ignore if `check` is True
+      - `address`: the (host, port) of `sock`
+      - `listeners`: An instance of :class:`~pymongo.monitoring.EventListeners`
+      - `max_bson_size`: The maximum encoded bson size for this server
+      - `read_concern`: The read concern for this command.
+      - `parse_write_concern_error`: Whether to parse the ``writeConcernError``
+        field in the command response.
+      - `collation`: The collation for this command.
+      - `compression_ctx`: optional compression Context.
+      - `use_op_msg`: True if we should use OP_MSG.
+      - `unacknowledged`: True if this is an unacknowledged command.
+      - `user_fields` (optional): Response fields that should be decoded
+        using the TypeDecoders from codec_options, passed to
+        bson._decode_all_selective.
+      - `exhaust_allowed`: True if we should enable OP_MSG exhaustAllowed.
+    """
+    name = next(iter(spec))
+    ns = dbname + ".$cmd"
+    speculative_hello = False
+
+    # Publish the original command document, perhaps with lsid and $clusterTime.
+    orig = spec
+    if is_mongos and not use_op_msg:
+        spec = message._maybe_add_read_preference(spec, read_preference)
+    if read_concern and not (session and session.in_transaction):
+        if read_concern.level:
+            spec["readConcern"] = read_concern.document
+        if session:
+            session._update_read_concern(spec, sock_info)
+    if collation is not None:
+        spec["collation"] = collation
+
+    publish = listeners is not None and listeners.enabled_for_commands
+    if publish:
+        start = datetime.datetime.now()
+        speculative_hello = _is_speculative_authenticate(name, spec)
+
+    if compression_ctx and name.lower() in _NO_COMPRESSION:
+        compression_ctx = None
+
+    if client and client._encrypter and not client._encrypter._bypass_auto_encryption:
+        spec = orig = client._encrypter.encrypt(dbname, spec, codec_options)
+
+    # Support CSOT
+    if client:
+        sock_info.apply_timeout(client, spec)
+    _csot.apply_write_concern(spec, write_concern)
+
+    if use_op_msg:
+        flags = _OpMsg.MORE_TO_COME if unacknowledged else 0
+        flags |= _OpMsg.EXHAUST_ALLOWED if exhaust_allowed else 0
+        request_id, msg, size, max_doc_size = message._op_msg(
+            flags, spec, dbname, read_preference, codec_options, ctx=compression_ctx
+        )
+        # If this is an unacknowledged write then make sure the encoded doc(s)
+        # are small enough, otherwise rely on the server to return an error.
+        if unacknowledged and max_bson_size is not None and max_doc_size > max_bson_size:
+            message._raise_document_too_large(name, size, max_bson_size)
+    else:
+        request_id, msg, size = message._query(
+            0, ns, 0, -1, spec, None, codec_options, compression_ctx
+        )
+
+    if max_bson_size is not None and size > max_bson_size + message._COMMAND_OVERHEAD:
+        message._raise_document_too_large(name, size, max_bson_size + message._COMMAND_OVERHEAD)
+
+    if publish:
+        encoding_duration = datetime.datetime.now() - start
+        listeners.publish_command_start(
+            orig, dbname, request_id, address, service_id=sock_info.service_id
+        )
+        start = datetime.datetime.now()
+
+    try:
+        sock_info.sock.sendall(msg)
+        if use_op_msg and unacknowledged:
+            # Unacknowledged, fake a successful command response.
+            reply = None
+            response_doc = {"ok": 1}
+        else:
+            reply = receive_message(sock_info, request_id)
+            sock_info.more_to_come = reply.more_to_come
+            unpacked_docs = reply.unpack_response(
+                codec_options=codec_options, user_fields=user_fields
+            )
+
+            response_doc = unpacked_docs[0]
+            if client:
+                client._process_response(response_doc, session)
+            if check:
+                helpers._check_command_response(
+                    response_doc,
+                    sock_info.max_wire_version,
+                    allowable_errors,
+                    parse_write_concern_error=parse_write_concern_error,
+                )
+    except Exception as exc:
+        if publish:
+            duration = (datetime.datetime.now() - start) + encoding_duration
+            if isinstance(exc, (NotPrimaryError, OperationFailure)):
+                failure = exc.details
+            else:
+                failure = message._convert_exception(exc)
+            listeners.publish_command_failure(
+                duration, failure, name, request_id, address, service_id=sock_info.service_id
+            )
+        raise
+    if publish:
+        duration = (datetime.datetime.now() - start) + encoding_duration
+        listeners.publish_command_success(
+            duration,
+            response_doc,
+            name,
+            request_id,
+            address,
+            service_id=sock_info.service_id,
+            speculative_hello=speculative_hello,
+        )
+
+    if client and client._encrypter and reply:
+        decrypted = client._encrypter.decrypt(reply.raw_command_response())
+        response_doc = _decode_all_selective(decrypted, codec_options, user_fields)[0]
+
+    return response_doc
+
+
+_UNPACK_COMPRESSION_HEADER = struct.Struct("<iiB").unpack
+
+
+def receive_message(sock_info, request_id, max_message_size=MAX_MESSAGE_SIZE):
+    """Receive a raw BSON message or raise socket.error."""
+    if _csot.get_timeout():
+        deadline = _csot.get_deadline()
+    else:
+        timeout = sock_info.sock.gettimeout()
+        if timeout:
+            deadline = time.monotonic() + timeout
+        else:
+            deadline = None
+    # Ignore the response's request id.
+    length, _, response_to, op_code = _UNPACK_HEADER(
+        _receive_data_on_socket(sock_info, 16, deadline)
+    )
+    # No request_id for exhaust cursor "getMore".
+    if request_id is not None:
+        if request_id != response_to:
+            raise ProtocolError("Got response id %r but expected %r" % (response_to, request_id))
+    if length <= 16:
+        raise ProtocolError(
+            "Message length (%r) not longer than standard message header size (16)" % (length,)
+        )
+    if length > max_message_size:
+        raise ProtocolError(
+            "Message length (%r) is larger than server max "
+            "message size (%r)" % (length, max_message_size)
+        )
+    if op_code == 2012:
+        op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(
+            _receive_data_on_socket(sock_info, 9, deadline)
+        )
+        data = decompress(_receive_data_on_socket(sock_info, length - 25, deadline), compressor_id)
+    else:
+        data = _receive_data_on_socket(sock_info, length - 16, deadline)
+
+    try:
+        unpack_reply = _UNPACK_REPLY[op_code]
+    except KeyError:
+        raise ProtocolError("Got opcode %r but expected %r" % (op_code, _UNPACK_REPLY.keys()))
+    return unpack_reply(data)
+
+
+_POLL_TIMEOUT = 0.5
+
+
+def wait_for_read(sock_info, deadline):
+    """Block until at least one byte is read, or a timeout, or a cancel."""
+    context = sock_info.cancel_context
+    # Only Monitor connections can be cancelled.
+    if context:
+        sock = sock_info.sock
+        timed_out = False
+        while True:
+            # SSLSocket can have buffered data which won't be caught by select.
+            if hasattr(sock, "pending") and sock.pending() > 0:
+                readable = True
+            else:
+                # Wait up to 500ms for the socket to become readable and then
+                # check for cancellation.
+                if deadline:
+                    remaining = deadline - time.monotonic()
+                    # When the timeout has expired perform one final check to
+                    # see if the socket is readable. This helps avoid spurious
+                    # timeouts on AWS Lambda and other FaaS environments.
+                    if remaining <= 0:
+                        timed_out = True
+                    timeout = max(min(remaining, _POLL_TIMEOUT), 0)
+                else:
+                    timeout = _POLL_TIMEOUT
+                readable = sock_info.socket_checker.select(sock, read=True, timeout=timeout)
+            if context.cancelled:
+                raise _OperationCancelled("hello cancelled")
+            if readable:
+                return
+            if timed_out:
+                raise socket.timeout("timed out")
+
+
+# Errors raised by sockets (and TLS sockets) when in non-blocking mode.
+BLOCKING_IO_ERRORS = (BlockingIOError,) + ssl_support.BLOCKING_IO_ERRORS
+
+
+def _receive_data_on_socket(sock_info, length, deadline):
+    buf = bytearray(length)
+    mv = memoryview(buf)
+    bytes_read = 0
+    while bytes_read < length:
+        try:
+            wait_for_read(sock_info, deadline)
+            # CSOT: Update timeout. When the timeout has expired perform one
+            # final non-blocking recv. This helps avoid spurious timeouts when
+            # the response is actually already buffered on the client.
+            if _csot.get_timeout():
+                sock_info.set_socket_timeout(max(deadline - time.monotonic(), 0))
+            chunk_length = sock_info.sock.recv_into(mv[bytes_read:])
+        except BLOCKING_IO_ERRORS:
+            raise socket.timeout("timed out")
+        except (IOError, OSError) as exc:  # noqa: B014
+            if _errno_from_exception(exc) == errno.EINTR:
+                continue
+            raise
+        if chunk_length == 0:
+            raise OSError("connection closed")
+
+        bytes_read += chunk_length
+
+    return mv

data/lib/pymongo/ocsp_cache.py

@@ -0,0 +1,87 @@
+# Copyright 2020-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for caching OCSP responses."""
+
+from collections import namedtuple
+from datetime import datetime as _datetime
+from threading import Lock
+
+
+class _OCSPCache(object):
+    """A cache for OCSP responses."""
+
+    CACHE_KEY_TYPE = namedtuple(  # type: ignore
+        "OcspResponseCacheKey",
+        ["hash_algorithm", "issuer_name_hash", "issuer_key_hash", "serial_number"],
+    )
+
+    def __init__(self):
+        self._data = {}
+        # Hold this lock when accessing _data.
+        self._lock = Lock()
+
+    def _get_cache_key(self, ocsp_request):
+        return self.CACHE_KEY_TYPE(
+            hash_algorithm=ocsp_request.hash_algorithm.name.lower(),
+            issuer_name_hash=ocsp_request.issuer_name_hash,
+            issuer_key_hash=ocsp_request.issuer_key_hash,
+            serial_number=ocsp_request.serial_number,
+        )
+
+    def __setitem__(self, key, value):
+        """Add/update a cache entry.
+
+        'key' is of type cryptography.x509.ocsp.OCSPRequest
+        'value' is of type cryptography.x509.ocsp.OCSPResponse
+
+        Validity of the OCSP response must be checked by caller.
+        """
+        with self._lock:
+            cache_key = self._get_cache_key(key)
+
+            # As per the OCSP protocol, if the response's nextUpdate field is
+            # not set, the responder is indicating that newer revocation
+            # information is available all the time.
+            if value.next_update is None:
+                self._data.pop(cache_key, None)
+                return
+
+            # Do nothing if the response is invalid.
+            if not (value.this_update <= _datetime.utcnow() < value.next_update):
+                return
+
+            # Cache new response OR update cached response if new response
+            # has longer validity.
+            cached_value = self._data.get(cache_key, None)
+            if cached_value is None or cached_value.next_update < value.next_update:
+                self._data[cache_key] = value
+
+    def __getitem__(self, item):
+        """Get a cache entry if it exists.
+
+        'item' is of type cryptography.x509.ocsp.OCSPRequest
+
+        Raises KeyError if the item is not in the cache.
+        """
+        with self._lock:
+            cache_key = self._get_cache_key(item)
+            value = self._data[cache_key]
+
+            # Return cached response if it is still valid.
+            if value.this_update <= _datetime.utcnow() < value.next_update:
+                return value
+
+            self._data.pop(cache_key, None)
+            raise KeyError(cache_key)