j2k-native 0.7.2

Pure-Rust JPEG 2000 and HTJ2K codec engine for j2k
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Aggregate accounting for allocations created while decoded tile storage is live.

use super::DecompositionStorage;
use crate::error::{Result, ValidationError};
use crate::{try_reserve_decode_elements, try_resize_decode_elements, DEFAULT_MAX_DECODE_BYTES};
use alloc::vec::Vec;
use core::mem::size_of;

pub(crate) struct DecodeAllocationBudget {
    live_bytes: usize,
}

impl DecodeAllocationBudget {
    pub(crate) fn for_storage(storage: &DecompositionStorage<'_>) -> Result<Self> {
        // The structural plan carries parsed metadata, channels, coefficient
        // and graph capacities, ROI arrays, and IDWT workspace. Segment
        // ownership is intentionally separate so retained capacity and packet
        // growth can be charged exactly once here.
        let segment_bytes = storage
            .segments
            .capacity()
            .checked_mul(size_of::<super::Segment<'_>>())
            .ok_or(ValidationError::ImageTooLarge)?;
        let live_bytes = storage
            .structural_workspace_bytes
            .checked_add(segment_bytes)
            .ok_or(ValidationError::ImageTooLarge)?;
        Self::from_live_bytes(live_bytes)
    }

    pub(crate) fn from_live_bytes(live_bytes: usize) -> Result<Self> {
        if live_bytes > DEFAULT_MAX_DECODE_BYTES {
            return Err(ValidationError::ImageTooLarge.into());
        }
        Ok(Self { live_bytes })
    }

    pub(crate) fn include_elements<T>(&mut self, count: usize) -> Result<()> {
        let additional = count
            .checked_mul(size_of::<T>())
            .ok_or(ValidationError::ImageTooLarge)?;
        self.include_bytes(additional)
    }

    pub(crate) fn include_bytes(&mut self, additional: usize) -> Result<()> {
        self.live_bytes = self
            .live_bytes
            .checked_add(additional)
            .ok_or(ValidationError::ImageTooLarge)?;
        if self.live_bytes > DEFAULT_MAX_DECODE_BYTES {
            return Err(ValidationError::ImageTooLarge.into());
        }
        Ok(())
    }

    pub(crate) fn include_capacity_overage<T>(
        &mut self,
        planned_count: usize,
        actual_capacity: usize,
    ) -> Result<()> {
        if actual_capacity > planned_count {
            self.include_elements::<T>(actual_capacity - planned_count)?;
        }
        Ok(())
    }

    pub(crate) fn reserve_new<T>(&mut self, values: &mut Vec<T>, target_len: usize) -> Result<()> {
        *values = Vec::new();
        self.include_elements::<T>(target_len)?;
        try_reserve_decode_elements(values, target_len)?;
        if let Err(error) = self.include_capacity_overage::<T>(target_len, values.capacity()) {
            *values = Vec::new();
            return Err(error);
        }
        Ok(())
    }

    pub(crate) fn resize_new<T: Clone>(
        &mut self,
        values: &mut Vec<T>,
        target_len: usize,
        value: T,
    ) -> Result<()> {
        *values = Vec::new();
        self.include_elements::<T>(target_len)?;
        try_resize_decode_elements(values, target_len, value)?;
        if let Err(error) = self.include_capacity_overage::<T>(target_len, values.capacity()) {
            *values = Vec::new();
            return Err(error);
        }
        Ok(())
    }

    pub(crate) fn release_elements<T>(&mut self, count: usize) -> Result<()> {
        let released = count
            .checked_mul(size_of::<T>())
            .ok_or(ValidationError::ImageTooLarge)?;
        self.live_bytes = self
            .live_bytes
            .checked_sub(released)
            .ok_or(ValidationError::ImageTooLarge)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::DecodeAllocationBudget;
    use crate::error::{DecodeError, ValidationError};
    use crate::DEFAULT_MAX_DECODE_BYTES;
    use alloc::vec::Vec;

    #[test]
    fn aggregate_budget_rejects_a_second_live_owner() {
        let mut budget =
            DecodeAllocationBudget::from_live_bytes(DEFAULT_MAX_DECODE_BYTES - size_of::<u16>())
                .expect("baseline fits");
        let mut owner: Vec<u16> = Vec::new();
        budget.reserve_new(&mut owner, 1).expect("first owner fits");

        let error = budget
            .include_elements::<u8>(1)
            .expect_err("second live owner exceeds cap");
        assert!(matches!(
            error,
            DecodeError::Validation(ValidationError::ImageTooLarge)
        ));
    }
}