Implement ALPN in tlslite.

third_party/tlslite/tlslite/messages.py

 # Authors: 
 #   Trevor Perrin
 #   Google - handling CertificateRequest.certificate_types
 #   Google (adapted by Sam Rushing and Marcelo Fernandez) - NPN support
 #   Dimitris Moraitis - Anon ciphersuites
 #   Yngve Pettersen (ported by Paul Sokolovsky) - TLS 1.2
 #
 # See the LICENSE file for legal information regarding use of this file.
 
 """Classes representing TLS messages."""
 
 from .utils.compat import *
 from .utils.cryptomath import *
 from .errors import *
 from .utils.codec import *
 from .constants import *
 from .x509 import X509
 from .x509certchain import X509CertChain
 from .utils.tackwrapper import *
 
+def parse_next_protos(b):
+    protos = []
+    while True:
+        if len(b) == 0:
+            break
+        l = b[0]
+        b = b[1:]
+        if len(b) < l:
+            raise BadNextProtos(len(b))
+        protos.append(b[:l])
+        b = b[l:]
+    return protos
+
+def next_protos_encoded(protocol_list):
+    b = bytearray()
+    for e in protocol_list:
+        if len(e) > 255 or len(e) == 0:
+            raise BadNextProtos(len(e))
+        b += bytearray( [len(e)] ) + bytearray(e)
+    return b
+
 class RecordHeader3(object):
     def __init__(self):
         self.type = 0
         self.version = (0,0)
         self.length = 0
         self.ssl2 = False
 
     def create(self, version, type, length):
         self.type = type
         self.version = version
@@ skipping 73 matching lines @@
         HandshakeMsg.__init__(self, HandshakeType.client_hello)
         self.ssl2 = ssl2
         self.client_version = (0,0)
         self.random = bytearray(32)
         self.session_id = bytearray(0)
         self.cipher_suites = []         # a list of 16-bit values
         self.certificate_types = [CertificateType.x509]
         self.compression_methods = []   # a list of 8-bit values
         self.srp_username = None        # a string
         self.tack = False
+        self.alpn_protos_advertised = None
         self.supports_npn = False
         self.server_name = bytearray(0)
         self.channel_id = False
         self.extended_master_secret = False
         self.tb_client_params = []
         self.support_signed_cert_timestamps = False
         self.status_request = False
 
     def create(self, version, random, session_id, cipher_suites,
                certificate_types=None, srpUsername=None,
-               tack=False, supports_npn=False, serverName=None):
+               tack=False, alpn_protos_advertised=None,
+               supports_npn=False, serverName=None):
         self.client_version = version
         self.random = random
         self.session_id = session_id
         self.cipher_suites = cipher_suites
         self.certificate_types = certificate_types
         self.compression_methods = [0]
         if srpUsername:
             self.srp_username = bytearray(srpUsername, "utf-8")
         self.tack = tack
+        self.alpn_protos_advertised = alpn_protos_advertised
         self.supports_npn = supports_npn
         if serverName:
             self.server_name = bytearray(serverName, "utf-8")
         return self
 
     def parse(self, p):
         if self.ssl2:
             self.client_version = (p.get(1), p.get(1))
             cipherSpecsLength = p.get(2)
             sessionIDLength = p.get(2)
@@ skipping 20 matching lines @@
                 while soFar != totalExtLength:
                     extType = p.get(2)
                     extLength = p.get(2)
                     index1 = p.index
                     if extType == ExtensionType.srp:
                         self.srp_username = p.getVarBytes(1)
                     elif extType == ExtensionType.cert_type:
                         self.certificate_types = p.getVarList(1, 1)
                     elif extType == ExtensionType.tack:
                         self.tack = True
+                    elif extType == ExtensionType.alpn:
+                        structLength = p.get(2)
+                        if structLength + 2 != extLength:
+                            raise SyntaxError()
+                        self.alpn_protos_advertised = parse_next_protos(p.getFixBytes(structLength))
                     elif extType == ExtensionType.supports_npn:
                         self.supports_npn = True
                     elif extType == ExtensionType.server_name:
                         serverNameListBytes = p.getFixBytes(extLength)
                         p2 = Parser(serverNameListBytes)
                         p2.startLengthCheck(2)
                         while 1:
                             if p2.atLengthCheck():
                                 break # no host_name, oh well
                             name_type = p2.get(1)
@@ skipping 52 matching lines @@
         w2 = Writer() # For Extensions
         if self.certificate_types and self.certificate_types != \
                 [CertificateType.x509]:
             w2.add(ExtensionType.cert_type, 2)
             w2.add(len(self.certificate_types)+1, 2)
             w2.addVarSeq(self.certificate_types, 1, 1)
         if self.srp_username:
             w2.add(ExtensionType.srp, 2)
             w2.add(len(self.srp_username)+1, 2)
             w2.addVarSeq(self.srp_username, 1, 1)
+        if self.alpn_protos_advertised is not None:
+            encoded_alpn_protos_advertised = next_protos_encoded(self.alpn_protos_advertised)
+            w2.add(ExtensionType.alpn, 2)
+            w2.add(len(encoded_alpn_protos_advertised) + 2, 2)
+            w2.add(len(encoded_alpn_protos_advertised), 2)
+            w2.addFixSeq(encoded_alpn_protos_advertised, 1)
         if self.supports_npn:
             w2.add(ExtensionType.supports_npn, 2)
             w2.add(0, 2)
         if self.server_name:
             w2.add(ExtensionType.server_name, 2)
             w2.add(len(self.server_name)+5, 2)
             w2.add(len(self.server_name)+3, 2)            
             w2.add(NameType.host_name, 1)
             w2.addVarSeq(self.server_name, 1, 2) 
         if self.tack:
             w2.add(ExtensionType.tack, 2)
             w2.add(0, 2)
         if len(w2.bytes):
             w.add(len(w2.bytes), 2)
             w.bytes += w2.bytes
         return self.postWrite(w)
 
 class BadNextProtos(Exception):
     def __init__(self, l):
         self.length = l
 
     def __str__(self):
         return 'Cannot encode a list of next protocols because it contains an element with invalid length %d. Element lengths must be 0 < x < 256' % self.length
 
+class InvalidALPNResponse(Exception):
+    def __init__(self, l):
+        self.length = l
+
+    def __str__(self):
+        return 'ALPN server response protocol list has invalid length %d.  It must be of length one.' % self.length
+
 class ServerHello(HandshakeMsg):
     def __init__(self):
         HandshakeMsg.__init__(self, HandshakeType.server_hello)
         self.server_version = (0,0)
         self.random = bytearray(32)
         self.session_id = bytearray(0)
         self.cipher_suite = 0
         self.certificate_type = CertificateType.x509
         self.compression_method = 0
         self.tackExt = None
+        self.alpn_proto_selected = None
         self.next_protos_advertised = None
         self.next_protos = None
         self.channel_id = False
         self.extended_master_secret = False
         self.tb_params = None
         self.signed_cert_timestamps = None
         self.status_request = False
 
     def create(self, version, random, session_id, cipher_suite,
-               certificate_type, tackExt, next_protos_advertised):
+               certificate_type, tackExt, alpn_proto_selected,
+               next_protos_advertised):
         self.server_version = version
         self.random = random
         self.session_id = session_id
         self.cipher_suite = cipher_suite
         self.certificate_type = certificate_type
         self.compression_method = 0
         self.tackExt = tackExt
+        self.alpn_proto_selected = alpn_proto_selected
         self.next_protos_advertised = next_protos_advertised
         return self
 
     def parse(self, p):
         p.startLengthCheck(3)
         self.server_version = (p.get(1), p.get(1))
         self.random = p.getFixBytes(32)
         self.session_id = p.getVarBytes(1)
         self.cipher_suite = p.get(2)
         self.compression_method = p.get(1)
         if not p.atLengthCheck():
             totalExtLength = p.get(2)
             soFar = 0
             while soFar != totalExtLength:
                 extType = p.get(2)
                 extLength = p.get(2)
                 if extType == ExtensionType.cert_type:
                     if extLength != 1:
                         raise SyntaxError()
                     self.certificate_type = p.get(1)
                 elif extType == ExtensionType.tack and tackpyLoaded:
                     self.tackExt = TackExtension(p.getFixBytes(extLength))
+                elif extType == ExtensionType.alpn:
+                    structLength = p.get(2)
+                    if structLength + 2 != extLength:
+                        raise SyntaxError()
+                    alpn_protos = parse_next_protos(p.getFixBytes(structLength))
+                    if len(alpn_protos) != 1:
+                        raise InvalidALPNResponse(len(alpn_protos));
+                    self.alpn_proto_selected = alpn_protos[0]
                 elif extType == ExtensionType.supports_npn:
-                    self.next_protos = self.__parse_next_protos(p.getFixBytes(extLength))
+                    self.next_protos = parse_next_protos(p.getFixBytes(extLength))
                 else:
                     p.getFixBytes(extLength)
                 soFar += 4 + extLength
         p.stopLengthCheck()
         return self
 
-    def __parse_next_protos(self, b):
-        protos = []
-        while True:
-            if len(b) == 0:
-                break
-            l = b[0]
-            b = b[1:]
-            if len(b) < l:
-                raise BadNextProtos(len(b))
-            protos.append(b[:l])
-            b = b[l:]
-        return protos
-
-    def __next_protos_encoded(self):
-        b = bytearray()
-        for e in self.next_protos_advertised:
-            if len(e) > 255 or len(e) == 0:
-                raise BadNextProtos(len(e))
-            b += bytearray( [len(e)] ) + bytearray(e)
-        return b
-
     def write(self):
         w = Writer()
         w.add(self.server_version[0], 1)
         w.add(self.server_version[1], 1)
         w.addFixSeq(self.random, 1)
         w.addVarSeq(self.session_id, 1, 1)
         w.add(self.cipher_suite, 2)
         w.add(self.compression_method, 1)
 
         w2 = Writer() # For Extensions
         if self.certificate_type and self.certificate_type != \
                 CertificateType.x509:
             w2.add(ExtensionType.cert_type, 2)
             w2.add(1, 2)
             w2.add(self.certificate_type, 1)
         if self.tackExt:
             b = self.tackExt.serialize()
             w2.add(ExtensionType.tack, 2)
             w2.add(len(b), 2)
             w2.bytes += b
+        if self.alpn_proto_selected is not None:
+            alpn_protos_single_element_list = [self.alpn_proto_selected]
+            encoded_alpn_protos_advertised = next_protos_encoded(alpn_protos_single_element_list)
+            w2.add(ExtensionType.alpn, 2)
+            w2.add(len(encoded_alpn_protos_advertised) + 2, 2)
+            w2.add(len(encoded_alpn_protos_advertised), 2)
+            w2.addFixSeq(encoded_alpn_protos_advertised, 1)
         if self.next_protos_advertised is not None:
-            encoded_next_protos_advertised = self.__next_protos_encoded()
+            encoded_next_protos_advertised = next_protos_encoded(self.next_protos_advertised)
             w2.add(ExtensionType.supports_npn, 2)
             w2.add(len(encoded_next_protos_advertised), 2)
             w2.addFixSeq(encoded_next_protos_advertised, 1)
         if self.channel_id:
             w2.add(ExtensionType.channel_id, 2)
             w2.add(0, 2)
         if self.extended_master_secret:
             w2.add(ExtensionType.extended_master_secret, 2)
             w2.add(0, 2)
         if self.tb_params:
@@ skipping 445 matching lines @@
         newMsg = ApplicationData().create(self.bytes[:1])
         self.bytes = self.bytes[1:]
         return newMsg
 
     def parse(self, p):
         self.bytes = p.bytes
         return self
 
     def write(self):
         return self.bytes