1use std::io::{Seek, SeekFrom, Write};
4
5use tiff_core::{ByteOrder, Tag};
6
7use crate::builder::ImageBuilder;
8use crate::compress;
9use crate::encoder;
10use crate::error::{Error, Result};
11use crate::sample::TiffWriteSample;
12
13const CLASSIC_TIFF_LIMIT: u64 = u32::MAX as u64;
14
15fn checked_len_u64(len: usize, context: &str) -> Result<u64> {
16 u64::try_from(len).map_err(|_| Error::Other(format!("{context} length exceeds u64::MAX")))
17}
18
19fn checked_add_u64(lhs: u64, rhs: u64, context: &str) -> Result<u64> {
20 lhs.checked_add(rhs)
21 .ok_or_else(|| Error::Other(format!("{context} overflow")))
22}
23
24fn classic_offset_u32(offset: u64) -> Result<u32> {
25 u32::try_from(offset).map_err(|_| Error::ClassicOffsetOverflow { offset })
26}
27
28fn classic_byte_count_u32(byte_count: u64) -> Result<u32> {
29 u32::try_from(byte_count).map_err(|_| Error::ClassicByteCountOverflow { byte_count })
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum TiffVariant {
35 Classic,
36 BigTiff,
37 Auto,
42}
43
44#[derive(Debug, Clone)]
46pub struct WriteOptions {
47 pub byte_order: ByteOrder,
48 pub variant: TiffVariant,
49}
50
51impl Default for WriteOptions {
52 fn default() -> Self {
53 Self {
54 byte_order: ByteOrder::LittleEndian,
55 variant: TiffVariant::Auto,
56 }
57 }
58}
59
60impl WriteOptions {
61 pub fn auto(_estimated_bytes: u64) -> Self {
67 Self {
68 byte_order: ByteOrder::LittleEndian,
69 variant: TiffVariant::Auto,
70 }
71 }
72}
73
74#[derive(Debug, Clone)]
76pub struct ImageHandle {
77 pub(crate) index: usize,
78}
79
80struct IfdState {
82 builder: ImageBuilder,
83 block_records: Vec<Option<(u64, u64)>>,
84}
85
86pub struct TiffWriter<W: Write + Seek> {
88 sink: W,
89 byte_order: ByteOrder,
90 requested_variant: TiffVariant,
91 header_offset: u64,
92 images: Vec<IfdState>,
93 finalized: bool,
94}
95
96impl<W: Write + Seek> TiffWriter<W> {
97 pub fn new(mut sink: W, options: WriteOptions) -> Result<Self> {
103 let header_offset = sink.stream_position()?;
104 let reserved_header_len = match options.variant {
105 TiffVariant::Classic => encoder::header_len(false),
106 TiffVariant::BigTiff | TiffVariant::Auto => encoder::header_len(true),
107 };
108 sink.write_all(&[0; encoder::BIGTIFF_HEADER_LEN as usize][..reserved_header_len as usize])?;
109
110 Ok(Self {
111 sink,
112 byte_order: options.byte_order,
113 requested_variant: options.variant,
114 header_offset,
115 images: Vec::new(),
116 finalized: false,
117 })
118 }
119
120 pub fn add_image(&mut self, builder: ImageBuilder) -> Result<ImageHandle> {
122 if self.finalized {
123 return Err(Error::AlreadyFinalized);
124 }
125 builder.validate()?;
126
127 let index = self.images.len();
128 self.images.push(IfdState {
129 block_records: vec![None; builder.block_count()],
130 builder,
131 });
132
133 Ok(ImageHandle { index })
134 }
135
136 pub fn write_header_prefix(&mut self, bytes: &[u8]) -> Result<()> {
142 if self.finalized {
143 return Err(Error::AlreadyFinalized);
144 }
145 if !self.images.is_empty() {
146 return Err(Error::Other(
147 "header prefix bytes must be written before adding images".into(),
148 ));
149 }
150
151 self.sink.seek(SeekFrom::End(0))?;
152 let prefix_end = checked_add_u64(
153 self.sink.stream_position()?,
154 checked_len_u64(bytes.len(), "header prefix")?,
155 "header prefix size",
156 )?;
157 if matches!(self.requested_variant, TiffVariant::Classic) {
158 classic_offset_u32(prefix_end)?;
159 }
160 self.sink.write_all(bytes)?;
161 Ok(())
162 }
163
164 pub fn write_block<T: TiffWriteSample>(
166 &mut self,
167 handle: &ImageHandle,
168 block_index: usize,
169 samples: &[T],
170 ) -> Result<()> {
171 if self.finalized {
172 return Err(Error::AlreadyFinalized);
173 }
174 let state = self
175 .images
176 .get(handle.index)
177 .ok_or(Error::Other("invalid image handle".into()))?;
178
179 let total_blocks = state.builder.block_count();
180 if block_index >= total_blocks {
181 return Err(Error::BlockIndexOutOfRange {
182 index: block_index,
183 total: total_blocks,
184 });
185 }
186
187 let expected = state.builder.block_sample_count(block_index);
188 if samples.len() != expected {
189 return Err(Error::BlockSizeMismatch {
190 index: block_index,
191 expected,
192 actual: samples.len(),
193 });
194 }
195
196 let compressed = if matches!(state.builder.compression, tiff_core::Compression::Lerc) {
197 let opts = state.builder.lerc_options.unwrap_or_default();
198 let block_width = state.builder.block_row_width() as u32;
199 let block_height = state.builder.block_height(block_index);
200 let depth = state.builder.block_samples_per_pixel() as u32;
201 compress::compress_block_lerc(
202 samples,
203 block_width,
204 block_height,
205 depth,
206 &opts,
207 block_index,
208 )?
209 } else {
210 compress::compress_block(
211 samples,
212 compress::BlockEncodingOptions {
213 byte_order: self.byte_order,
214 compression: state.builder.compression,
215 predictor: state.builder.predictor,
216 samples_per_pixel: state.builder.block_samples_per_pixel(),
217 row_width_pixels: state.builder.block_row_width(),
218 jpeg_options: state.builder.jpeg_options.as_ref(),
219 },
220 block_index,
221 )?
222 };
223
224 self.write_block_raw(handle, block_index, &compressed)
225 }
226
227 pub fn write_block_raw(
229 &mut self,
230 handle: &ImageHandle,
231 block_index: usize,
232 compressed_bytes: &[u8],
233 ) -> Result<()> {
234 if self.finalized {
235 return Err(Error::AlreadyFinalized);
236 }
237
238 let state = self
239 .images
240 .get(handle.index)
241 .ok_or(Error::Other("invalid image handle".into()))?;
242 let total = state.builder.block_count();
243 if block_index >= total {
244 return Err(Error::BlockIndexOutOfRange {
245 index: block_index,
246 total,
247 });
248 }
249
250 let offset = self.sink.seek(SeekFrom::End(0))?;
251 let byte_count = checked_len_u64(compressed_bytes.len(), "block payload")?;
252 if matches!(self.requested_variant, TiffVariant::Classic) {
253 classic_offset_u32(offset)?;
254 classic_byte_count_u32(byte_count)?;
255 }
256
257 self.sink.write_all(compressed_bytes)?;
258
259 let state = self
260 .images
261 .get_mut(handle.index)
262 .ok_or(Error::Other("invalid image handle".into()))?;
263 state.block_records[block_index] = Some((offset, byte_count));
264 Ok(())
265 }
266
267 pub fn write_block_raw_segmented(
270 &mut self,
271 handle: &ImageHandle,
272 block_index: usize,
273 prefix: &[u8],
274 payload: &[u8],
275 suffix: &[u8],
276 ) -> Result<()> {
277 if self.finalized {
278 return Err(Error::AlreadyFinalized);
279 }
280
281 let state = self
282 .images
283 .get(handle.index)
284 .ok_or(Error::Other("invalid image handle".into()))?;
285 let total = state.builder.block_count();
286 if block_index >= total {
287 return Err(Error::BlockIndexOutOfRange {
288 index: block_index,
289 total,
290 });
291 }
292
293 let start = self.sink.seek(SeekFrom::End(0))?;
294 let prefix_len = checked_len_u64(prefix.len(), "block prefix")?;
295 let byte_count = checked_len_u64(payload.len(), "block payload")?;
296 let suffix_len = checked_len_u64(suffix.len(), "block suffix")?;
297 let offset = checked_add_u64(start, prefix_len, "block offset")?;
298 let end = checked_add_u64(
299 checked_add_u64(offset, byte_count, "segmented block size")?,
300 suffix_len,
301 "segmented block size",
302 )?;
303 if matches!(self.requested_variant, TiffVariant::Classic) {
304 classic_offset_u32(offset)?;
305 classic_byte_count_u32(byte_count)?;
306 classic_offset_u32(end)?;
307 }
308
309 self.sink.write_all(prefix)?;
310 self.sink.write_all(payload)?;
311 self.sink.write_all(suffix)?;
312
313 let state = self
314 .images
315 .get_mut(handle.index)
316 .ok_or(Error::Other("invalid image handle".into()))?;
317 state.block_records[block_index] = Some((offset, byte_count));
318 Ok(())
319 }
320
321 fn choose_is_bigtiff(&mut self) -> Result<bool> {
322 match self.requested_variant {
323 TiffVariant::Classic => {
324 self.ensure_classic_layout()?;
325 Ok(false)
326 }
327 TiffVariant::BigTiff => Ok(true),
328 TiffVariant::Auto => Ok(!self.classic_layout_fits()?),
329 }
330 }
331
332 fn classic_layout_fits(&mut self) -> Result<bool> {
333 for state in &self.images {
334 for &(offset, byte_count) in state.block_records.iter().flatten() {
335 if offset > CLASSIC_TIFF_LIMIT || byte_count > CLASSIC_TIFF_LIMIT {
336 return Ok(false);
337 }
338 }
339 }
340
341 let mut current = self.sink.seek(SeekFrom::End(0))?;
342 for state in &self.images {
343 let tags = state.builder.build_tags(false);
344 current = checked_add_u64(
345 current,
346 encoder::estimate_ifd_size(self.byte_order, false, &tags),
347 "classic IFD layout",
348 )?;
349 if current > CLASSIC_TIFF_LIMIT {
350 return Ok(false);
351 }
352 }
353
354 Ok(true)
355 }
356
357 fn ensure_classic_layout(&mut self) -> Result<()> {
358 for state in &self.images {
359 for &(offset, byte_count) in state.block_records.iter().flatten() {
360 classic_offset_u32(offset)?;
361 classic_byte_count_u32(byte_count)?;
362 }
363 }
364
365 let mut current = self.sink.seek(SeekFrom::End(0))?;
366 for state in &self.images {
367 let tags = state.builder.build_tags(false);
368 current = checked_add_u64(
369 current,
370 encoder::estimate_ifd_size(self.byte_order, false, &tags),
371 "classic IFD layout",
372 )?;
373 classic_offset_u32(current)?;
374 }
375
376 Ok(())
377 }
378
379 fn write_final_ifds(
380 &mut self,
381 is_bigtiff: bool,
382 ) -> Result<Vec<(Vec<Tag>, encoder::IfdWriteResult)>> {
383 let mut results = Vec::with_capacity(self.images.len());
384 for state in &self.images {
385 let tags = state.builder.build_tags(is_bigtiff);
386 let (offsets_tag_code, byte_counts_tag_code) = state.builder.offset_tag_codes();
387 let ifd_result = encoder::write_ifd(
388 &mut self.sink,
389 self.byte_order,
390 is_bigtiff,
391 &tags,
392 offsets_tag_code,
393 byte_counts_tag_code,
394 state.builder.block_count(),
395 )?;
396 results.push((tags, ifd_result));
397 }
398 Ok(results)
399 }
400
401 pub fn finish(mut self) -> Result<W> {
403 if self.finalized {
404 return Err(Error::AlreadyFinalized);
405 }
406 self.finalized = true;
407
408 for state in &self.images {
409 let total = state.builder.block_count();
410 let written = state
411 .block_records
412 .iter()
413 .filter(|record| record.is_some())
414 .count();
415 if written != total {
416 return Err(Error::IncompleteImage { written, total });
417 }
418 }
419
420 let is_bigtiff = self.choose_is_bigtiff()?;
421
422 self.sink.seek(SeekFrom::Start(self.header_offset))?;
423 encoder::write_header(&mut self.sink, self.byte_order, is_bigtiff)?;
424
425 self.sink.seek(SeekFrom::End(0))?;
426 let ifd_results = self.write_final_ifds(is_bigtiff)?;
427
428 for (img_idx, state) in self.images.iter().enumerate() {
429 let offsets: Vec<u64> = state
430 .block_records
431 .iter()
432 .map(|record| record.unwrap().0)
433 .collect();
434 let byte_counts: Vec<u64> = state
435 .block_records
436 .iter()
437 .map(|record| record.unwrap().1)
438 .collect();
439
440 let (tags, ifd_result) = &ifd_results[img_idx];
441 let (offsets_tag_code, byte_counts_tag_code) = state.builder.offset_tag_codes();
442
443 if offsets.len() == 1 {
444 if let Some(off) = encoder::find_inline_tag_value_offset(
445 ifd_result.ifd_offset,
446 is_bigtiff,
447 tags,
448 offsets_tag_code,
449 ) {
450 self.sink.seek(SeekFrom::Start(off))?;
451 if is_bigtiff {
452 self.sink
453 .write_all(&self.byte_order.write_u64(offsets[0]))?;
454 } else {
455 self.sink.write_all(
456 &self.byte_order.write_u32(classic_offset_u32(offsets[0])?),
457 )?;
458 }
459 }
460 if let Some(off) = encoder::find_inline_tag_value_offset(
461 ifd_result.ifd_offset,
462 is_bigtiff,
463 tags,
464 byte_counts_tag_code,
465 ) {
466 self.sink.seek(SeekFrom::Start(off))?;
467 if is_bigtiff {
468 self.sink
469 .write_all(&self.byte_order.write_u64(byte_counts[0]))?;
470 } else {
471 self.sink.write_all(
472 &self
473 .byte_order
474 .write_u32(classic_byte_count_u32(byte_counts[0])?),
475 )?;
476 }
477 }
478 } else {
479 if let Some(off) = ifd_result.offsets_tag_data_offset {
480 encoder::patch_block_offsets(
481 &mut self.sink,
482 self.byte_order,
483 is_bigtiff,
484 off,
485 &offsets,
486 )?;
487 }
488 if let Some(off) = ifd_result.byte_counts_tag_data_offset {
489 encoder::patch_block_byte_counts(
490 &mut self.sink,
491 self.byte_order,
492 is_bigtiff,
493 off,
494 &byte_counts,
495 )?;
496 }
497 }
498
499 if img_idx == 0 {
500 encoder::patch_first_ifd(
501 &mut self.sink,
502 self.header_offset,
503 self.byte_order,
504 is_bigtiff,
505 ifd_result.ifd_offset,
506 )?;
507 } else {
508 let prev = &ifd_results[img_idx - 1].1;
509 encoder::patch_next_ifd(
510 &mut self.sink,
511 self.byte_order,
512 is_bigtiff,
513 prev.next_ifd_pointer_offset,
514 ifd_result.ifd_offset,
515 )?;
516 }
517 }
518
519 self.sink.seek(SeekFrom::End(0))?;
520 Ok(self.sink)
521 }
522}
523
524#[cfg(test)]
525mod tests {
526 use std::io::{self, Cursor, Seek, SeekFrom, Write};
527
528 use super::*;
529 use crate::builder::ImageBuilder;
530
531 #[derive(Default)]
532 struct CountingSink {
533 pos: u64,
534 len: u64,
535 }
536
537 impl Write for CountingSink {
538 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
539 self.pos += buf.len() as u64;
540 self.len = self.len.max(self.pos);
541 Ok(buf.len())
542 }
543
544 fn flush(&mut self) -> io::Result<()> {
545 Ok(())
546 }
547 }
548
549 impl Seek for CountingSink {
550 fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
551 let next = match pos {
552 SeekFrom::Start(offset) => offset as i128,
553 SeekFrom::End(delta) => self.len as i128 + delta as i128,
554 SeekFrom::Current(delta) => self.pos as i128 + delta as i128,
555 };
556 if next < 0 {
557 return Err(io::Error::new(
558 io::ErrorKind::InvalidInput,
559 "negative seek in CountingSink",
560 ));
561 }
562 self.pos = next as u64;
563 self.len = self.len.max(self.pos);
564 Ok(self.pos)
565 }
566 }
567
568 #[test]
569 fn auto_promotes_to_bigtiff_from_the_final_layout() {
570 let mut writer = TiffWriter::new(CountingSink::default(), WriteOptions::default()).unwrap();
571 let handle = writer
572 .add_image(ImageBuilder::new(1, 1).sample_type::<u8>().strips(1))
573 .unwrap();
574
575 writer
576 .sink
577 .seek(SeekFrom::Start(CLASSIC_TIFF_LIMIT + 1))
578 .unwrap();
579 writer.write_block_raw(&handle, 0, &[1]).unwrap();
580
581 assert!(writer.choose_is_bigtiff().unwrap());
582 }
583
584 #[test]
585 fn auto_keeps_classic_for_small_layouts() {
586 let mut writer = TiffWriter::new(Cursor::new(Vec::new()), WriteOptions::default()).unwrap();
587 let handle = writer
588 .add_image(ImageBuilder::new(1, 1).sample_type::<u8>().strips(1))
589 .unwrap();
590 writer.write_block(&handle, 0, &[7u8]).unwrap();
591
592 assert!(!writer.choose_is_bigtiff().unwrap());
593 }
594
595 #[test]
596 fn write_block_raw_validates_before_mutating_sink() {
597 let mut writer = TiffWriter::new(Cursor::new(Vec::new()), WriteOptions::default()).unwrap();
598 let handle = writer
599 .add_image(ImageBuilder::new(1, 1).sample_type::<u8>().strips(1))
600 .unwrap();
601
602 let len_before = writer.sink.get_ref().len();
603 let err = writer.write_block_raw(&handle, 1, &[1, 2, 3]).unwrap_err();
604
605 assert!(matches!(err, Error::BlockIndexOutOfRange { .. }));
606 assert_eq!(writer.sink.get_ref().len(), len_before);
607 }
608
609 #[test]
610 fn write_block_raw_segmented_validates_before_mutating_sink() {
611 let mut writer = TiffWriter::new(Cursor::new(Vec::new()), WriteOptions::default()).unwrap();
612 let handle = writer
613 .add_image(ImageBuilder::new(1, 1).sample_type::<u8>().strips(1))
614 .unwrap();
615
616 let len_before = writer.sink.get_ref().len();
617 let err = writer
618 .write_block_raw_segmented(&handle, 1, &[1, 2], &[3, 4], &[5, 6])
619 .unwrap_err();
620
621 assert!(matches!(err, Error::BlockIndexOutOfRange { .. }));
622 assert_eq!(writer.sink.get_ref().len(), len_before);
623 }
624}