]> git.sesse.net Git - fjl/blob - bytesource.c
Fix a bug where we could return too much data from the byte source. Add test.
[fjl] / bytesource.c
1 #include <stdbool.h>
2 #include <string.h>
3 #include <assert.h>
4
5 #include "choice.h"
6 #include "bytesource.h"
7
8 #define MARKER_CHAR 0xff
9 #define STUFF_MARKER 0x00
10
11 void init_byte_source(struct byte_source* source, raw_input_func_t* input_func, void* userdata)
12 {
13         // TODO: should this size be a different constant?
14         memset(source, 0, sizeof(*source));
15         source->bytes = (uint8_t*)malloc(BYTESOURCE_CHUNK_SIZE);
16         source->input_func = input_func;
17         source->userdata = userdata;
18 }
19
20 uint8_t byte_source_read_marker(struct byte_source* src)
21 {
22         // Refill until we have at least two bytes or EOF.
23         while (src->bytes_available < 2) {
24                 const unsigned bytes_to_read = 2 - src->bytes_available;
25                 const ssize_t bytes_read =
26                         (*src->input_func)(src->userdata,
27                                            src->bytes + src->bytes_available,
28                                            bytes_to_read);
29                 assert(bytes_read >= -1);
30                 assert(bytes_read <= (ssize_t)bytes_to_read);
31                 
32                 if (bytes_read == -1 || bytes_read == 0) {
33                         return 0x00;
34                 }
35
36                 src->bytes_available += bytes_read;
37         }
38
39         assert(src->bytes_available >= 2);
40         if (src->bytes[0] != MARKER_CHAR || src->bytes[1] == STUFF_MARKER) {
41                 return 0x00;
42         }
43
44         uint8_t ret = src->bytes[1];
45         memmove(src->bytes, src->bytes + 2, src->bytes_available - 2);
46         src->bytes_available -= 2;
47
48         return ret;
49 }
50
51 ssize_t byte_source_input_func(void* source, uint8_t* buf, size_t len)
52 {
53         struct byte_source* src = (struct byte_source*)source;
54
55         // If there's no data in the buffer (or only a partial marker), we have
56         // to read in more from our upstream src.
57         while (src->bytes_available == 0 ||
58                (src->bytes_available == 1 && src->bytes[0] == MARKER_CHAR)) {
59                 const unsigned space_left = BYTESOURCE_CHUNK_SIZE - src->bytes_available;
60                 const unsigned missing_data = len - src->bytes_available;
61                 const size_t bytes_to_read = (missing_data > space_left ? space_left : missing_data);
62                 assert(bytes_to_read <= BYTESOURCE_CHUNK_SIZE);
63                 const ssize_t bytes_read =
64                         (*src->input_func)(src->userdata,
65                                            src->bytes + src->bytes_available,
66                                            bytes_to_read);
67                 assert(bytes_read >= -1);
68                 assert(bytes_read <= (ssize_t)bytes_to_read);
69                 
70                 if (bytes_read == -1) {
71                         return -1;
72                 } else if (bytes_read == 0) {
73                         if (src->bytes_available == 1) {
74                                 // EOF in the middle of a marker => read error
75                                 return -1;
76                         } else {
77                                 assert(src->bytes_available == 0);
78                                 return 0;
79                         }
80                 }
81
82                 src->bytes_available += bytes_read;
83         }
84         
85         // Now unstuff as much as we can. First of all, if there's a 0xFF at the
86         // end of the buffer, we don't include it this time; the unstuff function
87         // will only give us an error since it can't decide if it's a marker or
88         // a stuffed 0xFF.
89         unsigned bytes_to_unstuff = src->bytes_available;
90         bool end_marker = false;
91         assert(bytes_to_unstuff > 0);
92         if (src->bytes[bytes_to_unstuff - 1] == 0xff) {
93                 --bytes_to_unstuff;
94                 end_marker = true;
95         }
96
97         int unstuffed_bytes = (*unstuff_choice)(buf, src->bytes, bytes_to_unstuff);
98         assert(unstuffed_bytes != 0);
99         if (unstuffed_bytes > 0) {
100                 // Fast path: No markers in the data. We can basically just
101                 // return it.
102                 if (end_marker) {
103                         src->bytes_available = 1;
104                         src->bytes[0] = 0xff;
105                 } else {
106                         src->bytes_available = 0;
107                         src->bytes[0] = 0;
108                 }
109                 return unstuffed_bytes;
110         }
111
112         // Slow path: There was a marker in the data. Unstuff manually until
113         // we hit the marker, then return that.
114         assert(unstuffed_bytes == -1);
115         unsigned bytes_read;
116         unsigned bytes_written = 0;
117         for (bytes_read = 0; bytes_read < src->bytes_available; ++bytes_read) {
118                 buf[bytes_written++] = src->bytes[bytes_read];
119                 if (src->bytes[bytes_read] != MARKER_CHAR) {
120                         continue;
121                 }
122
123                 assert(bytes_read < src->bytes_available);
124                 if (src->bytes[bytes_read + 1] == STUFF_MARKER) {
125                         // Skip the stuff byte.
126                         ++bytes_read;
127                         continue;
128                 } else {
129                         // OK, this is our marker.
130                         break;
131                 }       
132         }
133
134         memmove(src->bytes, src->bytes + bytes_read, src->bytes_available - bytes_read);
135         src->bytes_available -= bytes_read;
136         assert(bytes_written >= 1);
137         return bytes_written - 1;
138 }