jpeg_io.rs (45636B)
1 // Copyright 2022 Adobe. All rights reserved. 2 // This file is licensed to you under the Apache License, 3 // Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) 4 // or the MIT license (http://opensource.org/licenses/MIT), 5 // at your option. 6 7 // Unless required by applicable law or agreed to in writing, 8 // this software is distributed on an "AS IS" BASIS, WITHOUT 9 // WARRANTIES OR REPRESENTATIONS OF ANY KIND, either express or 10 // implied. See the LICENSE-MIT and LICENSE-APACHE files for the 11 // specific language governing permissions and limitations under 12 // each license. 13 14 use std::{ 15 collections::HashMap, 16 fs::File, 17 io::{BufReader, Cursor, Write}, 18 path::*, 19 }; 20 21 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; 22 use img_parts::{ 23 jpeg::{ 24 markers::{self, APP0, APP15, COM, DQT, DRI, P, RST0, RST7, SOF0, SOF15, SOS, Z}, 25 Jpeg, JpegSegment, 26 }, 27 Bytes, DynImage, 28 }; 29 use serde_bytes::ByteBuf; 30 use tempfile::Builder; 31 32 use crate::{ 33 assertions::{BoxMap, C2PA_BOXHASH}, 34 asset_io::{ 35 rename_or_copy, AssetBoxHash, AssetIO, CAIRead, CAIReadWrite, CAIReader, CAIWriter, 36 ComposedManifestRef, HashBlockObjectType, HashObjectPositions, RemoteRefEmbed, 37 RemoteRefEmbedType, 38 }, 39 error::{Error, Result}, 40 utils::xmp_inmemory_utils::{add_provenance, MIN_XMP}, 41 }; 42 43 static SUPPORTED_TYPES: [&str; 3] = ["jpg", "jpeg", "image/jpeg"]; 44 45 const XMP_SIGNATURE: &[u8] = b"http://ns.adobe.com/xap/1.0/"; 46 const XMP_SIGNATURE_BUFFER_SIZE: usize = XMP_SIGNATURE.len() + 1; // skip null or space char at end 47 48 const MAX_JPEG_MARKER_SIZE: usize = 64000; // technically it's 64K but a bit smaller is fine 49 50 const C2PA_MARKER: [u8; 4] = [0x63, 0x32, 0x70, 0x61]; 51 52 fn vec_compare(va: &[u8], vb: &[u8]) -> bool { 53 (va.len() == vb.len()) && // zip stops at the shortest 54 va.iter() 55 .zip(vb) 56 .all(|(a,b)| a == b) 57 } 58 59 // Return contents of APP1 segment if it is an XMP segment. 60 fn extract_xmp(seg: &JpegSegment) -> Option<String> { 61 let contents = seg.contents(); 62 if contents.starts_with(XMP_SIGNATURE) { 63 let rest = contents.slice(XMP_SIGNATURE_BUFFER_SIZE..); 64 String::from_utf8(rest.to_vec()).ok() 65 } else { 66 None 67 } 68 } 69 70 // Extract XMP from bytes. 71 fn xmp_from_bytes(asset_bytes: &[u8]) -> Option<String> { 72 if let Ok(jpeg) = Jpeg::from_bytes(Bytes::copy_from_slice(asset_bytes)) { 73 let segs = jpeg.segments_by_marker(markers::APP1); 74 let xmp: String = segs.filter_map(extract_xmp).collect(); 75 Some(xmp) 76 } else { 77 None 78 } 79 } 80 81 fn add_required_segs_to_stream( 82 input_stream: &mut dyn CAIRead, 83 output_stream: &mut dyn CAIReadWrite, 84 ) -> Result<()> { 85 let mut buf: Vec<u8> = Vec::new(); 86 input_stream.rewind()?; 87 input_stream.read_to_end(&mut buf).map_err(Error::IoError)?; 88 input_stream.rewind()?; 89 90 let dimg_opt = DynImage::from_bytes(buf.into()) 91 .map_err(|_err| Error::InvalidAsset("Could not parse input JPEG".to_owned()))?; 92 93 if let Some(DynImage::Jpeg(jpeg)) = dimg_opt { 94 // check for JUMBF Seg 95 let cai_app11 = get_cai_segments(&jpeg)?; // make sure we only check for C2PA segments 96 97 if cai_app11.is_empty() { 98 // create dummy JUMBF seg 99 let mut no_bytes: Vec<u8> = vec![0; 50]; // enough bytes to be valid 100 no_bytes.splice(16..20, C2PA_MARKER); // cai UUID signature 101 let aio = JpegIO {}; 102 aio.write_cai(input_stream, output_stream, &no_bytes)?; 103 } else { 104 // just clone 105 input_stream.rewind()?; 106 output_stream.rewind()?; 107 std::io::copy(input_stream, output_stream)?; 108 } 109 } else { 110 return Err(Error::UnsupportedType); 111 } 112 113 Ok(()) 114 } 115 116 // all cai specific segments 117 fn get_cai_segments(jpeg: &img_parts::jpeg::Jpeg) -> Result<Vec<usize>> { 118 let mut cai_segs: Vec<usize> = Vec::new(); 119 120 let segments = jpeg.segments(); 121 122 let mut cai_en: Vec<u8> = Vec::new(); 123 let mut cai_seg_cnt: u32 = 0; 124 125 for (i, segment) in segments.iter().enumerate() { 126 let raw_bytes = segment.contents(); 127 let seg_type = segment.marker(); 128 129 if raw_bytes.len() > 16 && seg_type == markers::APP11 { 130 // we need at least 16 bytes in each segment for CAI 131 let mut raw_vec = raw_bytes.to_vec(); 132 let _ci = raw_vec.as_mut_slice()[0..2].to_vec(); 133 let en = raw_vec.as_mut_slice()[2..4].to_vec(); 134 let mut z_vec = Cursor::new(raw_vec.as_mut_slice()[4..8].to_vec()); 135 let _z = z_vec.read_u32::<BigEndian>()?; 136 137 let is_cai_continuation = vec_compare(&cai_en, &en); 138 139 if cai_seg_cnt > 0 && is_cai_continuation { 140 cai_seg_cnt += 1; 141 cai_segs.push(i); 142 } else { 143 // check if this is a CAI JUMBF block 144 let jumb_type = &raw_vec.as_mut_slice()[24..28]; 145 let is_cai = vec_compare(&C2PA_MARKER, jumb_type); 146 if is_cai { 147 cai_segs.push(i); 148 cai_seg_cnt = 1; 149 cai_en.clone_from(&en); // store the identifier 150 } 151 } 152 } 153 } 154 155 Ok(cai_segs) 156 } 157 158 // delete cai segments 159 fn delete_cai_segments(jpeg: &mut img_parts::jpeg::Jpeg) -> Result<()> { 160 let cai_segs = get_cai_segments(jpeg)?; 161 let jpeg_segs = jpeg.segments_mut(); 162 163 // remove cai segments 164 for seg in cai_segs.iter().rev() { 165 jpeg_segs.remove(*seg); 166 } 167 Ok(()) 168 } 169 170 pub struct JpegIO {} 171 172 impl CAIReader for JpegIO { 173 fn read_cai(&self, asset_reader: &mut dyn CAIRead) -> Result<Vec<u8>> { 174 let mut buffer: Vec<u8> = Vec::new(); 175 176 let mut manifest_store_cnt = 0; 177 178 // load the bytes 179 let mut buf: Vec<u8> = Vec::new(); 180 181 asset_reader.rewind()?; 182 asset_reader.read_to_end(&mut buf).map_err(Error::IoError)?; 183 184 let dimg_opt = DynImage::from_bytes(buf.into()) 185 .map_err(|_err| Error::InvalidAsset("Could not parse input JPEG".to_owned()))?; 186 187 if let Some(dimg) = dimg_opt { 188 match dimg { 189 DynImage::Jpeg(jpeg) => { 190 let app11 = jpeg.segments_by_marker(markers::APP11); 191 let mut cai_en: Vec<u8> = Vec::new(); 192 let mut cai_seg_cnt: u32 = 0; 193 for segment in app11 { 194 let raw_bytes = segment.contents(); 195 if raw_bytes.len() > 16 { 196 // we need at least 16 bytes in each segment for CAI 197 let mut raw_vec = raw_bytes.to_vec(); 198 let _ci = raw_vec.as_mut_slice()[0..2].to_vec(); 199 let en = raw_vec.as_mut_slice()[2..4].to_vec(); 200 let mut z_vec = Cursor::new(raw_vec.as_mut_slice()[4..8].to_vec()); 201 let z = z_vec.read_u32::<BigEndian>()?; 202 203 let is_cai_continuation = vec_compare(&cai_en, &en); 204 205 if cai_seg_cnt > 0 && is_cai_continuation { 206 // make sure this is a cai segment for additional segments, 207 if z <= cai_seg_cnt { 208 // this a non contiguous segment with same "en" so a bad set of data 209 // reset and continue to search 210 cai_en = Vec::new(); 211 continue; 212 } 213 // take out LBox & TBox 214 buffer.append(&mut raw_vec.as_mut_slice()[16..].to_vec()); 215 216 cai_seg_cnt += 1; 217 } else if raw_vec.len() > 28 { 218 // must be at least 28 bytes for this to be a valid JUMBF box 219 // check if this is a CAI JUMBF block 220 let jumb_type = &raw_vec.as_mut_slice()[24..28]; 221 let is_cai = vec_compare(&C2PA_MARKER, jumb_type); 222 223 if is_cai { 224 if manifest_store_cnt == 1 { 225 return Err(Error::TooManyManifestStores); 226 } 227 228 buffer.append(&mut raw_vec.as_mut_slice()[8..].to_vec()); 229 cai_seg_cnt = 1; 230 cai_en.clone_from(&en); // store the identifier 231 232 manifest_store_cnt += 1; 233 } 234 } 235 } 236 } 237 } 238 _ => return Err(Error::InvalidAsset("Unknown image format".to_owned())), 239 }; 240 } else { 241 return Err(Error::UnsupportedType); 242 } 243 244 if buffer.is_empty() { 245 return Err(Error::JumbfNotFound); 246 } 247 248 Ok(buffer) 249 } 250 251 // Get XMP block 252 fn read_xmp(&self, asset_reader: &mut dyn CAIRead) -> Option<String> { 253 // load the bytes 254 let mut buf: Vec<u8> = Vec::new(); 255 match asset_reader.read_to_end(&mut buf) { 256 Ok(_) => xmp_from_bytes(&buf), 257 Err(_) => None, 258 } 259 } 260 } 261 262 impl CAIWriter for JpegIO { 263 fn write_cai( 264 &self, 265 input_stream: &mut dyn CAIRead, 266 output_stream: &mut dyn CAIReadWrite, 267 store_bytes: &[u8], 268 ) -> Result<()> { 269 let mut buf = Vec::new(); 270 // read the whole asset 271 input_stream.rewind()?; 272 input_stream.read_to_end(&mut buf).map_err(Error::IoError)?; 273 let mut jpeg = Jpeg::from_bytes(buf.into()).map_err(|_err| Error::EmbeddingError)?; 274 275 // remove existing CAI segments 276 delete_cai_segments(&mut jpeg)?; 277 278 let jumbf_len = store_bytes.len(); 279 let num_segments = (jumbf_len / MAX_JPEG_MARKER_SIZE) + 1; 280 let mut seg_chucks = store_bytes.chunks(MAX_JPEG_MARKER_SIZE); 281 282 for seg in 1..num_segments + 1 { 283 /* 284 If the size of the box payload is less than 2^32-8 bytes, 285 then all fields except the XLBox field, that is: Le, CI, En, Z, LBox and TBox, 286 shall be present in all JPEG XT marker segment representing this box, 287 regardless of whether the marker segments starts this box, 288 or continues a box started by a former JPEG XT Marker segment. 289 */ 290 // we need to prefix the JUMBF with the JPEG XT markers (ISO 19566-5) 291 // CI: JPEG extensions marker - JP 292 // En: Box Instance Number - 0x0001 293 // (NOTE: can be any unique ID, so we pick one that shouldn't conflict) 294 // Z: Packet sequence number - 0x00000001... 295 let ci = vec![0x4a, 0x50]; 296 let en = vec![0x02, 0x11]; 297 let z: u32 = u32::try_from(seg) 298 .map_err(|_| Error::InvalidAsset("Too many JUMBF segments".to_string()))?; //seg.to_be_bytes(); 299 300 let mut seg_data = Vec::new(); 301 seg_data.extend(ci); 302 seg_data.extend(en); 303 seg_data.extend(z.to_be_bytes()); 304 if seg > 1 { 305 // the LBox and TBox are already in the JUMBF 306 // but we need to duplicate them in all other segments 307 let lbox_tbox = &store_bytes[..8]; 308 seg_data.extend(lbox_tbox); 309 } 310 if seg_chucks.len() > 0 { 311 // make sure we have some... 312 if let Some(next_seg) = seg_chucks.next() { 313 seg_data.extend(next_seg); 314 } 315 } else { 316 seg_data.extend(store_bytes); 317 } 318 319 let seg_bytes = Bytes::from(seg_data); 320 let app11_segment = JpegSegment::new_with_contents(markers::APP11, seg_bytes); 321 jpeg.segments_mut().insert(seg, app11_segment); // we put this in the beginning... 322 } 323 324 output_stream.rewind()?; 325 jpeg.encoder() 326 .write_to(output_stream) 327 .map_err(|_err| Error::InvalidAsset("JPEG write error".to_owned()))?; 328 329 Ok(()) 330 } 331 332 fn get_object_locations_from_stream( 333 &self, 334 input_stream: &mut dyn CAIRead, 335 ) -> Result<Vec<HashObjectPositions>> { 336 let mut cai_en: Vec<u8> = Vec::new(); 337 let mut cai_seg_cnt: u32 = 0; 338 339 let mut positions: Vec<HashObjectPositions> = Vec::new(); 340 let mut curr_offset = 2; // start after JPEG marker 341 342 let output_vec: Vec<u8> = Vec::new(); 343 let mut output_stream = Cursor::new(output_vec); 344 // make sure the file has the required segments so we can generate all the required offsets 345 add_required_segs_to_stream(input_stream, &mut output_stream)?; 346 347 let buf: Vec<u8> = output_stream.into_inner(); 348 349 let dimg = DynImage::from_bytes(buf.into()) 350 .map_err(|e| Error::OtherError(Box::new(e)))? 351 .ok_or(Error::UnsupportedType)?; 352 353 match dimg { 354 DynImage::Jpeg(jpeg) => { 355 for seg in jpeg.segments() { 356 match seg.marker() { 357 markers::APP11 => { 358 // JUMBF marker 359 let raw_bytes = seg.contents(); 360 361 if raw_bytes.len() > 16 { 362 // we need at least 16 bytes in each segment for CAI 363 let mut raw_vec = raw_bytes.to_vec(); 364 let _ci = raw_vec.as_mut_slice()[0..2].to_vec(); 365 let en = raw_vec.as_mut_slice()[2..4].to_vec(); 366 367 let is_cai_continuation = vec_compare(&cai_en, &en); 368 369 if cai_seg_cnt > 0 && is_cai_continuation { 370 cai_seg_cnt += 1; 371 372 let v = HashObjectPositions { 373 offset: curr_offset, 374 length: seg.len_with_entropy(), 375 htype: HashBlockObjectType::Cai, 376 }; 377 positions.push(v); 378 } else { 379 // check if this is a CAI JUMBF block 380 let jumb_type = raw_vec.as_mut_slice()[24..28].to_vec(); 381 let is_cai = vec_compare(&C2PA_MARKER, &jumb_type); 382 if is_cai { 383 cai_seg_cnt = 1; 384 cai_en.clone_from(&en); // store the identifier 385 386 let v = HashObjectPositions { 387 offset: curr_offset, 388 length: seg.len_with_entropy(), 389 htype: HashBlockObjectType::Cai, 390 }; 391 392 positions.push(v); 393 } else { 394 // save other for completeness sake 395 let v = HashObjectPositions { 396 offset: curr_offset, 397 length: seg.len_with_entropy(), 398 htype: HashBlockObjectType::Other, 399 }; 400 positions.push(v); 401 } 402 } 403 } 404 } 405 markers::APP1 => { 406 // XMP marker or EXIF or Extra XMP 407 let v = HashObjectPositions { 408 offset: curr_offset, 409 length: seg.len_with_entropy(), 410 htype: HashBlockObjectType::Xmp, 411 }; 412 // todo: pick the app1 that is the xmp (not crucial as it gets hashed either way) 413 positions.push(v); 414 } 415 _ => { 416 // save other for completeness sake 417 let v = HashObjectPositions { 418 offset: curr_offset, 419 length: seg.len_with_entropy(), 420 htype: HashBlockObjectType::Other, 421 }; 422 423 positions.push(v); 424 } 425 } 426 curr_offset += seg.len_with_entropy(); 427 } 428 } 429 _ => return Err(Error::InvalidAsset("Unknown image format".to_owned())), 430 } 431 432 Ok(positions) 433 } 434 435 fn remove_cai_store_from_stream( 436 &self, 437 input_stream: &mut dyn CAIRead, 438 output_stream: &mut dyn CAIReadWrite, 439 ) -> Result<()> { 440 let mut buf = Vec::new(); 441 // read the whole asset 442 input_stream.rewind()?; 443 input_stream.read_to_end(&mut buf).map_err(Error::IoError)?; 444 let mut jpeg = Jpeg::from_bytes(buf.into()).map_err(|_err| Error::EmbeddingError)?; 445 446 // remove existing CAI segments 447 delete_cai_segments(&mut jpeg)?; 448 449 output_stream.rewind()?; 450 jpeg.encoder() 451 .write_to(output_stream) 452 .map_err(|_err| Error::InvalidAsset("JPEG write error".to_owned()))?; 453 454 Ok(()) 455 } 456 } 457 458 impl AssetIO for JpegIO { 459 fn read_cai_store(&self, asset_path: &Path) -> Result<Vec<u8>> { 460 let mut f = File::open(asset_path)?; 461 462 self.read_cai(&mut f) 463 } 464 465 fn save_cai_store(&self, asset_path: &std::path::Path, store_bytes: &[u8]) -> Result<()> { 466 let mut input_stream = std::fs::OpenOptions::new() 467 .read(true) 468 //.truncate(true) 469 .open(asset_path) 470 .map_err(Error::IoError)?; 471 472 let mut temp_file = Builder::new() 473 .prefix("c2pa_temp") 474 .rand_bytes(5) 475 .tempfile()?; 476 477 self.write_cai(&mut input_stream, &mut temp_file, store_bytes)?; 478 479 // copy temp file to asset 480 rename_or_copy(temp_file, asset_path) 481 } 482 483 fn get_object_locations(&self, asset_path: &Path) -> Result<Vec<HashObjectPositions>> { 484 let mut file = std::fs::OpenOptions::new() 485 .read(true) 486 .write(true) 487 .open(asset_path) 488 .map_err(Error::IoError)?; 489 490 self.get_object_locations_from_stream(&mut file) 491 } 492 493 fn remove_cai_store(&self, asset_path: &Path) -> Result<()> { 494 let input = std::fs::read(asset_path).map_err(Error::IoError)?; 495 496 let mut jpeg = Jpeg::from_bytes(input.into()).map_err(|_err| Error::EmbeddingError)?; 497 498 // remove existing CAI segments 499 delete_cai_segments(&mut jpeg)?; 500 501 // save updated file 502 let output = std::fs::OpenOptions::new() 503 .read(true) 504 .write(true) 505 .truncate(true) 506 .open(asset_path) 507 .map_err(Error::IoError)?; 508 509 jpeg.encoder() 510 .write_to(output) 511 .map_err(|_err| Error::InvalidAsset("JPEG write error".to_owned()))?; 512 513 Ok(()) 514 } 515 516 fn new(_asset_type: &str) -> Self 517 where 518 Self: Sized, 519 { 520 JpegIO {} 521 } 522 523 fn get_handler(&self, asset_type: &str) -> Box<dyn AssetIO> { 524 Box::new(JpegIO::new(asset_type)) 525 } 526 527 fn get_reader(&self) -> &dyn CAIReader { 528 self 529 } 530 531 fn get_writer(&self, asset_type: &str) -> Option<Box<dyn CAIWriter>> { 532 Some(Box::new(JpegIO::new(asset_type))) 533 } 534 535 fn remote_ref_writer_ref(&self) -> Option<&dyn RemoteRefEmbed> { 536 Some(self) 537 } 538 539 fn asset_box_hash_ref(&self) -> Option<&dyn AssetBoxHash> { 540 Some(self) 541 } 542 543 fn composed_data_ref(&self) -> Option<&dyn ComposedManifestRef> { 544 Some(self) 545 } 546 547 fn supported_types(&self) -> &[&str] { 548 &SUPPORTED_TYPES 549 } 550 } 551 552 impl RemoteRefEmbed for JpegIO { 553 #[allow(unused_variables)] 554 fn embed_reference( 555 &self, 556 asset_path: &Path, 557 embed_ref: crate::asset_io::RemoteRefEmbedType, 558 ) -> Result<()> { 559 match &embed_ref { 560 crate::asset_io::RemoteRefEmbedType::Xmp(_manifest_uri) => { 561 let mut file = std::fs::File::open(asset_path)?; 562 let mut temp = Cursor::new(Vec::new()); 563 self.embed_reference_to_stream(&mut file, &mut temp, embed_ref)?; 564 let mut output = std::fs::OpenOptions::new() 565 .read(true) 566 .write(true) 567 .truncate(true) 568 .open(asset_path) 569 .map_err(Error::IoError)?; 570 temp.set_position(0); 571 std::io::copy(&mut temp, &mut output).map_err(Error::IoError)?; 572 Ok(()) 573 } 574 crate::asset_io::RemoteRefEmbedType::StegoS(_) => Err(Error::UnsupportedType), 575 crate::asset_io::RemoteRefEmbedType::StegoB(_) => Err(Error::UnsupportedType), 576 crate::asset_io::RemoteRefEmbedType::Watermark(_) => Err(Error::UnsupportedType), 577 } 578 } 579 580 fn embed_reference_to_stream( 581 &self, 582 source_stream: &mut dyn CAIRead, 583 output_stream: &mut dyn CAIReadWrite, 584 embed_ref: RemoteRefEmbedType, 585 ) -> Result<()> { 586 match embed_ref { 587 crate::asset_io::RemoteRefEmbedType::Xmp(manifest_uri) => { 588 let mut buf = Vec::new(); 589 // read the whole asset 590 source_stream.rewind()?; 591 source_stream 592 .read_to_end(&mut buf) 593 .map_err(Error::IoError)?; 594 let mut jpeg = 595 Jpeg::from_bytes(buf.into()).map_err(|_err| Error::EmbeddingError)?; 596 597 // find any existing XMP segment and remember where it was 598 let mut xmp = MIN_XMP.to_string(); // default minimal XMP 599 let mut xmp_index = None; 600 let segments = jpeg.segments_mut(); 601 for (i, seg) in segments.iter().enumerate() { 602 if seg.marker() == markers::APP1 && seg.contents().starts_with(XMP_SIGNATURE) { 603 xmp = extract_xmp(seg).unwrap_or_else(|| xmp.clone()); 604 xmp_index = Some(i); 605 break; 606 } 607 } 608 // add provenance and JPEG XMP prefix 609 let xmp = format!( 610 "http://ns.adobe.com/xap/1.0/\0 {}", 611 add_provenance(&xmp, &manifest_uri)? 612 ); 613 let segment = JpegSegment::new_with_contents(markers::APP1, Bytes::from(xmp)); 614 // insert or add the segment 615 match xmp_index { 616 Some(i) => segments[i] = segment, 617 None => segments.insert(1, segment), 618 } 619 620 jpeg.encoder() 621 .write_to(output_stream) 622 .map_err(|_err| Error::InvalidAsset("JPEG write error".to_owned()))?; 623 Ok(()) 624 } 625 crate::asset_io::RemoteRefEmbedType::StegoS(_) => Err(Error::UnsupportedType), 626 crate::asset_io::RemoteRefEmbedType::StegoB(_) => Err(Error::UnsupportedType), 627 crate::asset_io::RemoteRefEmbedType::Watermark(_) => Err(Error::UnsupportedType), 628 } 629 } 630 } 631 632 const fn in_entropy(marker: u8) -> bool { 633 matches!(marker, RST0..=RST7 | Z) 634 } 635 636 // img-parts does not correctly return the true size of the SOS segment. This utility 637 // finds the correct break point for single image JPEGs. We will need a new JPEG decoder 638 // to handle those. Also this function can be removed if img-parts ever addresses this issue 639 // and support MPF JPEGs. 640 fn get_entropy_size(input_stream: &mut dyn CAIRead) -> Result<usize> { 641 // Search the entropy data looking for non entropy segment marker. The first valid seg marker before we hit 642 // end of the file. 643 644 let mut buf_reader = BufReader::new(input_stream); 645 646 let mut size = 0; 647 648 loop { 649 match buf_reader.read_u8() { 650 Ok(curr_byte) => { 651 if curr_byte == P { 652 let next_byte = buf_reader.read_u8()?; 653 654 if !in_entropy(next_byte) { 655 break; 656 } else { 657 size += 1; 658 } 659 } 660 size += 1; 661 } 662 Err(e) => return Err(Error::IoError(e)), 663 } 664 } 665 666 Ok(size) 667 } 668 669 const fn has_length(marker: u8) -> bool { 670 matches!(marker, RST0..=RST7 | APP0..=APP15 | SOF0..=SOF15 | SOS | COM | DQT | DRI) 671 } 672 673 fn get_seg_size(input_stream: &mut dyn CAIRead) -> Result<usize> { 674 let p = input_stream.read_u8()?; 675 let marker = if p == P { 676 input_stream.read_u8()? 677 } else { 678 return Err(Error::InvalidAsset( 679 "Cannot read segment marker".to_string(), 680 )); 681 }; 682 683 if has_length(marker) { 684 let val: usize = input_stream.read_u16::<BigEndian>()? as usize; 685 Ok(val + 2) 686 } else { 687 Ok(2) 688 } 689 } 690 691 fn make_box_maps(input_stream: &mut dyn CAIRead) -> Result<Vec<BoxMap>> { 692 let segment_names = HashMap::from([ 693 (0xe0u8, "APP0"), 694 (0xe1u8, "APP1"), 695 (0xe2u8, "APP2"), 696 (0xe3u8, "APP3"), 697 (0xe4u8, "APP4"), 698 (0xe5u8, "APP5"), 699 (0xe6u8, "APP6"), 700 (0xe7u8, "APP7"), 701 (0xe8u8, "APP8"), 702 (0xe9u8, "APP9"), 703 (0xeau8, "APP10"), 704 (0xebu8, "APP11"), 705 (0xecu8, "APP12"), 706 (0xedu8, "APP13"), 707 (0xeeu8, "APP14"), 708 (0xefu8, "APP15"), 709 (0xfeu8, "COM"), 710 (0xc4u8, "DHT"), 711 (0xdbu8, "DQT"), 712 (0xddu8, "DRI"), 713 (0xd9u8, "EOI"), 714 (0xd0u8, "RST0"), 715 (0xd1u8, "RST1"), 716 (0xd2u8, "RST2"), 717 (0xd3u8, "RST3"), 718 (0xd4u8, "RST4"), 719 (0xd5u8, "RST5"), 720 (0xd6u8, "RST6"), 721 (0xd7u8, "RST7"), 722 (0xc0u8, "SOF0"), 723 (0xc1u8, "SOF1"), 724 (0xc2u8, "SOF2"), 725 (0xd8u8, "SOI"), 726 (0xdau8, "SOS"), 727 (0xf0u8, "JPG0"), 728 (0xf1u8, "JPG1"), 729 (0xf2u8, "JPG2"), 730 (0xf3u8, "JPG3"), 731 (0xf4u8, "JPG4"), 732 (0xf5u8, "JPG5"), 733 (0xf6u8, "JPG6"), 734 (0xf7u8, "JPG7"), 735 (0xf8u8, "JPG8"), 736 (0xf9u8, "JPG9"), 737 (0xfau8, "JPG10"), 738 (0xfbu8, "JPG11"), 739 (0xfcu8, "JPG12"), 740 (0xfdu8, "JPG13"), 741 ]); 742 743 let mut box_maps = Vec::new(); 744 let mut cai_en: Vec<u8> = Vec::new(); 745 let mut cai_seg_cnt: u32 = 0; 746 let mut cai_index = 0; 747 748 input_stream.rewind()?; 749 750 let buf_reader = BufReader::new(input_stream); 751 let mut reader = jfifdump::Reader::new(buf_reader) 752 .map_err(|_e| Error::InvalidAsset("could not read JPEG segments".to_string()))?; 753 754 while let Ok(seg) = reader.next_segment() { 755 match seg.kind { 756 jfifdump::SegmentKind::Eoi => { 757 let bm = BoxMap { 758 names: vec!["EOI".to_string()], 759 alg: None, 760 hash: ByteBuf::from(Vec::new()), 761 pad: ByteBuf::from(Vec::new()), 762 range_start: seg.position, 763 range_len: 0, 764 }; 765 766 box_maps.push(bm); 767 } 768 jfifdump::SegmentKind::Soi => { 769 let bm = BoxMap { 770 names: vec!["SOI".to_string()], 771 alg: None, 772 hash: ByteBuf::from(Vec::new()), 773 pad: ByteBuf::from(Vec::new()), 774 range_start: seg.position, 775 range_len: 0, 776 }; 777 778 box_maps.push(bm); 779 } 780 jfifdump::SegmentKind::App { nr, data } if nr == 0x0b => { 781 let nr = nr | 0xe0; 782 783 // JUMBF marker 784 let raw_bytes = data; 785 786 if raw_bytes.len() > 16 { 787 // we need at least 16 bytes in each segment for CAI 788 let mut raw_vec = raw_bytes.to_vec(); 789 let _ci = raw_vec.as_mut_slice()[0..2].to_vec(); 790 let en = raw_vec.as_mut_slice()[2..4].to_vec(); 791 792 let is_cai_continuation = vec_compare(&cai_en, &en); 793 794 if cai_seg_cnt > 0 && is_cai_continuation { 795 cai_seg_cnt += 1; 796 797 let cai_bm = &mut box_maps[cai_index]; 798 cai_bm.range_len += raw_bytes.len() + 4; 799 } else { 800 // check if this is a CAI JUMBF block 801 let jumb_type = raw_vec.as_mut_slice()[24..28].to_vec(); 802 let is_cai = vec_compare(&C2PA_MARKER, &jumb_type); 803 if is_cai { 804 cai_seg_cnt = 1; 805 cai_en.clone_from(&en); // store the identifier 806 807 let c2pa_bm = BoxMap { 808 names: vec![C2PA_BOXHASH.to_string()], 809 alg: None, 810 hash: ByteBuf::from(Vec::new()), 811 pad: ByteBuf::from(Vec::new()), 812 range_start: seg.position, 813 range_len: raw_bytes.len() + 4, 814 }; 815 816 box_maps.push(c2pa_bm); 817 cai_index = box_maps.len() - 1; 818 } else { 819 let name = segment_names 820 .get(&nr) 821 .ok_or(Error::InvalidAsset("Unknown segment marker".to_owned()))?; 822 823 let bm = BoxMap { 824 names: vec![name.to_string()], 825 alg: None, 826 hash: ByteBuf::from(Vec::new()), 827 pad: ByteBuf::from(Vec::new()), 828 range_start: seg.position, 829 range_len: 0, 830 }; 831 832 box_maps.push(bm); 833 } 834 } 835 } 836 } 837 jfifdump::SegmentKind::App { nr, data } => { 838 let nr = nr | 0xe0; 839 let _data = data; 840 841 let name = segment_names 842 .get(&nr) 843 .ok_or(Error::InvalidAsset("Unknown segment marker".to_owned()))?; 844 845 let bm = BoxMap { 846 names: vec![name.to_string()], 847 alg: None, 848 hash: ByteBuf::from(Vec::new()), 849 pad: ByteBuf::from(Vec::new()), 850 range_start: seg.position, 851 range_len: 0, 852 }; 853 854 box_maps.push(bm); 855 } 856 jfifdump::SegmentKind::App0Jfif(_) => { 857 let bm = BoxMap { 858 names: vec!["APP0".to_string()], 859 alg: None, 860 hash: ByteBuf::from(Vec::new()), 861 pad: ByteBuf::from(Vec::new()), 862 range_start: seg.position, 863 range_len: 0, 864 }; 865 866 box_maps.push(bm); 867 } 868 jfifdump::SegmentKind::Dqt(_) => { 869 let bm = BoxMap { 870 names: vec!["DQT".to_string()], 871 alg: None, 872 hash: ByteBuf::from(Vec::new()), 873 pad: ByteBuf::from(Vec::new()), 874 range_start: seg.position, 875 range_len: 0, 876 }; 877 878 box_maps.push(bm); 879 } 880 jfifdump::SegmentKind::Dht(_) => { 881 let bm = BoxMap { 882 names: vec!["DHT".to_string()], 883 alg: None, 884 hash: ByteBuf::from(Vec::new()), 885 pad: ByteBuf::from(Vec::new()), 886 range_start: seg.position, 887 range_len: 0, 888 }; 889 890 box_maps.push(bm); 891 } 892 jfifdump::SegmentKind::Dac(_) => { 893 let bm = BoxMap { 894 names: vec!["DAC".to_string()], 895 alg: None, 896 hash: ByteBuf::from(Vec::new()), 897 pad: ByteBuf::from(Vec::new()), 898 range_start: seg.position, 899 range_len: 0, 900 }; 901 902 box_maps.push(bm); 903 } 904 jfifdump::SegmentKind::Frame(f) => { 905 let name = segment_names 906 .get(&f.sof) 907 .ok_or(Error::InvalidAsset("Unknown segment marker".to_owned()))?; 908 909 let bm = BoxMap { 910 names: vec![name.to_string()], 911 alg: None, 912 hash: ByteBuf::from(Vec::new()), 913 pad: ByteBuf::from(Vec::new()), 914 range_start: seg.position, 915 range_len: 0, 916 }; 917 918 box_maps.push(bm); 919 } 920 jfifdump::SegmentKind::Scan(_s) => { 921 let bm = BoxMap { 922 names: vec!["SOS".to_string()], 923 alg: None, 924 hash: ByteBuf::from(Vec::new()), 925 pad: ByteBuf::from(Vec::new()), 926 range_start: seg.position, 927 range_len: 0, 928 }; 929 930 box_maps.push(bm); 931 } 932 jfifdump::SegmentKind::Dri(_) => { 933 let bm = BoxMap { 934 names: vec!["DRI".to_string()], 935 alg: None, 936 hash: ByteBuf::from(Vec::new()), 937 pad: ByteBuf::from(Vec::new()), 938 range_start: seg.position, 939 range_len: 0, 940 }; 941 942 box_maps.push(bm); 943 } 944 jfifdump::SegmentKind::Rst(_r) => (), 945 jfifdump::SegmentKind::Comment(_) => { 946 let bm = BoxMap { 947 names: vec!["COM".to_string()], 948 alg: None, 949 hash: ByteBuf::from(Vec::new()), 950 pad: ByteBuf::from(Vec::new()), 951 range_start: seg.position, 952 range_len: 0, 953 }; 954 955 box_maps.push(bm); 956 } 957 jfifdump::SegmentKind::Unknown { marker, data: _ } => { 958 let name = segment_names 959 .get(&marker) 960 .ok_or(Error::InvalidAsset("Unknown segment marker".to_owned()))?; 961 962 let bm = BoxMap { 963 names: vec![name.to_string()], 964 alg: None, 965 hash: ByteBuf::from(Vec::new()), 966 pad: ByteBuf::from(Vec::new()), 967 range_start: seg.position, 968 range_len: 0, 969 }; 970 971 box_maps.push(bm); 972 } 973 } 974 } 975 976 Ok(box_maps) 977 } 978 979 impl AssetBoxHash for JpegIO { 980 fn get_box_map(&self, input_stream: &mut dyn CAIRead) -> Result<Vec<BoxMap>> { 981 let mut box_maps = make_box_maps(input_stream)?; 982 983 for bm in box_maps.iter_mut() { 984 if bm.names[0] == C2PA_BOXHASH { 985 continue; 986 } 987 988 input_stream.seek(std::io::SeekFrom::Start(bm.range_start as u64))?; 989 990 let size = if bm.names[0] == "SOS" { 991 let mut size = get_seg_size(input_stream)?; 992 993 input_stream.seek(std::io::SeekFrom::Start((bm.range_start + size) as u64))?; 994 995 size += get_entropy_size(input_stream)?; 996 997 size 998 } else { 999 get_seg_size(input_stream)? 1000 }; 1001 1002 bm.range_len = size; 1003 } 1004 1005 Ok(box_maps) 1006 } 1007 } 1008 1009 impl ComposedManifestRef for JpegIO { 1010 fn compose_manifest(&self, manifest_data: &[u8], _format: &str) -> Result<Vec<u8>> { 1011 let jumbf_len = manifest_data.len(); 1012 let num_segments = (jumbf_len / MAX_JPEG_MARKER_SIZE) + 1; 1013 let mut seg_chucks = manifest_data.chunks(MAX_JPEG_MARKER_SIZE); 1014 1015 let mut segments = Vec::new(); 1016 1017 for seg in 1..num_segments + 1 { 1018 /* 1019 If the size of the box payload is less than 2^32-8 bytes, 1020 then all fields except the XLBox field, that is: Le, CI, En, Z, LBox and TBox, 1021 shall be present in all JPEG XT marker segment representing this box, 1022 regardless of whether the marker segments starts this box, 1023 or continues a box started by a former JPEG XT Marker segment. 1024 */ 1025 // we need to prefix the JUMBF with the JPEG XT markers (ISO 19566-5) 1026 // CI: JPEG extensions marker - JP 1027 // En: Box Instance Number - 0x0001 1028 // (NOTE: can be any unique ID, so we pick one that shouldn't conflict) 1029 // Z: Packet sequence number - 0x00000001... 1030 let ci = vec![0x4a, 0x50]; 1031 let en = vec![0x02, 0x11]; 1032 let z: u32 = u32::try_from(seg) 1033 .map_err(|_| Error::InvalidAsset("Too many JUMBF segments".to_string()))?; //seg.to_be_bytes(); 1034 1035 let mut seg_data = Vec::new(); 1036 seg_data.extend(ci); 1037 seg_data.extend(en); 1038 seg_data.extend(z.to_be_bytes()); 1039 if seg > 1 { 1040 // the LBox and TBox are already in the JUMBF 1041 // but we need to duplicate them in all other segments 1042 let lbox_tbox = &manifest_data[..8]; 1043 seg_data.extend(lbox_tbox); 1044 } 1045 if seg_chucks.len() > 0 { 1046 // make sure we have some... 1047 if let Some(next_seg) = seg_chucks.next() { 1048 seg_data.extend(next_seg); 1049 } 1050 } else { 1051 seg_data.extend(manifest_data); 1052 } 1053 1054 let seg_bytes = Bytes::from(seg_data); 1055 let app11_segment = JpegSegment::new_with_contents(markers::APP11, seg_bytes); 1056 segments.push(app11_segment); 1057 } 1058 1059 let output = Vec::with_capacity(manifest_data.len() * 2); 1060 let mut out_stream = Cursor::new(output); 1061 1062 // right out segments 1063 for s in segments { 1064 // maker 1065 out_stream.write_u8(markers::P)?; 1066 out_stream.write_u8(s.marker())?; 1067 1068 //len 1069 out_stream.write_u16::<BigEndian>(s.contents().len() as u16 + 2)?; 1070 1071 // data 1072 out_stream.write_all(s.contents())?; 1073 } 1074 1075 Ok(out_stream.into_inner()) 1076 } 1077 } 1078 1079 #[cfg(test)] 1080 pub mod tests { 1081 #![allow(clippy::unwrap_used)] 1082 1083 use std::io::{Read, Seek}; 1084 1085 #[cfg(target_arch = "wasm32")] 1086 use wasm_bindgen_test::*; 1087 1088 use super::*; 1089 #[test] 1090 fn test_extract_xmp() { 1091 let contents = Bytes::from_static(b"http://ns.adobe.com/xap/1.0/\0stuff"); 1092 let seg = JpegSegment::new_with_contents(markers::APP1, contents); 1093 let result = extract_xmp(&seg); 1094 assert_eq!(result, Some("stuff".to_owned())); 1095 1096 let contents = Bytes::from_static(b"http://ns.adobe.com/xap/1.0/ stuff"); 1097 let seg = JpegSegment::new_with_contents(markers::APP1, contents); 1098 let result = extract_xmp(&seg); 1099 assert_eq!(result, Some("stuff".to_owned())); 1100 1101 let contents = Bytes::from_static(b"tiny"); 1102 let seg = JpegSegment::new_with_contents(markers::APP1, contents); 1103 let result = extract_xmp(&seg); 1104 assert_eq!(result, None); 1105 } 1106 1107 #[test] 1108 fn test_remove_c2pa() { 1109 let source = crate::utils::test::fixture_path("CA.jpg"); 1110 1111 let temp_dir = tempfile::tempdir().unwrap(); 1112 let output = crate::utils::test::temp_dir_path(&temp_dir, "CA_test.jpg"); 1113 1114 std::fs::copy(source, &output).unwrap(); 1115 let jpeg_io = JpegIO {}; 1116 1117 jpeg_io.remove_cai_store(&output).unwrap(); 1118 1119 // read back in asset, JumbfNotFound is expected since it was removed 1120 match jpeg_io.read_cai_store(&output) { 1121 Err(Error::JumbfNotFound) => (), 1122 _ => unreachable!(), 1123 } 1124 } 1125 1126 #[test] 1127 fn test_remove_c2pa_from_stream() { 1128 let source = crate::utils::test::fixture_path("CA.jpg"); 1129 1130 let source_bytes = std::fs::read(source).unwrap(); 1131 let mut source_stream = Cursor::new(source_bytes); 1132 1133 let jpeg_io = JpegIO {}; 1134 let jpg_writer = jpeg_io.get_writer("jpg").unwrap(); 1135 1136 let output_bytes = Vec::new(); 1137 let mut output_stream = Cursor::new(output_bytes); 1138 1139 jpg_writer 1140 .remove_cai_store_from_stream(&mut source_stream, &mut output_stream) 1141 .unwrap(); 1142 1143 // read back in asset, JumbfNotFound is expected since it was removed 1144 let jpg_reader = jpeg_io.get_reader(); 1145 match jpg_reader.read_cai(&mut output_stream) { 1146 Err(Error::JumbfNotFound) => (), 1147 _ => unreachable!(), 1148 } 1149 } 1150 1151 #[test] 1152 fn test_xmp_read_write() { 1153 let source = crate::utils::test::fixture_path("CA.jpg"); 1154 1155 let temp_dir = tempfile::tempdir().unwrap(); 1156 let output = crate::utils::test::temp_dir_path(&temp_dir, "CA_test.jpg"); 1157 1158 std::fs::copy(source, &output).unwrap(); 1159 1160 let test_msg = "this some test xmp data"; 1161 let handler = JpegIO::new(""); 1162 1163 // write xmp 1164 let assetio_handler = handler.get_handler("jpg"); 1165 1166 let remote_ref_handler = assetio_handler.remote_ref_writer_ref().unwrap(); 1167 1168 remote_ref_handler 1169 .embed_reference(&output, RemoteRefEmbedType::Xmp(test_msg.to_string())) 1170 .unwrap(); 1171 1172 // read back in XMP 1173 let mut file_reader = std::fs::File::open(&output).unwrap(); 1174 let read_xmp = assetio_handler 1175 .get_reader() 1176 .read_xmp(&mut file_reader) 1177 .unwrap(); 1178 1179 assert!(read_xmp.contains(test_msg)); 1180 } 1181 1182 #[cfg_attr(not(target_arch = "wasm32"), actix::test)] 1183 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)] 1184 async fn test_xmp_read_write_stream() { 1185 let source_bytes = include_bytes!("../../tests/fixtures/CA.jpg"); 1186 1187 let test_msg = "this some test xmp data"; 1188 let handler = JpegIO::new(""); 1189 1190 let assetio_handler = handler.get_handler("jpg"); 1191 1192 let remote_ref_handler = assetio_handler.remote_ref_writer_ref().unwrap(); 1193 1194 let mut source_stream = Cursor::new(source_bytes.to_vec()); 1195 let mut output_stream = Cursor::new(Vec::new()); 1196 remote_ref_handler 1197 .embed_reference_to_stream( 1198 &mut source_stream, 1199 &mut output_stream, 1200 RemoteRefEmbedType::Xmp(test_msg.to_string()), 1201 ) 1202 .unwrap(); 1203 1204 output_stream.set_position(0); 1205 1206 // read back in XMP 1207 let read_xmp = assetio_handler 1208 .get_reader() 1209 .read_xmp(&mut output_stream) 1210 .unwrap(); 1211 1212 output_stream.set_position(0); 1213 1214 //std::fs::write("../target/xmp_write.jpg", output_stream.into_inner()).unwrap(); 1215 1216 assert!(read_xmp.contains(test_msg)); 1217 } 1218 1219 #[test] 1220 fn test_embeddable_manifest() { 1221 let jpeg_io = JpegIO {}; 1222 1223 let source = crate::utils::test::fixture_path("CA.jpg"); 1224 1225 let ol = jpeg_io.get_object_locations(&source).unwrap(); 1226 1227 let cai_loc = ol 1228 .iter() 1229 .find(|o| o.htype == HashBlockObjectType::Cai) 1230 .unwrap(); 1231 let curr_manifest = jpeg_io.read_cai_store(&source).unwrap(); 1232 1233 let temp_dir = tempfile::tempdir().unwrap(); 1234 let output = crate::utils::test::temp_dir_path(&temp_dir, "CA_test.jpg"); 1235 1236 std::fs::copy(source, &output).unwrap(); 1237 1238 // remove existing 1239 jpeg_io.remove_cai_store(&output).unwrap(); 1240 1241 // generate new manifest data 1242 let em = jpeg_io 1243 .composed_data_ref() 1244 .unwrap() 1245 .compose_manifest(&curr_manifest, "jpeg") 1246 .unwrap(); 1247 1248 // insert new manifest 1249 let outbuf = Vec::new(); 1250 let mut out_stream = Cursor::new(outbuf); 1251 1252 let mut before = vec![0u8; cai_loc.offset]; 1253 let mut in_file = std::fs::File::open(&output).unwrap(); 1254 1255 // write before 1256 in_file.read_exact(before.as_mut_slice()).unwrap(); 1257 out_stream.write_all(&before).unwrap(); 1258 1259 // write composed bytes 1260 out_stream.write_all(&em).unwrap(); 1261 1262 // write bytes after 1263 let mut after_buf = Vec::new(); 1264 in_file.read_to_end(&mut after_buf).unwrap(); 1265 out_stream.write_all(&after_buf).unwrap(); 1266 1267 // read manifest back in from new in-memory JPEG 1268 out_stream.rewind().unwrap(); 1269 let restored_manifest = jpeg_io.read_cai(&mut out_stream).unwrap(); 1270 1271 assert_eq!(&curr_manifest, &restored_manifest); 1272 } 1273 }