prosody
6632acc96cf6
util/dbuffer.lua

Find changesets by keywords by author, files, the commit message, revision number or hash, or revset expression.

Software / code / prosody

Comparison

util/dbuffer.lua @ 11104:6632acc96cf6 0.11

util.dbuffer: Fix :sub() not working with partially-consumed chunks (thanks Zash for test case) This also appears to fix some bugs with chunk-encoded streams in net.http.parser.
author Matthew Wild <mwild1@gmail.com>
date Mon, 24 Aug 2020 16:18:13 +0100
child 11114:6a608ecb3471
child 11156:a8ef69f7fc35
comparison
equal deleted inserted replaced
11103:73b8aaf55775 11104:6632acc96cf6
1 local queue = require "util.queue";
2
3 local dbuffer_methods = {};
4 local dynamic_buffer_mt = { __index = dbuffer_methods };
5
6 function dbuffer_methods:write(data)
7 if self.max_size and #data + self._length > self.max_size then
8 return nil;
9 end
10 local ok = self.items:push(data);
11 if not ok then
12 self:collapse();
13 ok = self.items:push(data);
14 end
15 if not ok then
16 return nil;
17 end
18 self._length = self._length + #data;
19 return true;
20 end
21
22 function dbuffer_methods:read_chunk(requested_bytes)
23 local chunk, consumed = self.items:peek(), self.front_consumed;
24 if not chunk then return; end
25 local chunk_length = #chunk;
26 local remaining_chunk_length = chunk_length - consumed;
27 if not requested_bytes then
28 requested_bytes = remaining_chunk_length;
29 end
30 if remaining_chunk_length <= requested_bytes then
31 self.front_consumed = 0;
32 self._length = self._length - remaining_chunk_length;
33 self.items:pop();
34 assert(#chunk:sub(consumed + 1, -1) == remaining_chunk_length);
35 return chunk:sub(consumed + 1, -1), remaining_chunk_length;
36 end
37 local end_pos = consumed + requested_bytes;
38 self.front_consumed = end_pos;
39 self._length = self._length - requested_bytes;
40 assert(#chunk:sub(consumed + 1, end_pos) == requested_bytes);
41 return chunk:sub(consumed + 1, end_pos), requested_bytes;
42 end
43
44 function dbuffer_methods:read(requested_bytes)
45 local chunks;
46
47 if requested_bytes and requested_bytes > self._length then
48 return nil;
49 end
50
51 local chunk, read_bytes = self:read_chunk(requested_bytes);
52 if not requested_bytes then
53 return chunk;
54 elseif chunk then
55 requested_bytes = requested_bytes - read_bytes;
56 if requested_bytes == 0 then -- Already read everything we need
57 return chunk;
58 end
59 chunks = {};
60 else
61 return nil;
62 end
63
64 -- Need to keep reading more chunks
65 while chunk do
66 table.insert(chunks, chunk);
67 if requested_bytes > 0 then
68 chunk, read_bytes = self:read_chunk(requested_bytes);
69 requested_bytes = requested_bytes - read_bytes;
70 else
71 break;
72 end
73 end
74
75 return table.concat(chunks);
76 end
77
78 function dbuffer_methods:discard(requested_bytes)
79 if requested_bytes > self._length then
80 return nil;
81 end
82
83 local chunk, read_bytes = self:read_chunk(requested_bytes);
84 if chunk then
85 requested_bytes = requested_bytes - read_bytes;
86 if requested_bytes == 0 then -- Already read everything we need
87 return true;
88 end
89 else
90 return nil;
91 end
92
93 while chunk do
94 if requested_bytes > 0 then
95 chunk, read_bytes = self:read_chunk(requested_bytes);
96 requested_bytes = requested_bytes - read_bytes;
97 else
98 break;
99 end
100 end
101 return true;
102 end
103
104 function dbuffer_methods:sub(i, j)
105 if j == nil then
106 j = -1;
107 end
108 if j < 0 then
109 j = self._length + (j+1);
110 end
111 if i < 0 then
112 i = self._length + (i+1);
113 end
114 if i < 1 then
115 i = 1;
116 end
117 if j > self._length then
118 j = self._length;
119 end
120 if i > j then
121 return "";
122 end
123
124 self:collapse(j);
125
126 return self.items:peek():sub(self.front_consumed+1):sub(i, j);
127 end
128
129 function dbuffer_methods:byte(i, j)
130 i = i or 1;
131 j = j or i;
132 return string.byte(self:sub(i, j), 1, -1);
133 end
134
135 function dbuffer_methods:length()
136 return self._length;
137 end
138 dynamic_buffer_mt.__len = dbuffer_methods.length; -- support # operator
139
140 function dbuffer_methods:collapse(bytes)
141 bytes = bytes or self._length;
142
143 local front_chunk = self.items:peek();
144
145 if not front_chunk or #front_chunk - self.front_consumed >= bytes then
146 return;
147 end
148
149 local front_chunks = { front_chunk:sub(self.front_consumed+1) };
150 local front_bytes = #front_chunks[1];
151
152 while front_bytes < bytes do
153 self.items:pop();
154 local chunk = self.items:peek();
155 front_bytes = front_bytes + #chunk;
156 table.insert(front_chunks, chunk);
157 end
158 self.items:replace(table.concat(front_chunks));
159 self.front_consumed = 0;
160 end
161
162 local function new(max_size, max_chunks)
163 if max_size and max_size <= 0 then
164 return nil;
165 end
166 return setmetatable({
167 front_consumed = 0;
168 _length = 0;
169 max_size = max_size;
170 items = queue.new(max_chunks or 32);
171 }, dynamic_buffer_mt);
172 end
173
174 return {
175 new = new;
176 };