merge master, resolve simple conflicts

1 files changed, 242 insertions, 38 deletions

diff --git a/OpenSSL/SSL.py b/OpenSSL/SSL.py
index 0e645ef..36ef14a 100644
--- a/OpenSSL/SSL.py
+++ b/OpenSSL/SSL.py
@@ -1,18 +1,21 @@
 from sys import platform
 from functools import wraps, partial
-from itertools import count
+from itertools import count, chain
 from weakref import WeakValueDictionary
 from errno import errorcode
 
 from six import text_type as _text_type
 from six import integer_types as integer_types
+from six import int2byte, indexbytes
 
 from OpenSSL._util import (
     ffi as _ffi,
     lib as _lib,
     exception_from_error_queue as _exception_from_error_queue,
     native as _native,
-    warn_text as _warn_text)
+    warn_text as _warn_text,
+    path_string as _path_string,
+)
 
 from OpenSSL.crypto import (
     FILETYPE_PEM, _PassphraseHelper, PKey, X509Name, X509, X509Store)
@@ -165,10 +168,41 @@ class SysCallError(Error):
     pass
 
 
+class _CallbackExceptionHelper(object):
+    """
+    A base class for wrapper classes that allow for intelligent exception
+    handling in OpenSSL callbacks.
+
+    :ivar list _problems: Any exceptions that occurred while executing in a
+        context where they could not be raised in the normal way.  Typically
+        this is because OpenSSL has called into some Python code and requires a
+        return value.  The exceptions are saved to be raised later when it is
+        possible to do so.
+    """
+    def __init__(self):
+        self._problems = []
 
-class _VerifyHelper(object):
+
+    def raise_if_problem(self):
+        """
+        Raise an exception from the OpenSSL error queue or that was previously
+        captured whe running a callback.
+        """
+        if self._problems:
+            try:
+                _raise_current_error()
+            except Error:
+                pass
+            raise self._problems.pop(0)
+
+
+class _VerifyHelper(_CallbackExceptionHelper):
+    """
+    Wrap a callback such that it can be used as a certificate verification
+    callback.
+    """
     def __init__(self, callback):
-        self._problems = []
+        _CallbackExceptionHelper.__init__(self)
 
         @wraps(callback)
         def wrapper(ok, store_ctx):
@@ -197,14 +231,92 @@ class _VerifyHelper(object):
             "int (*)(int, X509_STORE_CTX *)", wrapper)
 
 
-    def raise_if_problem(self):
-        if self._problems:
+class _NpnAdvertiseHelper(_CallbackExceptionHelper):
+    """
+    Wrap a callback such that it can be used as an NPN advertisement callback.
+    """
+    def __init__(self, callback):
+        _CallbackExceptionHelper.__init__(self)
+
+        @wraps(callback)
+        def wrapper(ssl, out, outlen, arg):
             try:
-                _raise_current_error()
-            except Error:
-                pass
-            raise self._problems.pop(0)
+                conn = Connection._reverse_mapping[ssl]
+                protos = callback(conn)
+
+                # Join the protocols into a Python bytestring, length-prefixing
+                # each element.
+                protostr = b''.join(
+                    chain.from_iterable((int2byte(len(p)), p) for p in protos)
+                )
+
+                # Save our callback arguments on the connection object. This is
+                # done to make sure that they don't get freed before OpenSSL
+                # uses them. Then, return them appropriately in the output
+                # parameters.
+                conn._npn_advertise_callback_args = [
+                    _ffi.new("unsigned int *", len(protostr)),
+                    _ffi.new("unsigned char[]", protostr),
+                ]
+                outlen[0] = conn._npn_advertise_callback_args[0][0]
+                out[0] = conn._npn_advertise_callback_args[1]
+                return 0
+            except Exception as e:
+                self._problems.append(e)
+                return 2  # SSL_TLSEXT_ERR_ALERT_FATAL
+
+        self.callback = _ffi.callback(
+            "int (*)(SSL *, const unsigned char **, unsigned int *, void *)",
+            wrapper
+        )
+
+
+class _NpnSelectHelper(_CallbackExceptionHelper):
+    """
+    Wrap a callback such that it can be used as an NPN selection callback.
+    """
+    def __init__(self, callback):
+        _CallbackExceptionHelper.__init__(self)
+
+        @wraps(callback)
+        def wrapper(ssl, out, outlen, in_, inlen, arg):
+            try:
+                conn = Connection._reverse_mapping[ssl]
+
+                # The string passed to us is actually made up of multiple
+                # length-prefixed bytestrings. We need to split that into a
+                # list.
+                instr = _ffi.buffer(in_, inlen)[:]
+                protolist = []
+                while instr:
+                    l = indexbytes(instr, 0)
+                    proto = instr[1:l+1]
+                    protolist.append(proto)
+                    instr = instr[l+1:]
+
+                # Call the callback
+                outstr = callback(conn, protolist)
+
+                # Save our callback arguments on the connection object. This is
+                # done to make sure that they don't get freed before OpenSSL
+                # uses them. Then, return them appropriately in the output
+                # parameters.
+                conn._npn_select_callback_args = [
+                    _ffi.new("unsigned char *", len(outstr)),
+                    _ffi.new("unsigned char[]", outstr),
+                ]
+                outlen[0] = conn._npn_select_callback_args[0][0]
+                out[0] = conn._npn_select_callback_args[1]
+                return 0
+            except Exception as e:
+                self._problems.append(e)
+                return 2  # SSL_TLSEXT_ERR_ALERT_FATAL
 
+        self.callback = _ffi.callback(
+            "int (*)(SSL *, unsigned char **, unsigned char *, "
+                    "const unsigned char *, unsigned int, void *)",
+            wrapper
+        )
 
 
 def _asFileDescriptor(obj):
@@ -294,6 +406,10 @@ class Context(object):
         self._info_callback = None
         self._tlsext_servername_callback = None
         self._app_data = None
+        self._npn_advertise_helper = None
+        self._npn_advertise_callback = None
+        self._npn_select_helper = None
+        self._npn_select_callback = None
 
         # SSL_CTX_set_app_data(self->ctx, self);
         # SSL_CTX_set_mode(self->ctx, SSL_MODE_ENABLE_PARTIAL_WRITE |
@@ -307,19 +423,22 @@ class Context(object):
         Let SSL know where we can find trusted certificates for the certificate
         chain
 
-        :param cafile: In which file we can find the certificates
+        :param cafile: In which file we can find the certificates (``bytes`` or
+            ``unicode``).
         :param capath: In which directory we can find the certificates
+            (``bytes`` or ``unicode``).
+
         :return: None
         """
         if cafile is None:
             cafile = _ffi.NULL
-        elif not isinstance(cafile, bytes):
-            raise TypeError("cafile must be None or a byte string")
+        else:
+            cafile = _path_string(cafile)
 
         if capath is None:
             capath = _ffi.NULL
-        elif not isinstance(capath, bytes):
-            raise TypeError("capath must be None or a byte string")
+        else:
+            capath = _path_string(capath)
 
         load_result = _lib.SSL_CTX_load_verify_locations(self._context, cafile, capath)
         if not load_result:
@@ -369,15 +488,12 @@ class Context(object):
         """
         Load a certificate chain from a file
 
-        :param certfile: The name of the certificate chain file
+        :param certfile: The name of the certificate chain file (``bytes`` or
+            ``unicode``).
+
         :return: None
         """
-        if isinstance(certfile, _text_type):
-            # Perhaps sys.getfilesystemencoding() could be better?
-            certfile = certfile.encode("utf-8")
-
-        if not isinstance(certfile, bytes):
-            raise TypeError("certfile must be bytes or unicode")
+        certfile = _path_string(certfile)
 
         result = _lib.SSL_CTX_use_certificate_chain_file(self._context, certfile)
         if not result:
@@ -388,15 +504,13 @@ class Context(object):
         """
         Load a certificate from a file
 
-        :param certfile: The name of the certificate file
+        :param certfile: The name of the certificate file (``bytes`` or
+            ``unicode``).
         :param filetype: (optional) The encoding of the file, default is PEM
+
         :return: None
         """
-        if isinstance(certfile, _text_type):
-            # Perhaps sys.getfilesystemencoding() could be better?
-            certfile = certfile.encode("utf-8")
-        if not isinstance(certfile, bytes):
-            raise TypeError("certfile must be bytes or unicode")
+        certfile = _path_string(certfile)
         if not isinstance(filetype, integer_types):
             raise TypeError("filetype must be an integer")
 
@@ -450,16 +564,12 @@ class Context(object):
         """
         Load a private key from a file
 
-        :param keyfile: The name of the key file
+        :param keyfile: The name of the key file (``bytes`` or ``unicode``)
         :param filetype: (optional) The encoding of the file, default is PEM
+
         :return: None
         """
-        if isinstance(keyfile, _text_type):
-            # Perhaps sys.getfilesystemencoding() could be better?
-            keyfile = keyfile.encode("utf-8")
-
-        if not isinstance(keyfile, bytes):
-            raise TypeError("keyfile must be a byte string")
+        keyfile = _path_string(keyfile)
 
         if filetype is _unspecified:
             filetype = FILETYPE_PEM
@@ -595,11 +705,12 @@ class Context(object):
         """
         Load parameters for Ephemeral Diffie-Hellman
 
-        :param dhfile: The file to load EDH parameters from
+        :param dhfile: The file to load EDH parameters from (``bytes`` or
+            ``unicode``).
+
         :return: None
         """
-        if not isinstance(dhfile, bytes):
-            raise TypeError("dhfile must be a byte string")
+        dhfile = _path_string(dhfile)
 
         bio = _lib.BIO_new_file(dhfile, b"r")
         if bio == _ffi.NULL:
@@ -813,6 +924,39 @@ class Context(object):
         _lib.SSL_CTX_set_tlsext_servername_callback(
             self._context, self._tlsext_servername_callback)
 
+
+    def set_npn_advertise_callback(self, callback):
+        """
+        Specify a callback function that will be called when offering `Next
+        Protocol Negotiation
+        <https://technotes.googlecode.com/git/nextprotoneg.html>`_ as a server.
+
+        :param callback: The callback function.  It will be invoked with one
+            argument, the Connection instance.  It should return a list of
+            bytestrings representing the advertised protocols, like
+            ``[b'http/1.1', b'spdy/2']``.
+        """
+        self._npn_advertise_helper = _NpnAdvertiseHelper(callback)
+        self._npn_advertise_callback = self._npn_advertise_helper.callback
+        _lib.SSL_CTX_set_next_protos_advertised_cb(
+            self._context, self._npn_advertise_callback, _ffi.NULL)
+
+
+    def set_npn_select_callback(self, callback):
+        """
+        Specify a callback function that will be called when a server offers
+        Next Protocol Negotiation options.
+
+        :param callback: The callback function.  It will be invoked with two
+            arguments: the Connection, and a list of offered protocols as
+            bytestrings, e.g. ``[b'http/1.1', b'spdy/2']``.  It should return
+            one of those bytestrings, the chosen protocol.
+        """
+        self._npn_select_helper = _NpnSelectHelper(callback)
+        self._npn_select_callback = self._npn_select_helper.callback
+        _lib.SSL_CTX_set_next_proto_select_cb(
+            self._context, self._npn_select_callback, _ffi.NULL)
+
 ContextType = Context
 
 
@@ -837,6 +981,13 @@ class Connection(object):
         self._ssl = _ffi.gc(ssl, _lib.SSL_free)
         self._context = context
 
+        # References to strings used for Next Protocol Negotiation. OpenSSL's
+        # header files suggest that these might get copied at some point, but
+        # doesn't specify when, so we store them here to make sure they don't
+        # get freed before OpenSSL uses them.
+        self._npn_advertise_callback_args = None
+        self._npn_select_callback_args = None
+
         self._reverse_mapping[self._ssl] = self
 
         if socket is None:
@@ -871,6 +1022,10 @@ class Connection(object):
     def _raise_ssl_error(self, ssl, result):
         if self._context._verify_helper is not None:
             self._context._verify_helper.raise_if_problem()
+        if self._context._npn_advertise_helper is not None:
+            self._context._npn_advertise_helper.raise_if_problem()
+        if self._context._npn_select_helper is not None:
+            self._context._npn_select_helper.raise_if_problem()
 
         error = _lib.SSL_get_error(ssl, result)
         if error == _lib.SSL_ERROR_WANT_READ:
@@ -1036,6 +1191,45 @@ class Connection(object):
     read = recv
 
 
+    def recv_into(self, buffer, nbytes=None, flags=None):
+        """
+        Receive data on the connection and store the data into a buffer rather
+        than creating a new string.
+
+        :param buffer: The buffer to copy into.
+        :param nbytes: (optional) The maximum number of bytes to read into the
+            buffer. If not present, defaults to the size of the buffer. If
+            larger than the size of the buffer, is reduced to the size of the
+            buffer.
+        :param flags: (optional) Included for compatibility with the socket
+            API, the value is ignored.
+        :return: The number of bytes read into the buffer.
+        """
+        if nbytes is None:
+            nbytes = len(buffer)
+        else:
+            nbytes = min(nbytes, len(buffer))
+
+        # We need to create a temporary buffer. This is annoying, it would be
+        # better if we could pass memoryviews straight into the SSL_read call,
+        # but right now we can't. Revisit this if CFFI gets that ability.
+        buf = _ffi.new("char[]", nbytes)
+        result = _lib.SSL_read(self._ssl, buf, nbytes)
+        self._raise_ssl_error(self._ssl, result)
+
+        # This strange line is all to avoid a memory copy. The buffer protocol
+        # should allow us to assign a CFFI buffer to the LHS of this line, but
+        # on CPython 3.3+ that segfaults. As a workaround, we can temporarily
+        # wrap it in a memoryview, except on Python 2.6 which doesn't have a
+        # memoryview type.
+        try:
+            buffer[:result] = memoryview(_ffi.buffer(buf, result))
+        except NameError:
+            buffer[:result] = _ffi.buffer(buf, result)
+
+        return result
+
+
     def _handle_bio_errors(self, bio, result):
         if _lib.BIO_should_retry(bio):
             if _lib.BIO_should_read(bio):
@@ -1560,6 +1754,16 @@ class Connection(object):
             version =_ffi.string(_lib.SSL_CIPHER_get_version(cipher))
             return version.decode("utf-8")
 
+    def get_next_proto_negotiated(self):
+        """
+        Get the protocol that was negotiated by NPN.
+        """
+        data = _ffi.new("unsigned char **")
+        data_len = _ffi.new("unsigned int *")
+
+        _lib.SSL_get0_next_proto_negotiated(self._ssl, data, data_len)
+
+        return _ffi.buffer(data[0], data_len[0])[:]
 
 
 ConnectionType = Connection