capsule.py (3755B)
1 from enum import IntEnum 2 from typing import Iterator, Optional 3 4 from aioquic.buffer import UINT_VAR_MAX_SIZE, Buffer, BufferReadError 5 6 7 class CapsuleType(IntEnum): 8 # Defined in 9 # https://datatracker.ietf.org/doc/html/draft-ietf-masque-h3-datagram-04#section-8.2 10 DATAGRAM_DRAFT04 = 0xff37a0 11 REGISTER_DATAGRAM_CONTEXT_DRAFT04 = 0xff37a1 12 REGISTER_DATAGRAM_NO_CONTEXT_DRAFT04 = 0xff37a2 13 CLOSE_DATAGRAM_CONTEXT_DRAFT04 = 0xff37a3 14 # Defined in 15 # https://datatracker.ietf.org/doc/html/rfc9297#section-5.4 16 DATAGRAM_RFC = 0x00 17 # Defined in 18 # https://www.ietf.org/archive/id/draft-ietf-webtrans-http3-01.html. 19 CLOSE_WEBTRANSPORT_SESSION = 0x2843 20 21 22 class H3Capsule: 23 """ 24 Represents the Capsule concept defined in 25 https://ietf-wg-masque.github.io/draft-ietf-masque-h3-datagram/draft-ietf-masque-h3-datagram.html#name-capsules. 26 """ 27 28 def __init__(self, type: int, data: bytes) -> None: 29 """ 30 :param type the type of this Capsule. We don't use CapsuleType here 31 because this may be a capsule of an unknown type. 32 :param data the payload 33 """ 34 self.type = type 35 self.data = data 36 37 def encode(self) -> bytes: 38 """ 39 Encodes this H3Capsule and return the bytes. 40 """ 41 buffer = Buffer(capacity=len(self.data) + 2 * UINT_VAR_MAX_SIZE) 42 buffer.push_uint_var(self.type) 43 buffer.push_uint_var(len(self.data)) 44 buffer.push_bytes(self.data) 45 return buffer.data 46 47 48 class H3CapsuleDecoder: 49 """ 50 A decoder of H3Capsule. This is a streaming decoder and can handle multiple 51 decoders. 52 """ 53 54 def __init__(self) -> None: 55 self._buffer: Optional[Buffer] = None 56 self._type: Optional[int] = None 57 self._length: Optional[int] = None 58 self._final: bool = False 59 60 def append(self, data: bytes) -> None: 61 """ 62 Appends the given bytes to this decoder. 63 """ 64 assert not self._final 65 66 if len(data) == 0: 67 return 68 if self._buffer: 69 remaining = self._buffer.pull_bytes( 70 self._buffer.capacity - self._buffer.tell()) 71 self._buffer = Buffer(data=(remaining + data)) 72 else: 73 self._buffer = Buffer(data=data) 74 75 def final(self) -> None: 76 """ 77 Pushes the end-of-stream mark to this decoder. After calling this, 78 calling append() will be invalid. 79 """ 80 self._final = True 81 82 def __iter__(self) -> Iterator[H3Capsule]: 83 """ 84 Yields decoded capsules. 85 """ 86 try: 87 while self._buffer is not None: 88 if self._type is None: 89 self._type = self._buffer.pull_uint_var() 90 if self._length is None: 91 self._length = self._buffer.pull_uint_var() 92 if self._buffer.capacity - self._buffer.tell() < self._length: 93 if self._final: 94 raise ValueError('insufficient buffer') 95 return 96 capsule = H3Capsule( 97 self._type, self._buffer.pull_bytes(self._length)) 98 self._type = None 99 self._length = None 100 if self._buffer.tell() == self._buffer.capacity: 101 self._buffer = None 102 yield capsule 103 except BufferReadError as e: 104 if self._final: 105 raise e 106 if not self._buffer: 107 return 108 size = self._buffer.capacity - self._buffer.tell() 109 if size >= UINT_VAR_MAX_SIZE: 110 raise e 111 # Ignore the error because there may not be sufficient input. 112 return