rustfs_rio/
hash_reader.rs

1// Copyright 2024 RustFS Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! HashReader implementation with generic support
16//!
17//! This module provides a generic `HashReader<R>` that can wrap any type implementing
18//! `AsyncRead + Unpin + Send + Sync + 'static + EtagResolvable`.
19//!
20//! ## Migration from the original Reader enum
21//!
22//! The original `HashReader::new` method that worked with the `Reader` enum
23//! has been replaced with a generic approach. To preserve the original logic:
24//!
25//! ### Original logic (before generics):
26//! ```ignore
27//! // Original code would do:
28//! // 1. Check if inner is already a HashReader
29//! // 2. If size > 0, wrap with HardLimitReader  
30//! // 3. If !diskable_md5, wrap with EtagReader
31//! // 4. Create HashReader with the wrapped reader
32//!
33//! let reader = HashReader::new(inner, size, actual_size, etag, diskable_md5)?;
34//! ```
35//!
36//! ### New generic approach:
37//! ```rust
38//! use rustfs_rio::{HashReader, HardLimitReader, EtagReader};
39//! use tokio::io::BufReader;
40//! use std::io::Cursor;
41//! use rustfs_rio::WarpReader;
42//!
43//! # tokio_test::block_on(async {
44//! let data = b"hello world";
45//! let reader = BufReader::new(Cursor::new(&data[..]));
46//! let reader = Box::new(WarpReader::new(reader));
47//! let size = data.len() as i64;
48//! let actual_size = size;
49//! let etag = None;
50//! let diskable_md5 = false;
51//!
52//! // Method 1: Simple creation (recommended for most cases)
53//! let hash_reader = HashReader::new(reader, size, actual_size, etag.clone(), diskable_md5).unwrap();
54//!
55//! // Method 2: With manual wrapping to recreate original logic
56//! let reader2 = BufReader::new(Cursor::new(&data[..]));
57//! let reader2 = Box::new(WarpReader::new(reader2));
58//! let wrapped_reader: Box<dyn rustfs_rio::Reader> = if size > 0 {
59//!     if !diskable_md5 {
60//!         // Wrap with both HardLimitReader and EtagReader
61//!         let hard_limit = HardLimitReader::new(reader2, size);
62//!         Box::new(EtagReader::new(Box::new(hard_limit), etag.clone()))
63//!     } else {
64//!         // Only wrap with HardLimitReader
65//!         Box::new(HardLimitReader::new(reader2, size))
66//!     }
67//! } else if !diskable_md5 {
68//!     // Only wrap with EtagReader
69//!     Box::new(EtagReader::new(reader2, etag.clone()))
70//! } else {
71//!     // No wrapping needed
72//!     reader2
73//! };
74//! let hash_reader2 = HashReader::new(wrapped_reader, size, actual_size, etag, diskable_md5).unwrap();
75//! # });
76//! ```
77//!
78//! ## HashReader Detection
79//!
80//! The `HashReaderDetector` trait allows detection of existing HashReader instances:
81//!
82//! ```rust
83//! use rustfs_rio::{HashReader, HashReaderDetector};
84//! use tokio::io::BufReader;
85//! use std::io::Cursor;
86//! use rustfs_rio::WarpReader;
87//!
88//! # tokio_test::block_on(async {
89//! let data = b"test";
90//! let reader = BufReader::new(Cursor::new(&data[..]));
91//! let hash_reader = HashReader::new(Box::new(WarpReader::new(reader)), 4, 4, None, false).unwrap();
92//!
93//! // Check if a type is a HashReader
94//! assert!(hash_reader.is_hash_reader());
95//!
96//! // Use new for compatibility (though it's simpler to use new() directly)
97//! let reader2 = BufReader::new(Cursor::new(&data[..]));
98//! let result = HashReader::new(Box::new(WarpReader::new(reader2)), 4, 4, None, false);
99//! assert!(result.is_ok());
100//! # });
101//! ```
102
103use pin_project_lite::pin_project;
104use std::pin::Pin;
105use std::task::{Context, Poll};
106use tokio::io::{AsyncRead, ReadBuf};
107
108use crate::compress_index::{Index, TryGetIndex};
109use crate::{EtagReader, EtagResolvable, HardLimitReader, HashReaderDetector, Reader};
110
111/// Trait for mutable operations on HashReader
112pub trait HashReaderMut {
113    fn bytes_read(&self) -> u64;
114    fn checksum(&self) -> &Option<String>;
115    fn set_checksum(&mut self, checksum: Option<String>);
116    fn size(&self) -> i64;
117    fn set_size(&mut self, size: i64);
118    fn actual_size(&self) -> i64;
119    fn set_actual_size(&mut self, actual_size: i64);
120}
121
122pin_project! {
123
124    pub struct HashReader {
125        #[pin]
126        pub inner: Box<dyn Reader>,
127        pub size: i64,
128        checksum: Option<String>,
129        pub actual_size: i64,
130        pub diskable_md5: bool,
131        bytes_read: u64,
132        // TODO: content_hash
133    }
134
135}
136
137impl HashReader {
138    pub fn new(
139        mut inner: Box<dyn Reader>,
140        size: i64,
141        actual_size: i64,
142        md5: Option<String>,
143        diskable_md5: bool,
144    ) -> std::io::Result<Self> {
145        // Check if it's already a HashReader and update its parameters
146        if let Some(existing_hash_reader) = inner.as_hash_reader_mut() {
147            if existing_hash_reader.bytes_read() > 0 {
148                return Err(std::io::Error::new(
149                    std::io::ErrorKind::InvalidData,
150                    "Cannot create HashReader from an already read HashReader",
151                ));
152            }
153
154            if let Some(checksum) = existing_hash_reader.checksum() {
155                if let Some(ref md5) = md5 {
156                    if checksum != md5 {
157                        return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "HashReader checksum mismatch"));
158                    }
159                }
160            }
161
162            if existing_hash_reader.size() > 0 && size > 0 && existing_hash_reader.size() != size {
163                return Err(std::io::Error::new(
164                    std::io::ErrorKind::InvalidData,
165                    format!("HashReader size mismatch: expected {}, got {}", existing_hash_reader.size(), size),
166                ));
167            }
168
169            existing_hash_reader.set_checksum(md5.clone());
170
171            if existing_hash_reader.size() < 0 && size >= 0 {
172                existing_hash_reader.set_size(size);
173            }
174
175            if existing_hash_reader.actual_size() <= 0 && actual_size >= 0 {
176                existing_hash_reader.set_actual_size(actual_size);
177            }
178
179            return Ok(Self {
180                inner,
181                size,
182                checksum: md5,
183                actual_size,
184                diskable_md5,
185                bytes_read: 0,
186            });
187        }
188
189        if size > 0 {
190            let hr = HardLimitReader::new(inner, size);
191            inner = Box::new(hr);
192            if !diskable_md5 && !inner.is_hash_reader() {
193                let er = EtagReader::new(inner, md5.clone());
194                inner = Box::new(er);
195            }
196        } else if !diskable_md5 {
197            let er = EtagReader::new(inner, md5.clone());
198            inner = Box::new(er);
199        }
200        Ok(Self {
201            inner,
202            size,
203            checksum: md5,
204            actual_size,
205            diskable_md5,
206            bytes_read: 0,
207        })
208    }
209
210    /// Update HashReader parameters
211    pub fn update_params(&mut self, size: i64, actual_size: i64, etag: Option<String>) {
212        if self.size < 0 && size >= 0 {
213            self.size = size;
214        }
215
216        if self.actual_size <= 0 && actual_size > 0 {
217            self.actual_size = actual_size;
218        }
219
220        if etag.is_some() {
221            self.checksum = etag;
222        }
223    }
224
225    pub fn size(&self) -> i64 {
226        self.size
227    }
228    pub fn actual_size(&self) -> i64 {
229        self.actual_size
230    }
231}
232
233impl HashReaderMut for HashReader {
234    fn bytes_read(&self) -> u64 {
235        self.bytes_read
236    }
237
238    fn checksum(&self) -> &Option<String> {
239        &self.checksum
240    }
241
242    fn set_checksum(&mut self, checksum: Option<String>) {
243        self.checksum = checksum;
244    }
245
246    fn size(&self) -> i64 {
247        self.size
248    }
249
250    fn set_size(&mut self, size: i64) {
251        self.size = size;
252    }
253
254    fn actual_size(&self) -> i64 {
255        self.actual_size
256    }
257
258    fn set_actual_size(&mut self, actual_size: i64) {
259        self.actual_size = actual_size;
260    }
261}
262
263impl AsyncRead for HashReader {
264    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
265        let this = self.project();
266        let poll = this.inner.poll_read(cx, buf);
267        if let Poll::Ready(Ok(())) = &poll {
268            let filled = buf.filled().len();
269            *this.bytes_read += filled as u64;
270
271            if filled == 0 {
272                // EOF
273                // TODO: check content_hash
274            }
275        }
276        poll
277    }
278}
279
280impl EtagResolvable for HashReader {
281    fn try_resolve_etag(&mut self) -> Option<String> {
282        if self.diskable_md5 {
283            return None;
284        }
285        if let Some(etag) = self.inner.try_resolve_etag() {
286            return Some(etag);
287        }
288        // If no etag from inner and we have a stored checksum, return it
289        self.checksum.clone()
290    }
291}
292
293impl HashReaderDetector for HashReader {
294    fn is_hash_reader(&self) -> bool {
295        true
296    }
297
298    fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
299        Some(self)
300    }
301}
302
303impl TryGetIndex for HashReader {
304    fn try_get_index(&self) -> Option<&Index> {
305        self.inner.try_get_index()
306    }
307}
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312    use crate::{DecryptReader, WarpReader, encrypt_reader};
313    use std::io::Cursor;
314    use tokio::io::{AsyncReadExt, BufReader};
315
316    #[tokio::test]
317    async fn test_hashreader_wrapping_logic() {
318        let data = b"hello world";
319        let size = data.len() as i64;
320        let actual_size = size;
321        let etag = None;
322
323        // Test 1: Simple creation
324        let reader1 = BufReader::new(Cursor::new(&data[..]));
325        let reader1 = Box::new(WarpReader::new(reader1));
326        let hash_reader1 = HashReader::new(reader1, size, actual_size, etag.clone(), false).unwrap();
327        assert_eq!(hash_reader1.size(), size);
328        assert_eq!(hash_reader1.actual_size(), actual_size);
329
330        // Test 2: With HardLimitReader wrapping
331        let reader2 = BufReader::new(Cursor::new(&data[..]));
332        let reader2 = Box::new(WarpReader::new(reader2));
333        let hard_limit = HardLimitReader::new(reader2, size);
334        let hard_limit = Box::new(hard_limit);
335        let hash_reader2 = HashReader::new(hard_limit, size, actual_size, etag.clone(), false).unwrap();
336        assert_eq!(hash_reader2.size(), size);
337        assert_eq!(hash_reader2.actual_size(), actual_size);
338
339        // Test 3: With EtagReader wrapping
340        let reader3 = BufReader::new(Cursor::new(&data[..]));
341        let reader3 = Box::new(WarpReader::new(reader3));
342        let etag_reader = EtagReader::new(reader3, etag.clone());
343        let etag_reader = Box::new(etag_reader);
344        let hash_reader3 = HashReader::new(etag_reader, size, actual_size, etag.clone(), false).unwrap();
345        assert_eq!(hash_reader3.size(), size);
346        assert_eq!(hash_reader3.actual_size(), actual_size);
347    }
348
349    #[tokio::test]
350    async fn test_hashreader_etag_basic() {
351        let data = b"hello hashreader";
352        let reader = BufReader::new(Cursor::new(&data[..]));
353        let reader = Box::new(WarpReader::new(reader));
354        let mut hash_reader = HashReader::new(reader, data.len() as i64, data.len() as i64, None, false).unwrap();
355        let mut buf = Vec::new();
356        let _ = hash_reader.read_to_end(&mut buf).await.unwrap();
357        // Since we removed EtagReader integration, etag might be None
358        let _etag = hash_reader.try_resolve_etag();
359        // Just check that we can call etag() without error
360        assert_eq!(buf, data);
361    }
362
363    #[tokio::test]
364    async fn test_hashreader_diskable_md5() {
365        let data = b"no etag";
366        let reader = BufReader::new(Cursor::new(&data[..]));
367        let reader = Box::new(WarpReader::new(reader));
368        let mut hash_reader = HashReader::new(reader, data.len() as i64, data.len() as i64, None, true).unwrap();
369        let mut buf = Vec::new();
370        let _ = hash_reader.read_to_end(&mut buf).await.unwrap();
371        // Etag should be None when diskable_md5 is true
372        let etag = hash_reader.try_resolve_etag();
373        assert!(etag.is_none());
374        assert_eq!(buf, data);
375    }
376
377    #[tokio::test]
378    async fn test_hashreader_new_logic() {
379        let data = b"test data";
380        let reader = BufReader::new(Cursor::new(&data[..]));
381        let reader = Box::new(WarpReader::new(reader));
382        // Create a HashReader first
383        let hash_reader =
384            HashReader::new(reader, data.len() as i64, data.len() as i64, Some("test_etag".to_string()), false).unwrap();
385        let hash_reader = Box::new(WarpReader::new(hash_reader));
386        // Now try to create another HashReader from the existing one using new
387        let result = HashReader::new(hash_reader, data.len() as i64, data.len() as i64, Some("test_etag".to_string()), false);
388
389        assert!(result.is_ok());
390        let final_reader = result.unwrap();
391        assert_eq!(final_reader.checksum, Some("test_etag".to_string()));
392        assert_eq!(final_reader.size(), data.len() as i64);
393    }
394
395    #[tokio::test]
396    async fn test_for_wrapping_readers() {
397        use crate::{CompressReader, DecompressReader};
398        use md5::{Digest, Md5};
399        use rand::Rng;
400        use rand::RngCore;
401        use rustfs_utils::compress::CompressionAlgorithm;
402
403        // Generate 1MB random data
404        let size = 1024 * 1024;
405        let mut data = vec![0u8; size];
406        rand::rng().fill(&mut data[..]);
407
408        let mut hasher = Md5::new();
409        hasher.update(&data);
410
411        let expected = format!("{:x}", hasher.finalize());
412
413        println!("expected: {expected}");
414
415        let reader = Cursor::new(data.clone());
416        let reader = BufReader::new(reader);
417
418        // 启用压缩测试
419        let is_compress = true;
420        let size = data.len() as i64;
421        let actual_size = data.len() as i64;
422
423        let reader = Box::new(WarpReader::new(reader));
424        // 创建 HashReader
425        let mut hr = HashReader::new(reader, size, actual_size, Some(expected.clone()), false).unwrap();
426
427        // 如果启用压缩，先压缩数据
428        let compressed_data = if is_compress {
429            let mut compressed_buf = Vec::new();
430            let compress_reader = CompressReader::new(hr, CompressionAlgorithm::Gzip);
431            let mut compress_reader = compress_reader;
432            compress_reader.read_to_end(&mut compressed_buf).await.unwrap();
433
434            println!("Original size: {}, Compressed size: {}", data.len(), compressed_buf.len());
435
436            compressed_buf
437        } else {
438            // 如果不压缩，直接读取原始数据
439            let mut buf = Vec::new();
440            hr.read_to_end(&mut buf).await.unwrap();
441            buf
442        };
443
444        let mut key = [0u8; 32];
445        let mut nonce = [0u8; 12];
446        rand::rng().fill_bytes(&mut key);
447        rand::rng().fill_bytes(&mut nonce);
448
449        let is_encrypt = true;
450
451        if is_encrypt {
452            // 加密压缩后的数据
453            let encrypt_reader = encrypt_reader::EncryptReader::new(WarpReader::new(Cursor::new(compressed_data)), key, nonce);
454            let mut encrypted_data = Vec::new();
455            let mut encrypt_reader = encrypt_reader;
456            encrypt_reader.read_to_end(&mut encrypted_data).await.unwrap();
457
458            println!("Encrypted size: {}", encrypted_data.len());
459
460            // 解密数据
461            let decrypt_reader = DecryptReader::new(WarpReader::new(Cursor::new(encrypted_data)), key, nonce);
462            let mut decrypt_reader = decrypt_reader;
463            let mut decrypted_data = Vec::new();
464            decrypt_reader.read_to_end(&mut decrypted_data).await.unwrap();
465
466            if is_compress {
467                // 如果使用了压缩，需要解压缩
468                let decompress_reader =
469                    DecompressReader::new(WarpReader::new(Cursor::new(decrypted_data)), CompressionAlgorithm::Gzip);
470                let mut decompress_reader = decompress_reader;
471                let mut final_data = Vec::new();
472                decompress_reader.read_to_end(&mut final_data).await.unwrap();
473
474                println!("Final decompressed size: {}", final_data.len());
475                assert_eq!(final_data.len() as i64, actual_size);
476                assert_eq!(&final_data, &data);
477            } else {
478                // 如果没有压缩，直接比较解密后的数据
479                assert_eq!(decrypted_data.len() as i64, actual_size);
480                assert_eq!(&decrypted_data, &data);
481            }
482            return;
483        }
484
485        // 如果不加密，直接处理压缩/解压缩
486        if is_compress {
487            let decompress_reader =
488                DecompressReader::new(WarpReader::new(Cursor::new(compressed_data)), CompressionAlgorithm::Gzip);
489            let mut decompress_reader = decompress_reader;
490            let mut decompressed = Vec::new();
491            decompress_reader.read_to_end(&mut decompressed).await.unwrap();
492
493            assert_eq!(decompressed.len() as i64, actual_size);
494            assert_eq!(&decompressed, &data);
495        } else {
496            assert_eq!(compressed_data.len() as i64, actual_size);
497            assert_eq!(&compressed_data, &data);
498        }
499
500        // 验证 etag（注意：压缩会改变数据，所以这里的 etag 验证可能需要调整）
501        println!("Test completed successfully with compression: {is_compress}, encryption: {is_encrypt}");
502    }
503
504    #[tokio::test]
505    async fn test_compression_with_compressible_data() {
506        use crate::{CompressReader, DecompressReader};
507        use rustfs_utils::compress::CompressionAlgorithm;
508
509        // Create highly compressible data (repeated pattern)
510        let pattern = b"Hello, World! This is a test pattern that should compress well. ";
511        let repeat_count = 16384; // 16K repetitions
512        let mut data = Vec::new();
513        for _ in 0..repeat_count {
514            data.extend_from_slice(pattern);
515        }
516
517        println!("Original data size: {} bytes", data.len());
518
519        let reader = BufReader::new(Cursor::new(data.clone()));
520        let reader = Box::new(WarpReader::new(reader));
521        let hash_reader = HashReader::new(reader, data.len() as i64, data.len() as i64, None, false).unwrap();
522
523        // Test compression
524        let compress_reader = CompressReader::new(hash_reader, CompressionAlgorithm::Gzip);
525        let mut compressed_data = Vec::new();
526        let mut compress_reader = compress_reader;
527        compress_reader.read_to_end(&mut compressed_data).await.unwrap();
528
529        println!("Compressed data size: {} bytes", compressed_data.len());
530        println!("Compression ratio: {:.2}%", (compressed_data.len() as f64 / data.len() as f64) * 100.0);
531
532        // Verify compression actually reduced size for this compressible data
533        assert!(compressed_data.len() < data.len(), "Compression should reduce size for repetitive data");
534
535        // Test decompression
536        let decompress_reader = DecompressReader::new(Cursor::new(compressed_data), CompressionAlgorithm::Gzip);
537        let mut decompressed_data = Vec::new();
538        let mut decompress_reader = decompress_reader;
539        decompress_reader.read_to_end(&mut decompressed_data).await.unwrap();
540
541        // Verify decompressed data matches original
542        assert_eq!(decompressed_data.len(), data.len());
543        assert_eq!(&decompressed_data, &data);
544
545        println!("Compression/decompression test passed successfully!");
546    }
547
548    #[tokio::test]
549    async fn test_compression_algorithms() {
550        use crate::{CompressReader, DecompressReader};
551        use rustfs_utils::compress::CompressionAlgorithm;
552
553        let data = b"This is test data for compression algorithm testing. ".repeat(1000);
554        println!("Testing with {} bytes of data", data.len());
555
556        let algorithms = vec![
557            CompressionAlgorithm::Gzip,
558            CompressionAlgorithm::Deflate,
559            CompressionAlgorithm::Zstd,
560        ];
561
562        for algorithm in algorithms {
563            println!("\nTesting algorithm: {algorithm:?}");
564
565            let reader = BufReader::new(Cursor::new(data.clone()));
566            let reader = Box::new(WarpReader::new(reader));
567            let hash_reader = HashReader::new(reader, data.len() as i64, data.len() as i64, None, false).unwrap();
568
569            // Compress
570            let compress_reader = CompressReader::new(hash_reader, algorithm);
571            let mut compressed_data = Vec::new();
572            let mut compress_reader = compress_reader;
573            compress_reader.read_to_end(&mut compressed_data).await.unwrap();
574
575            println!(
576                "  Compressed size: {} bytes (ratio: {:.2}%)",
577                compressed_data.len(),
578                (compressed_data.len() as f64 / data.len() as f64) * 100.0
579            );
580
581            // Decompress
582            let decompress_reader = DecompressReader::new(Cursor::new(compressed_data), algorithm);
583            let mut decompressed_data = Vec::new();
584            let mut decompress_reader = decompress_reader;
585            decompress_reader.read_to_end(&mut decompressed_data).await.unwrap();
586
587            // Verify
588            assert_eq!(decompressed_data.len(), data.len());
589            assert_eq!(&decompressed_data, &data);
590            println!("  ✓ Algorithm {algorithm:?} test passed");
591        }
592    }
593}
rustfs_rio/hash_reader.rs

rustfs_rio/
hash_reader.rs
