okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

Aly Raffaufrefactor(runtime): decouple core storage from work9dd3a56

main
branchesmainall refs →
6.5 KiB213 linesraw
1//// The kernel reads image headers itself so `show_image` can refuse an image
2//// over the model's edge, or one cut short, at the call, and the daemon reads
3//// them again as the authority. Two parsers in two languages drift apart
4//// silently, and E2E only ever sends one well-formed PNG, so these fixtures
5//// hold both to one answer.
6
7import albedo/daemon/image
8import albedo/openai_api/types
9
10import albedo/daemon/store
11import albedo/harness/extensions/python/kernel as python
12import gleam/bit_array
13import gleam/int
14import gleam/json
15import gleam/list
16import gleam/option.{type Option, None, Some}
17import gleam/string
18import gleeunit/should
19import harness/kernel_fixture
20
21/// Name, canonical base64, and the MIME type and size its header says.
22const fixtures = [
23  #(
24    "png",
25    "iVBORw0KGgoAAAANSUhEUgAACQAAAAS8CAIAAABW4yLdAAAAAElFTkSuQmCC",
26    Some(#("image/png", 2304, 1212)),
27  ),
28  #(
29    "png bytes after iend",
30    "iVBORw0KGgoAAAANSUhEUgAACQAAAAS8CAIAAABW4yLdAAAAAElFTkSuQmCCdGFpbA==",
31    Some(#("image/png", 2304, 1212)),
32  ),
33  #(
34    "png stops before iend",
35    "iVBORw0KGgoAAAANSUhEUgAACQAAAAS8CAIAAABW4yLd",
36    None,
37  ),
38  #("png header only", "iVBORw0KGgoAAAANSUhEUgAACQAAAAS8CAIAAAA=", None),
39  #("png zero width", "iVBORw0KGgoAAAANSUhEUgAAAAAAAAAFCAIAAAA=", None),
40  #("png truncated", "iVBORw0KGgoAAAANSUhEUgAAAAM=", None),
41  #(
42    "jpeg baseline",
43    "/9j/4AAQSkZJRgABAQAAAQABAAD/wAAICAHgAoAD/9k=",
44    Some(#("image/jpeg", 640, 480)),
45  ),
46  #(
47    "jpeg progressive",
48    "/9j/wgAICAu4D6AD/9k=",
49    Some(#("image/jpeg", 4000, 3000)),
50  ),
51  #(
52    "jpeg fill bytes and restart",
53    "/9j/0P///+AAEEpGSUYAAQEAAAEAAQAAAAD/wQAICAAJAAcD/9k=",
54    Some(#("image/jpeg", 7, 9)),
55  ),
56  #(
57    "jpeg without end of image",
58    "/9j/4AAQSkZJRgABAQAAAQABAAD/wAAICAHgAoAD",
59    None,
60  ),
61  #("jpeg scan before frame", "/9j/2gAEAAD/wAAICAABAAED", None),
62  #("jpeg truncated segment", "/9j/4QAWeHh4eHh4", None),
63  #("jpeg zero height", "/9j/wAAICAAAAAoD", None),
64  #("jpeg short length", "/9j/4AAB/8AACAgAAgACAw==", None),
65  #(
66    "webp vp8x",
67    "UklGRhYAAABXRUJQVlA4WAoAAAAAAAAAtwsAEwAA",
68    Some(#("image/webp", 3000, 20)),
69  ),
70  #(
71    "webp vp8l",
72    "UklGRhIAAABXRUJQVlA4TAYAAAAvAMgDAAA=",
73    Some(#("image/webp", 2049, 16)),
74  ),
75  #(
76    "webp vp8",
77    "UklGRhYAAABXRUJQVlA4IAoAAAAAAACdASqABzgE",
78    Some(#("image/webp", 1920, 1080)),
79  ),
80  #(
81    "webp chunk before image",
82    "UklGRh4AAABXRUJQSUNDUAMAAABhYmMAVlA4TAYAAAAvBEABAAA=",
83    Some(#("image/webp", 5, 6)),
84  ),
85  #(
86    "webp vp8 without start code",
87    "UklGRhYAAABXRUJQVlA4IAoAAAAAAAAAAAAFAAUA",
88    None,
89  ),
90  #("webp riff size mismatch", "UklGRhIAAABXRUJQVlA4TAYAAAAvBEABAAAAAA==", None),
91  #("webp no image chunk", "UklGRhAAAABXRUJQRVhJRgQAAABhYmNk", None),
92  #("garbage", "bm90IGFuIGltYWdlIGF0IGFsbA==", None),
93]
94
95pub fn the_kernel_reads_image_headers_the_way_the_daemon_does_test() -> Nil {
96  let expected = list.map(fixtures, fn(fixture) { #(fixture.0, fixture.2) })
97  list.map(fixtures, fn(fixture) {
98    #(fixture.0, case image.inspect(fixture.1) {
99      Ok(#(mime, width, height, _)) -> Some(#(mime, width, height))
100      Error(_) -> None
101    })
102  })
103  |> should.equal(expected)
104
105  let assert Ok(store) = store.start(":memory:", "")
106  let assert Ok(kernel) = kernel_fixture.local(store, "/tmp")
107  let encoded = json.array(fixtures, fn(fixture) { json.string(fixture.1) })
108  let assert Ok(read) =
109    python.execute(
110      kernel,
111      "headers",
112      "import base64, json\nfrom albedo_capture import image_size as size\n"
113        <> "for data in json.loads("
114        <> string.inspect(json.to_string(encoded))
115        <> "):\n    print(json.dumps(size(base64.b64decode(data))))",
116      10_000,
117    )
118  list.zip(
119    list.map(fixtures, fn(fixture) { fixture.0 }),
120    string.split(string.trim_end(read.output), "\n"),
121  )
122  |> should.equal(
123    list.map(expected, fn(pair) { #(pair.0, header_json(pair.1)) }),
124  )
125  let assert Ok(_) = python.stop(kernel)
126  store.close(store)
127}
128
129fn header_json(header: Option(#(String, Int, Int))) -> String {
130  case header {
131    Some(#(mime, width, height)) ->
132      "[\""
133      <> mime
134      <> "\", "
135      <> int.to_string(width)
136      <> ", "
137      <> int.to_string(height)
138      <> "]"
139    None -> "null"
140  }
141}
142
143@external(erlang, "albedo_images", "encode_base64")
144fn encode_base64(bytes: BitArray) -> String
145
146/// A whole 2x3 PNG: signature, IHDR, and IEND, in 45 bytes.
147const png = "iVBORw0KGgoAAAANSUhEUgAAAAIAAAADCAIAAAA2iEnWAAAAAElFTkSuQmCC"
148
149pub fn canonical_base64_and_metadata_must_agree_test() -> Nil {
150  image.inspect(png <> "AA==") |> should.equal(Ok(#("image/png", 2, 3, 46)))
151  image.inspect(png <> "AB==") |> should.equal(Error(Nil))
152  image.inspect(png <> "AAA=") |> should.equal(Ok(#("image/png", 2, 3, 47)))
153  image.inspect(png <> "AAB=") |> should.equal(Error(Nil))
154  image.inspect(png <> "AA==AAAA") |> should.equal(Error(Nil))
155  image.valid_payload("image/png", png, 2, 3, 45) |> should.be_true
156  image.valid_payload("image/png", png, 2, 4, 45) |> should.be_false
157  image.valid_payload("image/jpeg", png, 2, 3, 45) |> should.be_false
158  image.valid_payload("image/png", png, 2, 3, 46) |> should.be_false
159}
160
161pub fn late_jpeg_headers_and_webp_riff_sizes_require_full_decode_test() -> Nil {
162  let padding = bit_array.from_string(string.repeat("x", 60_000))
163  let jpeg = <<
164    0xFF,
165    0xD8,
166    0xFF,
167    0xE0,
168    60_002:size(16),
169    padding:bits,
170    0xFF,
171    0xC0,
172    8:size(16),
173    8,
174    3:size(16),
175    2:size(16),
176    3,
177  >>
178  let data = encode_base64(jpeg)
179  image.valid_payload("image/jpeg", data, 2, 3, bit_array.byte_size(jpeg))
180  |> should.be_true
181  let chunks = <<
182    "JUNK":utf8,
183    60_000:little-size(32),
184    padding:bits,
185    "VP8L":utf8,
186    6:little-size(32),
187    0x2F,
188    81_924:little-size(32),
189    0,
190  >>
191  let size = 4 + bit_array.byte_size(chunks)
192  let webp = <<"RIFF":utf8, size:little-size(32), "WEBP":utf8, chunks:bits>>
193  image.valid_payload(
194    "image/webp",
195    encode_base64(webp),
196    5,
197    6,
198    bit_array.byte_size(webp),
199  )
200  |> should.be_true
201  let bad_size = <<"RIFF":utf8, 4:little-size(32), "WEBP":utf8, chunks:bits>>
202  image.inspect(encode_base64(bad_size)) |> should.equal(Error(Nil))
203}
204
205pub fn canonical_payload_size_enforces_the_decoded_limit_test() -> Nil {
206  let assert Ok(whole) = bit_array.base64_decode(png)
207  let tail = types.max_image_bytes - bit_array.byte_size(whole)
208  let padding = bit_array.from_string(string.repeat("x", tail))
209  let bytes = <<whole:bits, padding:bits>>
210  image.inspect(encode_base64(bytes))
211  |> should.equal(Ok(#("image/png", 2, 3, types.max_image_bytes)))
212  image.inspect(encode_base64(<<bytes:bits, 0>>)) |> should.equal(Error(Nil))
213}