1 #[cfg(feature="validation")]
2 use core;
3 use std::io::{self, Error, ErrorKind, Read, Write};
4 use super::{Rebox, HeapAllocator, IoWriterWrapper};
5 use brotli::{DecompressorWriterCustomIo, CustomWrite};
6 use alloc_no_stdlib::{Allocator, SliceWrapper};
7 use brotli::enc::BrotliEncoderParams;
8 #[cfg(feature="validation")]
9 use sha2::{Sha256, Digest};
10 #[cfg(feature="validation")]
11 type Checksum = Sha256;
12
13 struct Tee<OutputA:Write, OutputB:Write>(OutputA, OutputB);
14 impl <OutputA:Write, OutputB:Write> Write for Tee<OutputA, OutputB> {
write(&mut self, data:&[u8]) -> Result<usize, io::Error>15 fn write(&mut self, data:&[u8]) -> Result<usize, io::Error> {
16 match self.0.write(data) {
17 Err(err) => return Err(err),
18 Ok(size) => match self.1.write_all(&data[..size]) {
19 Ok(_) => Ok(size),
20 Err(err) => Err(err),
21 }
22 }
23 }
flush(&mut self) -> Result<(), io::Error>24 fn flush(&mut self) -> Result<(), io::Error> {
25 match self.0.flush() {
26 Err(err) => return Err(err),
27 Ok(_) => {
28 loop {
29 match self.1.flush() {
30 Err(e) => match e.kind() {
31 ErrorKind::Interrupted => continue,
32 _ => return Err(e),
33 },
34 Ok(e) => return Ok(e),
35 }
36 }
37 },
38 }
39 }
40 }
41
42 struct DecompressAndValidate<'a, OutputType:Write+'a>(
43 DecompressorWriterCustomIo<io::Error,
44 IoWriterWrapper<'a, OutputType>,
45 Rebox<u8>, // buffer type
46 HeapAllocator,
47 HeapAllocator,
48 HeapAllocator>
49 );
50
51 impl<'a, OutputType:Write> Write for DecompressAndValidate<'a, OutputType> {
write(&mut self, data:&[u8]) -> Result<usize, io::Error>52 fn write(&mut self, data:&[u8]) -> Result<usize, io::Error> {
53 self.0.write(data)
54 }
flush(&mut self) -> Result<(), io::Error>55 fn flush(&mut self) -> Result<(), io::Error> {
56 self.0.flush()
57 }
58 }
59
60 #[cfg(not(feature="validation"))]
make_sha_writer() -> io::Sink61 fn make_sha_writer() -> io::Sink {
62 io::sink()
63 }
64 #[cfg(not(feature="validation"))]
make_sha_reader<InputType:Read>(r:&mut InputType) -> &mut InputType65 fn make_sha_reader<InputType:Read>(r:&mut InputType) -> &mut InputType {
66 r
67 }
68
69 #[cfg(not(feature="validation"))]
sha_ok<InputType:Read>(_writer: &mut io::Sink, _reader: &mut InputType) -> bool70 fn sha_ok<InputType:Read>(_writer: &mut io::Sink, _reader: &mut InputType) -> bool {
71 false
72 }
73
74 #[cfg(feature="validation")]
75 struct ShaReader<'a, InputType:Read+'a> {
76 reader:&'a mut InputType,
77 checksum:Checksum,
78 }
79 #[cfg(feature="validation")]
80 impl<'a, InputType:Read+'a> Read for ShaReader<'a, InputType> {
read(&mut self, data: &mut [u8]) -> Result<usize, io::Error>81 fn read(&mut self, data: &mut [u8]) -> Result<usize, io::Error> {
82 match self.reader.read(data) {
83 Err(e) => Err(e),
84 Ok(size) => {
85 self.checksum.input(&data[..size]);
86 Ok(size)
87 },
88 }
89 }
90 }
91 #[cfg(feature="validation")]
make_sha_reader<InputType:Read>(r:&mut InputType) -> ShaReader<InputType>92 fn make_sha_reader<InputType:Read>(r:&mut InputType) -> ShaReader<InputType> {
93 ShaReader{
94 reader:r,
95 checksum:Checksum::default(),
96 }
97 }
98
99 #[cfg(feature="validation")]
sha_ok<InputType:Read>(writer: &mut ShaWriter, reader: &mut ShaReader<InputType>) -> bool100 fn sha_ok<InputType:Read>(writer: &mut ShaWriter, reader: &mut ShaReader<InputType>) -> bool {
101 core::mem::replace(&mut writer.0,
102 Checksum::default()).result() == core::mem::replace(&mut reader.checksum,
103 Checksum::default()).result()
104 }
105 #[cfg(feature="validation")]
106 #[derive(Default)]
107 struct ShaWriter(Checksum);
108 #[cfg(feature="validation")]
109 impl Write for ShaWriter {
write(&mut self, data:&[u8]) -> Result<usize, io::Error>110 fn write(&mut self, data:&[u8]) -> Result<usize, io::Error> {
111 self.0.input(data);
112 Ok(data.len())
113 }
flush(&mut self) -> Result<(), io::Error>114 fn flush(&mut self) -> Result<(), io::Error> {
115 Ok(())
116 }
117 }
118 #[cfg(feature="validation")]
make_sha_writer() -> ShaWriter119 fn make_sha_writer() -> ShaWriter {
120 ShaWriter::default()
121 }
122 #[cfg(feature="validation")]
123 const VALIDATION_FAILED: &'static str = "Validation failed";
124 #[cfg(not(feature="validation"))]
125 const VALIDATION_FAILED: &'static str = "Validation module not enabled: build with cargo build --features=validation";
126
compress_validate<InputType:Read, OutputType: Write>(r: &mut InputType, w: &mut OutputType, buffer_size: usize, params: &BrotliEncoderParams, custom_dictionary:Rebox<u8>, num_threads: usize) -> Result<(), io::Error>127 pub fn compress_validate<InputType:Read, OutputType: Write>(r: &mut InputType,
128 w: &mut OutputType,
129 buffer_size: usize,
130 params: &BrotliEncoderParams,
131 custom_dictionary:Rebox<u8>,
132 num_threads: usize) -> Result<(), io::Error> {
133 let mut m8 = HeapAllocator::default();
134 let buffer = <HeapAllocator as Allocator<u8>>::alloc_cell(&mut m8, buffer_size);
135 // FIXME: could reuse the dictionary to seed the compressor, but that violates the abstraction right now
136 // also dictionaries are not very popular since they are mostly an internal concept, given their deprecation in
137 // the standard brotli spec
138 let mut dict = Vec::<u8>::new();
139
140 dict.extend_from_slice(custom_dictionary.slice());
141 let mut sha_writer = make_sha_writer();
142 let mut sha_reader = make_sha_reader(r);
143 let ret;
144 {
145 let validate_writer = DecompressAndValidate(DecompressorWriterCustomIo::new_with_custom_dictionary(
146 IoWriterWrapper(&mut sha_writer),
147 buffer,
148 m8,
149 HeapAllocator::default(),
150 HeapAllocator::default(),
151 custom_dictionary,
152 Error::new(ErrorKind::InvalidData,
153 "Invalid Data")));
154 let mut overarching_writer = Tee(validate_writer, w);
155 ret = super::compress(&mut sha_reader, &mut overarching_writer, buffer_size, params, &dict[..], num_threads);
156 }
157 match ret {
158 Ok(_ret) => {
159 if sha_ok(&mut sha_writer, &mut sha_reader) {
160 return Ok(());
161 } else {
162 return Err(Error::new(ErrorKind::InvalidData, VALIDATION_FAILED));
163 }
164 },
165 Err(e) => Err(e),
166 }
167 }
168