Coverage for aiocoap/blockwise.py: 92%

1# SPDX-FileCopyrightText: Christian Amsüss and the aiocoap contributors 

2# 

3# SPDX-License-Identifier: MIT 

4 

5"""Helpers for the implementation of RFC7959 blockwise transfers""" 

6 

7from typing import Awaitable, Callable 

8 

9from . import numbers 

10from .numbers.optionnumbers import OptionNumber 

11from .numbers import codes 

12from .error import ConstructionRenderableError 

13from .message import Message 

14from .optiontypes import BlockOption 

15from .util.asyncio.timeoutdict import TimeoutDict 

16 

17 

18def _extract_block_key(message): 

19 """Extract a key that hashes equally for all blocks of a blockwise 

20 operation from a request message. 

21 

22 See discussion at <https://mailarchive.ietf.org/arch/msg/core/I-6LzAL6lIUVDA6_g9YM3Zjhg8E>. 

23 """ 

24 

25 return ( 

26 message.remote.blockwise_key, 

27 message.code, 

28 message.get_cache_key( 

29 [ 

30 OptionNumber.BLOCK1, 

31 OptionNumber.BLOCK2, 

32 OptionNumber.OBSERVE, 

33 ] 

34 ), 

35 ) 

36 

37 

38class ContinueException(ConstructionRenderableError): 

39 """Not an error in the CoAP sense, but an error in the processing sense, 

40 indicating that no complete request message is available for processing. 

41 

42 It reflects back the request's block1 option when rendered. 

43 """ 

44 

45 def __init__(self, block1): 

46 self.block1 = block1 

47 

48 def to_message(self): 

49 m = super().to_message() 

50 m.opt.block1 = self.block1 

51 return m 

52 

53 code = codes.CONTINUE 

54 

55 

56class IncompleteException(ConstructionRenderableError): 

57 code = codes.REQUEST_ENTITY_INCOMPLETE 

58 

59 

60class Block1Spool: 

61 def __init__(self): 

62 # FIXME: introduce an actual parameter here 

63 self._assemblies = TimeoutDict(numbers.TransportTuning().MAX_TRANSMIT_WAIT) 

64 

65 def feed_and_take(self, req: Message) -> Message: 

66 """Assemble the request into the spool. This either produces a 

67 reassembled request message, or raises either a Continue or a Request 

68 Entity Incomplete exception. 

69 

70 Requests without block1 are simply passed through.""" 

71 

72 if req.opt.block1 is None: 

73 return req 

74 

75 block_key = _extract_block_key(req) 

76 

77 if req.opt.block1.block_number == 0: 

78 # silently discarding any old incomplete operation 

79 self._assemblies[block_key] = req 

80 else: 

81 try: 

82 self._assemblies[block_key]._append_request_block(req) 

83 except KeyError: 

84 # KeyError: Received unmatched blockwise response 

85 # ValueError: Failed to assemble -- gaps or overlaps in data 

86 raise IncompleteException from None 

87 

88 if req.opt.block1.more: 

89 raise ContinueException(req.opt.block1) 

90 else: 

91 return self._assemblies[block_key] 

92 # which happens to carry the last block's block1 option 

93 

94 

95class Block2Cache: 

96 """A cache of responses to a give block key. 

97 

98 Use this when result rendering is expensive, not idempotent or has varying 

99 output -- otherwise it's often better to calculate the full response again 

100 and serve chunks. 

101 """ 

102 

103 def __init__(self): 

104 # FIXME: introduce an actual parameter here 

105 self._completes = TimeoutDict(numbers.TransportTuning().MAX_TRANSMIT_WAIT) 

106 

107 async def extract_or_insert( 

108 self, req: Message, response_builder: Callable[[], Awaitable[Message]] 

109 ): 

110 """Given a request message, 

111 

112 * if it is querying a particular block, look it up in the cache or 

113 raise Request Entity Incomplete. 

114 * otherwise, 

115 * await the response builder 

116 * return the response if it doesn't need chunking, or 

117 * return the first chunk and store it for later use 

118 

119 """ 

120 block_key = _extract_block_key(req) 

121 

122 if req.opt.block2 is None or req.opt.block2.block_number == 0: 

123 assembled = await response_builder() 

124 else: 

125 try: 

126 assembled = self._completes[block_key] 

127 except KeyError: 

128 raise IncompleteException from None 

129 

130 if ( 

131 len(assembled.payload) > req.remote.maximum_payload_size 

132 or req.opt.block2 is not None 

133 and len(assembled.payload) > req.opt.block2.size 

134 ): 

135 self._completes[block_key] = assembled 

136 

137 block2 = req.opt.block2 or BlockOption.BlockwiseTuple( 

138 0, 0, req.remote.maximum_block_size_exp 

139 ) 

140 return assembled._extract_block( 

141 block2.block_number, 

142 block2.size_exponent, 

143 req.remote.maximum_payload_size, 

144 ) 

145 else: 

146 return assembled