Merge branch 'main' of github.com:rustfs/s3-rustfs into feature/observability-metrics

# Conflicts:
#	.github/workflows/build.yml
#	.github/workflows/ci.yml
#	Cargo.lock
#	Cargo.toml
#	appauth/src/token.rs
#	crates/config/src/config.rs
#	crates/event-notifier/examples/simple.rs
#	crates/event-notifier/src/global.rs
#	crates/event-notifier/src/lib.rs
#	crates/event-notifier/src/notifier.rs
#	crates/event-notifier/src/store.rs
#	crates/filemeta/src/filemeta.rs
#	crates/notify/examples/webhook.rs
#	crates/utils/Cargo.toml
#	ecstore/Cargo.toml
#	ecstore/src/cmd/bucket_replication.rs
#	ecstore/src/config/com.rs
#	ecstore/src/disk/error.rs
#	ecstore/src/disk/mod.rs
#	ecstore/src/set_disk.rs
#	ecstore/src/store_api.rs
#	ecstore/src/store_list_objects.rs
#	iam/Cargo.toml
#	iam/src/manager.rs
#	policy/Cargo.toml
#	rustfs/src/admin/rpc.rs
#	rustfs/src/main.rs
#	rustfs/src/storage/mod.rs
This commit is contained in:
houseme
2025-06-19 13:16:48 +08:00
249 changed files with 25137 additions and 11731 deletions
+2 -2
View File
@@ -200,7 +200,7 @@ mod tests {
// Test port related constants
assert_eq!(DEFAULT_PORT, 9000);
assert_eq!(DEFAULT_CONSOLE_PORT, 9002);
assert_eq!(DEFAULT_CONSOLE_PORT, 9001);
assert_ne!(DEFAULT_PORT, DEFAULT_CONSOLE_PORT, "Main port and console port should be different");
}
@@ -215,7 +215,7 @@ mod tests {
"Address should contain the default port"
);
assert_eq!(DEFAULT_CONSOLE_ADDRESS, ":9002");
assert_eq!(DEFAULT_CONSOLE_ADDRESS, ":9001");
assert!(DEFAULT_CONSOLE_ADDRESS.starts_with(':'), "Console address should start with colon");
assert!(
DEFAULT_CONSOLE_ADDRESS.contains(&DEFAULT_CONSOLE_PORT.to_string()),
+32
View File
@@ -0,0 +1,32 @@
[package]
name = "rustfs-filemeta"
edition.workspace = true
license.workspace = true
repository.workspace = true
rust-version.workspace = true
version.workspace = true
[dependencies]
crc32fast = "1.4.2"
rmp.workspace = true
rmp-serde.workspace = true
serde.workspace = true
time.workspace = true
uuid = { workspace = true, features = ["v4", "fast-rng", "serde"] }
tokio = { workspace = true, features = ["io-util", "macros", "sync"] }
xxhash-rust = { version = "0.8.15", features = ["xxh64"] }
bytes.workspace = true
rustfs-utils = {workspace = true, features= ["hash"]}
byteorder = "1.5.0"
tracing.workspace = true
thiserror.workspace = true
[dev-dependencies]
criterion = { version = "0.5", features = ["html_reports"] }
[[bench]]
name = "xl_meta_bench"
harness = false
[lints]
workspace = true
+238
View File
@@ -0,0 +1,238 @@
# RustFS FileMeta
A high-performance Rust implementation of xl-storage-format-v2, providing complete compatibility with S3-compatible metadata format while offering enhanced performance and safety.
## Overview
This crate implements the XL (Erasure Coded) metadata format used for distributed object storage. It provides:
- **Full S3 Compatibility**: 100% compatible with xl.meta file format
- **High Performance**: Optimized for speed with sub-microsecond parsing times
- **Memory Safety**: Written in safe Rust with comprehensive error handling
- **Comprehensive Testing**: Extensive test suite with real metadata validation
- **Cross-Platform**: Supports multiple CPU architectures (x86_64, aarch64)
## Features
### Core Functionality
- ✅ XL v2 file format parsing and serialization
- ✅ MessagePack-based metadata encoding/decoding
- ✅ Version management with modification time sorting
- ✅ Erasure coding information storage
- ✅ Inline data support for small objects
- ✅ CRC32 integrity verification using xxHash64
- ✅ Delete marker handling
- ✅ Legacy version support
### Advanced Features
- ✅ Signature calculation for version integrity
- ✅ Metadata validation and compatibility checking
- ✅ Version statistics and analytics
- ✅ Async I/O support with tokio
- ✅ Comprehensive error handling
- ✅ Performance benchmarking
## Performance
Based on our benchmarks:
| Operation | Time | Description |
|-----------|------|-------------|
| Parse Real xl.meta | ~255 ns | Parse authentic xl metadata |
| Parse Complex xl.meta | ~1.1 µs | Parse multi-version metadata |
| Serialize Real xl.meta | ~659 ns | Serialize to xl format |
| Round-trip Real xl.meta | ~1.3 µs | Parse + serialize cycle |
| Version Statistics | ~5.2 ns | Calculate version stats |
| Integrity Validation | ~7.8 ns | Validate metadata integrity |
## Usage
### Basic Usage
```rust
use rustfs_filemeta::file_meta::FileMeta;
// Load metadata from bytes
let metadata = FileMeta::load(&xl_meta_bytes)?;
// Access version information
for version in &metadata.versions {
println!("Version ID: {:?}", version.header.version_id);
println!("Mod Time: {:?}", version.header.mod_time);
}
// Serialize back to bytes
let serialized = metadata.marshal_msg()?;
```
### Advanced Usage
```rust
use rustfs_filemeta::file_meta::FileMeta;
// Load with validation
let mut metadata = FileMeta::load(&xl_meta_bytes)?;
// Validate integrity
metadata.validate_integrity()?;
// Check xl format compatibility
if metadata.is_compatible_with_meta() {
println!("Compatible with xl format");
}
// Get version statistics
let stats = metadata.get_version_stats();
println!("Total versions: {}", stats.total_versions);
println!("Object versions: {}", stats.object_versions);
println!("Delete markers: {}", stats.delete_markers);
```
### Working with FileInfo
```rust
use rustfs_filemeta::fileinfo::FileInfo;
use rustfs_filemeta::file_meta::FileMetaVersion;
// Convert FileInfo to metadata version
let file_info = FileInfo::new("bucket", "object.txt");
let meta_version = FileMetaVersion::from(file_info);
// Add version to metadata
metadata.add_version(file_info)?;
```
## Data Structures
### FileMeta
The main metadata container that holds all versions and inline data:
```rust
pub struct FileMeta {
pub versions: Vec<FileMetaShallowVersion>,
pub data: InlineData,
pub meta_ver: u8,
}
```
### FileMetaVersion
Represents a single object version:
```rust
pub struct FileMetaVersion {
pub version_type: VersionType,
pub object: Option<MetaObject>,
pub delete_marker: Option<MetaDeleteMarker>,
pub write_version: u64,
}
```
### MetaObject
Contains object-specific metadata including erasure coding information:
```rust
pub struct MetaObject {
pub version_id: Option<Uuid>,
pub data_dir: Option<Uuid>,
pub erasure_algorithm: ErasureAlgo,
pub erasure_m: usize,
pub erasure_n: usize,
// ... additional fields
}
```
## File Format Compatibility
This implementation is fully compatible with xl-storage-format-v2:
- **Header Format**: XL2 v1 format with proper version checking
- **Serialization**: MessagePack encoding identical to standard format
- **Checksums**: xxHash64-based CRC validation
- **Version Types**: Support for Object, Delete, and Legacy versions
- **Inline Data**: Compatible inline data storage for small objects
## Testing
The crate includes comprehensive tests with real xl metadata:
```bash
# Run all tests
cargo test
# Run benchmarks
cargo bench
# Run with coverage
cargo test --features coverage
```
### Test Coverage
- ✅ Real xl.meta file compatibility
- ✅ Complex multi-version scenarios
- ✅ Error handling and recovery
- ✅ Inline data processing
- ✅ Signature calculation
- ✅ Round-trip serialization
- ✅ Performance benchmarks
- ✅ Edge cases and boundary conditions
## Architecture
The crate follows a modular design:
```
src/
├── file_meta.rs # Core metadata structures and logic
├── file_meta_inline.rs # Inline data handling
├── fileinfo.rs # File information structures
├── test_data.rs # Test data generation
└── lib.rs # Public API exports
```
## Error Handling
Comprehensive error handling with detailed error messages:
```rust
use rustfs_filemeta::error::Error;
match FileMeta::load(&invalid_data) {
Ok(metadata) => { /* process metadata */ },
Err(Error::InvalidFormat(msg)) => {
eprintln!("Invalid format: {}", msg);
},
Err(Error::CorruptedData(msg)) => {
eprintln!("Corrupted data: {}", msg);
},
Err(e) => {
eprintln!("Other error: {}", e);
}
}
```
## Dependencies
- `rmp` - MessagePack serialization
- `uuid` - UUID handling
- `time` - Date/time operations
- `xxhash-rust` - Fast hashing
- `tokio` - Async runtime (optional)
- `criterion` - Benchmarking (dev dependency)
## Contributing
1. Fork the repository
2. Create a feature branch
3. Add tests for new functionality
4. Ensure all tests pass
5. Submit a pull request
## License
This project is licensed under the Apache License 2.0 - see the LICENSE file for details.
## Acknowledgments
- Original xl-storage-format-v2 implementation contributors
- Rust community for excellent crates and tooling
- Contributors and testers who helped improve this implementation
+95
View File
@@ -0,0 +1,95 @@
use criterion::{Criterion, black_box, criterion_group, criterion_main};
use rustfs_filemeta::{FileMeta, test_data::*};
fn bench_create_real_xlmeta(c: &mut Criterion) {
c.bench_function("create_real_xlmeta", |b| b.iter(|| black_box(create_real_xlmeta().unwrap())));
}
fn bench_create_complex_xlmeta(c: &mut Criterion) {
c.bench_function("create_complex_xlmeta", |b| b.iter(|| black_box(create_complex_xlmeta().unwrap())));
}
fn bench_parse_real_xlmeta(c: &mut Criterion) {
let data = create_real_xlmeta().unwrap();
c.bench_function("parse_real_xlmeta", |b| b.iter(|| black_box(FileMeta::load(&data).unwrap())));
}
fn bench_parse_complex_xlmeta(c: &mut Criterion) {
let data = create_complex_xlmeta().unwrap();
c.bench_function("parse_complex_xlmeta", |b| b.iter(|| black_box(FileMeta::load(&data).unwrap())));
}
fn bench_serialize_real_xlmeta(c: &mut Criterion) {
let data = create_real_xlmeta().unwrap();
let fm = FileMeta::load(&data).unwrap();
c.bench_function("serialize_real_xlmeta", |b| b.iter(|| black_box(fm.marshal_msg().unwrap())));
}
fn bench_serialize_complex_xlmeta(c: &mut Criterion) {
let data = create_complex_xlmeta().unwrap();
let fm = FileMeta::load(&data).unwrap();
c.bench_function("serialize_complex_xlmeta", |b| b.iter(|| black_box(fm.marshal_msg().unwrap())));
}
fn bench_round_trip_real_xlmeta(c: &mut Criterion) {
let original_data = create_real_xlmeta().unwrap();
c.bench_function("round_trip_real_xlmeta", |b| {
b.iter(|| {
let fm = FileMeta::load(&original_data).unwrap();
let serialized = fm.marshal_msg().unwrap();
black_box(FileMeta::load(&serialized).unwrap())
})
});
}
fn bench_round_trip_complex_xlmeta(c: &mut Criterion) {
let original_data = create_complex_xlmeta().unwrap();
c.bench_function("round_trip_complex_xlmeta", |b| {
b.iter(|| {
let fm = FileMeta::load(&original_data).unwrap();
let serialized = fm.marshal_msg().unwrap();
black_box(FileMeta::load(&serialized).unwrap())
})
});
}
fn bench_version_stats(c: &mut Criterion) {
let data = create_complex_xlmeta().unwrap();
let fm = FileMeta::load(&data).unwrap();
c.bench_function("version_stats", |b| b.iter(|| black_box(fm.get_version_stats())));
}
fn bench_validate_integrity(c: &mut Criterion) {
let data = create_real_xlmeta().unwrap();
let fm = FileMeta::load(&data).unwrap();
c.bench_function("validate_integrity", |b| {
b.iter(|| {
fm.validate_integrity().unwrap();
black_box(())
})
});
}
criterion_group!(
benches,
bench_create_real_xlmeta,
bench_create_complex_xlmeta,
bench_parse_real_xlmeta,
bench_parse_complex_xlmeta,
bench_serialize_real_xlmeta,
bench_serialize_complex_xlmeta,
bench_round_trip_real_xlmeta,
bench_round_trip_complex_xlmeta,
bench_version_stats,
bench_validate_integrity
);
criterion_main!(benches);
+569
View File
@@ -0,0 +1,569 @@
pub type Result<T> = core::result::Result<T, Error>;
#[derive(thiserror::Error, Debug)]
pub enum Error {
#[error("File not found")]
FileNotFound,
#[error("File version not found")]
FileVersionNotFound,
#[error("Volume not found")]
VolumeNotFound,
#[error("File corrupt")]
FileCorrupt,
#[error("Done for now")]
DoneForNow,
#[error("Method not allowed")]
MethodNotAllowed,
#[error("Unexpected error")]
Unexpected,
#[error("I/O error: {0}")]
Io(std::io::Error),
#[error("rmp serde decode error: {0}")]
RmpSerdeDecode(String),
#[error("rmp serde encode error: {0}")]
RmpSerdeEncode(String),
#[error("Invalid UTF-8: {0}")]
FromUtf8(String),
#[error("rmp decode value read error: {0}")]
RmpDecodeValueRead(String),
#[error("rmp encode value write error: {0}")]
RmpEncodeValueWrite(String),
#[error("rmp decode num value read error: {0}")]
RmpDecodeNumValueRead(String),
#[error("rmp decode marker read error: {0}")]
RmpDecodeMarkerRead(String),
#[error("time component range error: {0}")]
TimeComponentRange(String),
#[error("uuid parse error: {0}")]
UuidParse(String),
}
impl Error {
pub fn other<E>(error: E) -> Error
where
E: Into<Box<dyn std::error::Error + Send + Sync>>,
{
std::io::Error::other(error).into()
}
}
impl PartialEq for Error {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
(Error::FileCorrupt, Error::FileCorrupt) => true,
(Error::DoneForNow, Error::DoneForNow) => true,
(Error::MethodNotAllowed, Error::MethodNotAllowed) => true,
(Error::FileNotFound, Error::FileNotFound) => true,
(Error::FileVersionNotFound, Error::FileVersionNotFound) => true,
(Error::VolumeNotFound, Error::VolumeNotFound) => true,
(Error::Io(e1), Error::Io(e2)) => e1.kind() == e2.kind() && e1.to_string() == e2.to_string(),
(Error::RmpSerdeDecode(e1), Error::RmpSerdeDecode(e2)) => e1 == e2,
(Error::RmpSerdeEncode(e1), Error::RmpSerdeEncode(e2)) => e1 == e2,
(Error::RmpDecodeValueRead(e1), Error::RmpDecodeValueRead(e2)) => e1 == e2,
(Error::RmpEncodeValueWrite(e1), Error::RmpEncodeValueWrite(e2)) => e1 == e2,
(Error::RmpDecodeNumValueRead(e1), Error::RmpDecodeNumValueRead(e2)) => e1 == e2,
(Error::TimeComponentRange(e1), Error::TimeComponentRange(e2)) => e1 == e2,
(Error::UuidParse(e1), Error::UuidParse(e2)) => e1 == e2,
(Error::Unexpected, Error::Unexpected) => true,
(a, b) => a.to_string() == b.to_string(),
}
}
}
impl Clone for Error {
fn clone(&self) -> Self {
match self {
Error::FileNotFound => Error::FileNotFound,
Error::FileVersionNotFound => Error::FileVersionNotFound,
Error::FileCorrupt => Error::FileCorrupt,
Error::DoneForNow => Error::DoneForNow,
Error::MethodNotAllowed => Error::MethodNotAllowed,
Error::VolumeNotFound => Error::VolumeNotFound,
Error::Io(e) => Error::Io(std::io::Error::new(e.kind(), e.to_string())),
Error::RmpSerdeDecode(s) => Error::RmpSerdeDecode(s.clone()),
Error::RmpSerdeEncode(s) => Error::RmpSerdeEncode(s.clone()),
Error::FromUtf8(s) => Error::FromUtf8(s.clone()),
Error::RmpDecodeValueRead(s) => Error::RmpDecodeValueRead(s.clone()),
Error::RmpEncodeValueWrite(s) => Error::RmpEncodeValueWrite(s.clone()),
Error::RmpDecodeNumValueRead(s) => Error::RmpDecodeNumValueRead(s.clone()),
Error::RmpDecodeMarkerRead(s) => Error::RmpDecodeMarkerRead(s.clone()),
Error::TimeComponentRange(s) => Error::TimeComponentRange(s.clone()),
Error::UuidParse(s) => Error::UuidParse(s.clone()),
Error::Unexpected => Error::Unexpected,
}
}
}
impl From<std::io::Error> for Error {
fn from(e: std::io::Error) -> Self {
match e.kind() {
std::io::ErrorKind::UnexpectedEof => Error::Unexpected,
_ => Error::Io(e),
}
}
}
impl From<Error> for std::io::Error {
fn from(e: Error) -> Self {
match e {
Error::Unexpected => std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "Unexpected EOF"),
Error::Io(e) => e,
_ => std::io::Error::other(e.to_string()),
}
}
}
impl From<rmp_serde::decode::Error> for Error {
fn from(e: rmp_serde::decode::Error) -> Self {
Error::RmpSerdeDecode(e.to_string())
}
}
impl From<rmp_serde::encode::Error> for Error {
fn from(e: rmp_serde::encode::Error) -> Self {
Error::RmpSerdeEncode(e.to_string())
}
}
impl From<std::string::FromUtf8Error> for Error {
fn from(e: std::string::FromUtf8Error) -> Self {
Error::FromUtf8(e.to_string())
}
}
impl From<rmp::decode::ValueReadError> for Error {
fn from(e: rmp::decode::ValueReadError) -> Self {
Error::RmpDecodeValueRead(e.to_string())
}
}
impl From<rmp::encode::ValueWriteError> for Error {
fn from(e: rmp::encode::ValueWriteError) -> Self {
Error::RmpEncodeValueWrite(e.to_string())
}
}
impl From<rmp::decode::NumValueReadError> for Error {
fn from(e: rmp::decode::NumValueReadError) -> Self {
Error::RmpDecodeNumValueRead(e.to_string())
}
}
impl From<time::error::ComponentRange> for Error {
fn from(e: time::error::ComponentRange) -> Self {
Error::TimeComponentRange(e.to_string())
}
}
impl From<uuid::Error> for Error {
fn from(e: uuid::Error) -> Self {
Error::UuidParse(e.to_string())
}
}
impl From<rmp::decode::MarkerReadError> for Error {
fn from(e: rmp::decode::MarkerReadError) -> Self {
let serr = format!("{:?}", e);
Error::RmpDecodeMarkerRead(serr)
}
}
pub fn is_io_eof(e: &Error) -> bool {
match e {
Error::Io(e) => e.kind() == std::io::ErrorKind::UnexpectedEof,
_ => false,
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{Error as IoError, ErrorKind};
#[test]
fn test_filemeta_error_from_io_error() {
let io_error = IoError::new(ErrorKind::PermissionDenied, "permission denied");
let filemeta_error: Error = io_error.into();
match filemeta_error {
Error::Io(inner_io) => {
assert_eq!(inner_io.kind(), ErrorKind::PermissionDenied);
assert!(inner_io.to_string().contains("permission denied"));
}
_ => panic!("Expected Io variant"),
}
}
#[test]
fn test_filemeta_error_other_function() {
let custom_error = "Custom filemeta error";
let filemeta_error = Error::other(custom_error);
match filemeta_error {
Error::Io(io_error) => {
assert!(io_error.to_string().contains(custom_error));
assert_eq!(io_error.kind(), ErrorKind::Other);
}
_ => panic!("Expected Io variant"),
}
}
#[test]
fn test_filemeta_error_conversions() {
// Test various error conversions
let serde_decode_err =
rmp_serde::decode::Error::InvalidMarkerRead(std::io::Error::new(ErrorKind::InvalidData, "invalid"));
let filemeta_error: Error = serde_decode_err.into();
assert!(matches!(filemeta_error, Error::RmpSerdeDecode(_)));
// Test with string-based error that we can actually create
let encode_error_string = "test encode error";
let filemeta_error = Error::RmpSerdeEncode(encode_error_string.to_string());
assert!(matches!(filemeta_error, Error::RmpSerdeEncode(_)));
let utf8_err = std::string::String::from_utf8(vec![0xFF]).unwrap_err();
let filemeta_error: Error = utf8_err.into();
assert!(matches!(filemeta_error, Error::FromUtf8(_)));
}
#[test]
fn test_filemeta_error_clone() {
let test_cases = vec![
Error::FileNotFound,
Error::FileVersionNotFound,
Error::VolumeNotFound,
Error::FileCorrupt,
Error::DoneForNow,
Error::MethodNotAllowed,
Error::Unexpected,
Error::Io(IoError::new(ErrorKind::NotFound, "test")),
Error::RmpSerdeDecode("test decode error".to_string()),
Error::RmpSerdeEncode("test encode error".to_string()),
Error::FromUtf8("test utf8 error".to_string()),
Error::RmpDecodeValueRead("test value read error".to_string()),
Error::RmpEncodeValueWrite("test value write error".to_string()),
Error::RmpDecodeNumValueRead("test num read error".to_string()),
Error::RmpDecodeMarkerRead("test marker read error".to_string()),
Error::TimeComponentRange("test time error".to_string()),
Error::UuidParse("test uuid error".to_string()),
];
for original_error in test_cases {
let cloned_error = original_error.clone();
assert_eq!(original_error, cloned_error);
}
}
#[test]
fn test_filemeta_error_partial_eq() {
// Test equality for simple variants
assert_eq!(Error::FileNotFound, Error::FileNotFound);
assert_ne!(Error::FileNotFound, Error::FileVersionNotFound);
// Test equality for Io variants
let io1 = Error::Io(IoError::new(ErrorKind::NotFound, "test"));
let io2 = Error::Io(IoError::new(ErrorKind::NotFound, "test"));
let io3 = Error::Io(IoError::new(ErrorKind::PermissionDenied, "test"));
assert_eq!(io1, io2);
assert_ne!(io1, io3);
// Test equality for string variants
let decode1 = Error::RmpSerdeDecode("error message".to_string());
let decode2 = Error::RmpSerdeDecode("error message".to_string());
let decode3 = Error::RmpSerdeDecode("different message".to_string());
assert_eq!(decode1, decode2);
assert_ne!(decode1, decode3);
}
#[test]
fn test_filemeta_error_display() {
let test_cases = vec![
(Error::FileNotFound, "File not found"),
(Error::FileVersionNotFound, "File version not found"),
(Error::VolumeNotFound, "Volume not found"),
(Error::FileCorrupt, "File corrupt"),
(Error::DoneForNow, "Done for now"),
(Error::MethodNotAllowed, "Method not allowed"),
(Error::Unexpected, "Unexpected error"),
(Error::RmpSerdeDecode("test".to_string()), "rmp serde decode error: test"),
(Error::RmpSerdeEncode("test".to_string()), "rmp serde encode error: test"),
(Error::FromUtf8("test".to_string()), "Invalid UTF-8: test"),
(Error::TimeComponentRange("test".to_string()), "time component range error: test"),
(Error::UuidParse("test".to_string()), "uuid parse error: test"),
];
for (error, expected_message) in test_cases {
assert_eq!(error.to_string(), expected_message);
}
}
#[test]
fn test_rmp_conversions() {
// Test rmp value read error (this one works since it has the same signature)
let value_read_err = rmp::decode::ValueReadError::InvalidMarkerRead(std::io::Error::new(ErrorKind::InvalidData, "test"));
let filemeta_error: Error = value_read_err.into();
assert!(matches!(filemeta_error, Error::RmpDecodeValueRead(_)));
// Test rmp num value read error
let num_value_err =
rmp::decode::NumValueReadError::InvalidMarkerRead(std::io::Error::new(ErrorKind::InvalidData, "test"));
let filemeta_error: Error = num_value_err.into();
assert!(matches!(filemeta_error, Error::RmpDecodeNumValueRead(_)));
}
#[test]
fn test_time_and_uuid_conversions() {
// Test time component range error
use time::{Date, Month};
let time_result = Date::from_calendar_date(2023, Month::January, 32); // Invalid day
assert!(time_result.is_err());
let time_error = time_result.unwrap_err();
let filemeta_error: Error = time_error.into();
assert!(matches!(filemeta_error, Error::TimeComponentRange(_)));
// Test UUID parse error
let uuid_result = uuid::Uuid::parse_str("invalid-uuid");
assert!(uuid_result.is_err());
let uuid_error = uuid_result.unwrap_err();
let filemeta_error: Error = uuid_error.into();
assert!(matches!(filemeta_error, Error::UuidParse(_)));
}
#[test]
fn test_marker_read_error_conversion() {
// Test rmp marker read error conversion
let marker_err = rmp::decode::MarkerReadError(std::io::Error::new(ErrorKind::InvalidData, "marker test"));
let filemeta_error: Error = marker_err.into();
assert!(matches!(filemeta_error, Error::RmpDecodeMarkerRead(_)));
assert!(filemeta_error.to_string().contains("marker"));
}
#[test]
fn test_is_io_eof_function() {
// Test is_io_eof helper function
let eof_error = Error::Io(IoError::new(ErrorKind::UnexpectedEof, "eof"));
assert!(is_io_eof(&eof_error));
let not_eof_error = Error::Io(IoError::new(ErrorKind::NotFound, "not found"));
assert!(!is_io_eof(&not_eof_error));
let non_io_error = Error::FileNotFound;
assert!(!is_io_eof(&non_io_error));
}
#[test]
fn test_filemeta_error_to_io_error_conversion() {
// Test conversion from FileMeta Error to io::Error through other function
let original_io_error = IoError::new(ErrorKind::InvalidData, "test data");
let filemeta_error = Error::other(original_io_error);
match filemeta_error {
Error::Io(io_err) => {
assert_eq!(io_err.kind(), ErrorKind::Other);
assert!(io_err.to_string().contains("test data"));
}
_ => panic!("Expected Io variant"),
}
}
#[test]
fn test_filemeta_error_roundtrip_conversion() {
// Test roundtrip conversion: io::Error -> FileMeta Error -> io::Error
let original_io_error = IoError::new(ErrorKind::PermissionDenied, "permission test");
// Convert to FileMeta Error
let filemeta_error: Error = original_io_error.into();
// Extract the io::Error back
match filemeta_error {
Error::Io(extracted_io_error) => {
assert_eq!(extracted_io_error.kind(), ErrorKind::PermissionDenied);
assert!(extracted_io_error.to_string().contains("permission test"));
}
_ => panic!("Expected Io variant"),
}
}
#[test]
fn test_filemeta_error_io_error_kinds_preservation() {
let io_error_kinds = vec![
ErrorKind::NotFound,
ErrorKind::PermissionDenied,
ErrorKind::ConnectionRefused,
ErrorKind::ConnectionReset,
ErrorKind::ConnectionAborted,
ErrorKind::NotConnected,
ErrorKind::AddrInUse,
ErrorKind::AddrNotAvailable,
ErrorKind::BrokenPipe,
ErrorKind::AlreadyExists,
ErrorKind::WouldBlock,
ErrorKind::InvalidInput,
ErrorKind::InvalidData,
ErrorKind::TimedOut,
ErrorKind::WriteZero,
ErrorKind::Interrupted,
ErrorKind::UnexpectedEof,
ErrorKind::Other,
];
for kind in io_error_kinds {
let io_error = IoError::new(kind, format!("test error for {:?}", kind));
let filemeta_error: Error = io_error.into();
match filemeta_error {
Error::Unexpected => {
assert_eq!(kind, ErrorKind::UnexpectedEof);
}
Error::Io(extracted_io_error) => {
assert_eq!(extracted_io_error.kind(), kind);
assert!(extracted_io_error.to_string().contains("test error"));
}
_ => panic!("Expected Io variant for kind {:?}", kind),
}
}
}
#[test]
fn test_filemeta_error_downcast_chain() {
// Test error downcast chain functionality
let original_io_error = IoError::new(ErrorKind::InvalidData, "original error");
let filemeta_error = Error::other(original_io_error);
// The error should be wrapped as an Io variant
if let Error::Io(io_err) = filemeta_error {
// The wrapped error should be Other kind (from std::io::Error::other)
assert_eq!(io_err.kind(), ErrorKind::Other);
// But the message should still contain the original error information
assert!(io_err.to_string().contains("original error"));
} else {
panic!("Expected Io variant");
}
}
#[test]
fn test_filemeta_error_maintains_error_information() {
let test_cases = vec![
(ErrorKind::NotFound, "file not found"),
(ErrorKind::PermissionDenied, "access denied"),
(ErrorKind::InvalidData, "corrupt data"),
(ErrorKind::TimedOut, "operation timed out"),
];
for (kind, message) in test_cases {
let io_error = IoError::new(kind, message);
let error_message = io_error.to_string();
let filemeta_error: Error = io_error.into();
match filemeta_error {
Error::Io(extracted_io_error) => {
assert_eq!(extracted_io_error.kind(), kind);
assert_eq!(extracted_io_error.to_string(), error_message);
}
_ => panic!("Expected Io variant"),
}
}
}
#[test]
fn test_filemeta_error_complex_conversion_chain() {
// Test conversion from string error types that we can actually create
// Test with UUID error conversion
let uuid_result = uuid::Uuid::parse_str("invalid-uuid-format");
assert!(uuid_result.is_err());
let uuid_error = uuid_result.unwrap_err();
let filemeta_error: Error = uuid_error.into();
match filemeta_error {
Error::UuidParse(message) => {
assert!(message.contains("invalid"));
}
_ => panic!("Expected UuidParse variant"),
}
// Test with time error conversion
use time::{Date, Month};
let time_result = Date::from_calendar_date(2023, Month::January, 32); // Invalid day
assert!(time_result.is_err());
let time_error = time_result.unwrap_err();
let filemeta_error2: Error = time_error.into();
match filemeta_error2 {
Error::TimeComponentRange(message) => {
assert!(message.contains("range"));
}
_ => panic!("Expected TimeComponentRange variant"),
}
// Test with UTF8 error conversion
let utf8_result = std::string::String::from_utf8(vec![0xFF]);
assert!(utf8_result.is_err());
let utf8_error = utf8_result.unwrap_err();
let filemeta_error3: Error = utf8_error.into();
match filemeta_error3 {
Error::FromUtf8(message) => {
assert!(message.contains("utf"));
}
_ => panic!("Expected FromUtf8 variant"),
}
}
#[test]
fn test_filemeta_error_equality_with_io_errors() {
// Test equality comparison for Io variants
let io_error1 = IoError::new(ErrorKind::NotFound, "test message");
let io_error2 = IoError::new(ErrorKind::NotFound, "test message");
let io_error3 = IoError::new(ErrorKind::PermissionDenied, "test message");
let io_error4 = IoError::new(ErrorKind::NotFound, "different message");
let filemeta_error1 = Error::Io(io_error1);
let filemeta_error2 = Error::Io(io_error2);
let filemeta_error3 = Error::Io(io_error3);
let filemeta_error4 = Error::Io(io_error4);
// Same kind and message should be equal
assert_eq!(filemeta_error1, filemeta_error2);
// Different kinds should not be equal
assert_ne!(filemeta_error1, filemeta_error3);
// Different messages should not be equal
assert_ne!(filemeta_error1, filemeta_error4);
}
#[test]
fn test_filemeta_error_clone_io_variants() {
let io_error = IoError::new(ErrorKind::ConnectionReset, "connection lost");
let original_error = Error::Io(io_error);
let cloned_error = original_error.clone();
// Cloned error should be equal to original
assert_eq!(original_error, cloned_error);
// Both should maintain the same properties
match (original_error, cloned_error) {
(Error::Io(orig_io), Error::Io(cloned_io)) => {
assert_eq!(orig_io.kind(), cloned_io.kind());
assert_eq!(orig_io.to_string(), cloned_io.to_string());
}
_ => panic!("Both should be Io variants"),
}
}
}
+457
View File
@@ -0,0 +1,457 @@
use crate::error::{Error, Result};
use crate::headers::RESERVED_METADATA_PREFIX_LOWER;
use crate::headers::RUSTFS_HEALING;
use bytes::Bytes;
use rmp_serde::Serializer;
use rustfs_utils::HashAlgorithm;
use serde::Deserialize;
use serde::Serialize;
use std::collections::HashMap;
use time::OffsetDateTime;
use uuid::Uuid;
pub const ERASURE_ALGORITHM: &str = "rs-vandermonde";
pub const BLOCK_SIZE_V2: usize = 1024 * 1024; // 1M
// Additional constants from Go version
pub const NULL_VERSION_ID: &str = "null";
// pub const RUSTFS_ERASURE_UPGRADED: &str = "x-rustfs-internal-erasure-upgraded";
#[derive(Serialize, Deserialize, Debug, PartialEq, Clone, Default)]
pub struct ObjectPartInfo {
pub etag: String,
pub number: usize,
pub size: usize,
pub actual_size: i64, // Original data size
pub mod_time: Option<OffsetDateTime>,
// Index holds the index of the part in the erasure coding
pub index: Option<Bytes>,
// Checksums holds checksums of the part
pub checksums: Option<HashMap<String, String>>,
}
#[derive(Serialize, Deserialize, Debug, PartialEq, Default, Clone)]
// ChecksumInfo - carries checksums of individual scattered parts per disk.
pub struct ChecksumInfo {
pub part_number: usize,
pub algorithm: HashAlgorithm,
pub hash: Bytes,
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Default, Clone)]
pub enum ErasureAlgo {
#[default]
Invalid = 0,
ReedSolomon = 1,
}
impl ErasureAlgo {
pub fn valid(&self) -> bool {
*self > ErasureAlgo::Invalid
}
pub fn to_u8(&self) -> u8 {
match self {
ErasureAlgo::Invalid => 0,
ErasureAlgo::ReedSolomon => 1,
}
}
pub fn from_u8(u: u8) -> Self {
match u {
1 => ErasureAlgo::ReedSolomon,
_ => ErasureAlgo::Invalid,
}
}
}
impl std::fmt::Display for ErasureAlgo {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ErasureAlgo::Invalid => write!(f, "Invalid"),
ErasureAlgo::ReedSolomon => write!(f, "{}", ERASURE_ALGORITHM),
}
}
}
#[derive(Serialize, Deserialize, Debug, PartialEq, Default, Clone)]
// ErasureInfo holds erasure coding and bitrot related information.
pub struct ErasureInfo {
// Algorithm is the String representation of erasure-coding-algorithm
pub algorithm: String,
// DataBlocks is the number of data blocks for erasure-coding
pub data_blocks: usize,
// ParityBlocks is the number of parity blocks for erasure-coding
pub parity_blocks: usize,
// BlockSize is the size of one erasure-coded block
pub block_size: usize,
// Index is the index of the current disk
pub index: usize,
// Distribution is the distribution of the data and parity blocks
pub distribution: Vec<usize>,
// Checksums holds all bitrot checksums of all erasure encoded blocks
pub checksums: Vec<ChecksumInfo>,
}
pub fn calc_shard_size(block_size: usize, data_shards: usize) -> usize {
(block_size.div_ceil(data_shards) + 1) & !1
}
impl ErasureInfo {
pub fn get_checksum_info(&self, part_number: usize) -> ChecksumInfo {
for sum in &self.checksums {
if sum.part_number == part_number {
return sum.clone();
}
}
ChecksumInfo {
algorithm: HashAlgorithm::HighwayHash256S,
..Default::default()
}
}
/// Calculate the size of each shard.
pub fn shard_size(&self) -> usize {
calc_shard_size(self.block_size, self.data_blocks)
}
/// Calculate the total erasure file size for a given original size.
// Returns the final erasure size from the original size
pub fn shard_file_size(&self, total_length: i64) -> i64 {
if total_length == 0 {
return 0;
}
if total_length < 0 {
return total_length;
}
let total_length = total_length as usize;
let num_shards = total_length / self.block_size;
let last_block_size = total_length % self.block_size;
let last_shard_size = calc_shard_size(last_block_size, self.data_blocks);
(num_shards * self.shard_size() + last_shard_size) as i64
}
/// Check if this ErasureInfo equals another ErasureInfo
pub fn equals(&self, other: &ErasureInfo) -> bool {
self.algorithm == other.algorithm
&& self.data_blocks == other.data_blocks
&& self.parity_blocks == other.parity_blocks
&& self.block_size == other.block_size
&& self.index == other.index
&& self.distribution == other.distribution
}
}
// #[derive(Debug, Clone)]
#[derive(Serialize, Deserialize, Debug, PartialEq, Clone, Default)]
pub struct FileInfo {
pub volume: String,
pub name: String,
pub version_id: Option<Uuid>,
pub is_latest: bool,
pub deleted: bool,
// Transition related fields
pub transition_status: Option<String>,
pub transitioned_obj_name: Option<String>,
pub transition_tier: Option<String>,
pub transition_version_id: Option<String>,
pub expire_restored: bool,
pub data_dir: Option<Uuid>,
pub mod_time: Option<OffsetDateTime>,
pub size: i64,
// File mode bits
pub mode: Option<u32>,
// WrittenByVersion is the unix time stamp of the version that created this version of the object
pub written_by_version: Option<u64>,
pub metadata: HashMap<String, String>,
pub parts: Vec<ObjectPartInfo>,
pub erasure: ErasureInfo,
// MarkDeleted marks this version as deleted
pub mark_deleted: bool,
// ReplicationState - Internal replication state to be passed back in ObjectInfo
// pub replication_state: Option<ReplicationState>, // TODO: implement ReplicationState
pub data: Option<Bytes>,
pub num_versions: usize,
pub successor_mod_time: Option<OffsetDateTime>,
pub fresh: bool,
pub idx: usize,
// Combined checksum when object was uploaded
pub checksum: Option<Bytes>,
pub versioned: bool,
}
impl FileInfo {
pub fn new(object: &str, data_blocks: usize, parity_blocks: usize) -> Self {
let indexs = {
let cardinality = data_blocks + parity_blocks;
let mut nums = vec![0; cardinality];
let key_crc = crc32fast::hash(object.as_bytes());
let start = key_crc as usize % cardinality;
for i in 1..=cardinality {
nums[i - 1] = 1 + ((start + i) % cardinality);
}
nums
};
Self {
erasure: ErasureInfo {
algorithm: String::from(ERASURE_ALGORITHM),
data_blocks,
parity_blocks,
block_size: BLOCK_SIZE_V2,
distribution: indexs,
..Default::default()
},
..Default::default()
}
}
pub fn is_valid(&self) -> bool {
if self.deleted {
return true;
}
let data_blocks = self.erasure.data_blocks;
let parity_blocks = self.erasure.parity_blocks;
(data_blocks >= parity_blocks)
&& (data_blocks > 0)
&& (self.erasure.index > 0
&& self.erasure.index <= data_blocks + parity_blocks
&& self.erasure.distribution.len() == (data_blocks + parity_blocks))
}
pub fn get_etag(&self) -> Option<String> {
self.metadata.get("etag").cloned()
}
pub fn write_quorum(&self, quorum: usize) -> usize {
if self.deleted {
return quorum;
}
if self.erasure.data_blocks == self.erasure.parity_blocks {
return self.erasure.data_blocks + 1;
}
self.erasure.data_blocks
}
pub fn marshal_msg(&self) -> Result<Vec<u8>> {
let mut buf = Vec::new();
self.serialize(&mut Serializer::new(&mut buf))?;
Ok(buf)
}
pub fn unmarshal(buf: &[u8]) -> Result<Self> {
let t: FileInfo = rmp_serde::from_slice(buf)?;
Ok(t)
}
pub fn add_object_part(
&mut self,
num: usize,
etag: String,
part_size: usize,
mod_time: Option<OffsetDateTime>,
actual_size: i64,
index: Option<Bytes>,
) {
let part = ObjectPartInfo {
etag,
number: num,
size: part_size,
mod_time,
actual_size,
index,
checksums: None,
};
for p in self.parts.iter_mut() {
if p.number == num {
*p = part;
return;
}
}
self.parts.push(part);
self.parts.sort_by(|a, b| a.number.cmp(&b.number));
}
// to_part_offset gets the part index where offset is located, returns part index and offset
pub fn to_part_offset(&self, offset: usize) -> Result<(usize, usize)> {
if offset == 0 {
return Ok((0, 0));
}
let mut part_offset = offset;
for (i, part) in self.parts.iter().enumerate() {
let part_index = i;
if part_offset < part.size {
return Ok((part_index, part_offset));
}
part_offset -= part.size
}
Err(Error::other("part not found"))
}
pub fn set_healing(&mut self) {
self.metadata.insert(RUSTFS_HEALING.to_string(), "true".to_string());
}
pub fn set_inline_data(&mut self) {
self.metadata
.insert(format!("{}inline-data", RESERVED_METADATA_PREFIX_LOWER).to_owned(), "true".to_owned());
}
pub fn set_data_moved(&mut self) {
self.metadata
.insert(format!("{}data-moved", RESERVED_METADATA_PREFIX_LOWER).to_owned(), "true".to_owned());
}
pub fn inline_data(&self) -> bool {
self.metadata
.contains_key(format!("{}inline-data", RESERVED_METADATA_PREFIX_LOWER).as_str())
&& !self.is_remote()
}
/// Check if the object is compressed
pub fn is_compressed(&self) -> bool {
self.metadata
.contains_key(&format!("{}compression", RESERVED_METADATA_PREFIX_LOWER))
}
/// Check if the object is remote (transitioned to another tier)
pub fn is_remote(&self) -> bool {
!self.transition_tier.as_ref().is_none_or(|s| s.is_empty())
}
/// Get the data directory for this object
pub fn get_data_dir(&self) -> String {
if self.deleted {
return "delete-marker".to_string();
}
self.data_dir.map_or("".to_string(), |dir| dir.to_string())
}
/// Read quorum returns expected read quorum for this FileInfo
pub fn read_quorum(&self, dquorum: usize) -> usize {
if self.deleted {
return dquorum;
}
self.erasure.data_blocks
}
/// Create a shallow copy with minimal information for READ MRF checks
pub fn shallow_copy(&self) -> Self {
Self {
volume: self.volume.clone(),
name: self.name.clone(),
version_id: self.version_id,
deleted: self.deleted,
erasure: self.erasure.clone(),
..Default::default()
}
}
/// Check if this FileInfo equals another FileInfo
pub fn equals(&self, other: &FileInfo) -> bool {
// Check if both are compressed or both are not compressed
if self.is_compressed() != other.is_compressed() {
return false;
}
// Check transition info
if !self.transition_info_equals(other) {
return false;
}
// Check mod time
if self.mod_time != other.mod_time {
return false;
}
// Check erasure info
self.erasure.equals(&other.erasure)
}
/// Check if transition related information are equal
pub fn transition_info_equals(&self, other: &FileInfo) -> bool {
self.transition_status == other.transition_status
&& self.transition_tier == other.transition_tier
&& self.transitioned_obj_name == other.transitioned_obj_name
&& self.transition_version_id == other.transition_version_id
}
/// Check if metadata maps are equal
pub fn metadata_equals(&self, other: &FileInfo) -> bool {
if self.metadata.len() != other.metadata.len() {
return false;
}
for (k, v) in &self.metadata {
if other.metadata.get(k) != Some(v) {
return false;
}
}
true
}
/// Check if replication related fields are equal
pub fn replication_info_equals(&self, other: &FileInfo) -> bool {
self.mark_deleted == other.mark_deleted
// TODO: Add replication_state comparison when implemented
// && self.replication_state == other.replication_state
}
}
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct FileInfoVersions {
// Name of the volume.
pub volume: String,
// Name of the file.
pub name: String,
// Represents the latest mod time of the
// latest version.
pub latest_mod_time: Option<OffsetDateTime>,
pub versions: Vec<FileInfo>,
pub free_versions: Vec<FileInfo>,
}
impl FileInfoVersions {
pub fn find_version_index(&self, v: &str) -> Option<usize> {
if v.is_empty() {
return None;
}
let vid = Uuid::parse_str(v).unwrap_or_default();
self.versions.iter().position(|v| v.version_id == Some(vid))
}
/// Calculate the total size of all versions for this object
pub fn size(&self) -> i64 {
self.versions.iter().map(|v| v.size).sum()
}
}
#[derive(Default, Serialize, Deserialize)]
pub struct RawFileInfo {
pub buf: Vec<u8>,
}
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct FilesInfo {
pub files: Vec<FileInfo>,
pub is_truncated: bool,
}
File diff suppressed because it is too large Load Diff
+238
View File
@@ -0,0 +1,238 @@
use crate::error::{Error, Result};
use serde::{Deserialize, Serialize};
use std::io::{Cursor, Read};
use uuid::Uuid;
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct InlineData(Vec<u8>);
const INLINE_DATA_VER: u8 = 1;
impl InlineData {
pub fn new() -> Self {
Self(Vec::new())
}
pub fn update(&mut self, buf: &[u8]) {
self.0 = buf.to_vec()
}
pub fn as_slice(&self) -> &[u8] {
self.0.as_slice()
}
pub fn version_ok(&self) -> bool {
if self.0.is_empty() {
return true;
}
self.0[0] > 0 && self.0[0] <= INLINE_DATA_VER
}
pub fn after_version(&self) -> &[u8] {
if self.0.is_empty() { &self.0 } else { &self.0[1..] }
}
pub fn find(&self, key: &str) -> Result<Option<Vec<u8>>> {
if self.0.is_empty() || !self.version_ok() {
return Ok(None);
}
let buf = self.after_version();
let mut cur = Cursor::new(buf);
let mut fields_len = rmp::decode::read_map_len(&mut cur)?;
while fields_len > 0 {
fields_len -= 1;
let str_len = rmp::decode::read_str_len(&mut cur)?;
let mut field_buff = vec![0u8; str_len as usize];
cur.read_exact(&mut field_buff)?;
let field = String::from_utf8(field_buff)?;
let bin_len = rmp::decode::read_bin_len(&mut cur)? as usize;
let start = cur.position() as usize;
let end = start + bin_len;
cur.set_position(end as u64);
if field.as_str() == key {
let buf = &buf[start..end];
return Ok(Some(buf.to_vec()));
}
}
Ok(None)
}
pub fn validate(&self) -> Result<()> {
if self.0.is_empty() {
return Ok(());
}
let mut cur = Cursor::new(self.after_version());
let mut fields_len = rmp::decode::read_map_len(&mut cur)?;
while fields_len > 0 {
fields_len -= 1;
let str_len = rmp::decode::read_str_len(&mut cur)?;
let mut field_buff = vec![0u8; str_len as usize];
cur.read_exact(&mut field_buff)?;
let field = String::from_utf8(field_buff)?;
if field.is_empty() {
return Err(Error::other("InlineData key empty"));
}
let bin_len = rmp::decode::read_bin_len(&mut cur)? as usize;
let start = cur.position() as usize;
let end = start + bin_len;
cur.set_position(end as u64);
}
Ok(())
}
pub fn replace(&mut self, key: &str, value: Vec<u8>) -> Result<()> {
if self.after_version().is_empty() {
let mut keys = Vec::with_capacity(1);
let mut values = Vec::with_capacity(1);
keys.push(key.to_owned());
values.push(value);
return self.serialize(keys, values);
}
let buf = self.after_version();
let mut cur = Cursor::new(buf);
let mut fields_len = rmp::decode::read_map_len(&mut cur)? as usize;
let mut keys = Vec::with_capacity(fields_len + 1);
let mut values = Vec::with_capacity(fields_len + 1);
let mut replaced = false;
while fields_len > 0 {
fields_len -= 1;
let str_len = rmp::decode::read_str_len(&mut cur)?;
let mut field_buff = vec![0u8; str_len as usize];
cur.read_exact(&mut field_buff)?;
let find_key = String::from_utf8(field_buff)?;
let bin_len = rmp::decode::read_bin_len(&mut cur)? as usize;
let start = cur.position() as usize;
let end = start + bin_len;
cur.set_position(end as u64);
let find_value = &buf[start..end];
if find_key.as_str() == key {
values.push(value.clone());
replaced = true
} else {
values.push(find_value.to_vec());
}
keys.push(find_key);
}
if !replaced {
keys.push(key.to_owned());
values.push(value);
}
self.serialize(keys, values)
}
pub fn remove(&mut self, remove_keys: Vec<Uuid>) -> Result<bool> {
let buf = self.after_version();
let mut cur = Cursor::new(buf);
let mut fields_len = rmp::decode::read_map_len(&mut cur)? as usize;
let mut keys = Vec::with_capacity(fields_len + 1);
let mut values = Vec::with_capacity(fields_len + 1);
let remove_key = |found_key: &str| {
for key in remove_keys.iter() {
if key.to_string().as_str() == found_key {
return true;
}
}
false
};
let mut found = false;
while fields_len > 0 {
fields_len -= 1;
let str_len = rmp::decode::read_str_len(&mut cur)?;
let mut field_buff = vec![0u8; str_len as usize];
cur.read_exact(&mut field_buff)?;
let find_key = String::from_utf8(field_buff)?;
let bin_len = rmp::decode::read_bin_len(&mut cur)? as usize;
let start = cur.position() as usize;
let end = start + bin_len;
cur.set_position(end as u64);
let find_value = &buf[start..end];
if !remove_key(&find_key) {
values.push(find_value.to_vec());
keys.push(find_key);
} else {
found = true;
}
}
if !found {
return Ok(false);
}
if keys.is_empty() {
self.0 = Vec::new();
return Ok(true);
}
self.serialize(keys, values)?;
Ok(true)
}
fn serialize(&mut self, keys: Vec<String>, values: Vec<Vec<u8>>) -> Result<()> {
assert_eq!(keys.len(), values.len(), "InlineData serialize: keys/values not match");
if keys.is_empty() {
self.0 = Vec::new();
return Ok(());
}
let mut wr = Vec::new();
wr.push(INLINE_DATA_VER);
let map_len = keys.len();
rmp::encode::write_map_len(&mut wr, map_len as u32)?;
for i in 0..map_len {
rmp::encode::write_str(&mut wr, keys[i].as_str())?;
rmp::encode::write_bin(&mut wr, values[i].as_slice())?;
}
self.0 = wr;
Ok(())
}
}
+23
View File
@@ -0,0 +1,23 @@
pub const AMZ_META_UNENCRYPTED_CONTENT_LENGTH: &str = "X-Amz-Meta-X-Amz-Unencrypted-Content-Length";
pub const AMZ_META_UNENCRYPTED_CONTENT_MD5: &str = "X-Amz-Meta-X-Amz-Unencrypted-Content-Md5";
pub const AMZ_STORAGE_CLASS: &str = "x-amz-storage-class";
pub const RESERVED_METADATA_PREFIX: &str = "X-RustFS-Internal-";
pub const RESERVED_METADATA_PREFIX_LOWER: &str = "x-rustfs-internal-";
pub const RUSTFS_HEALING: &str = "X-Rustfs-Internal-healing";
// pub const RUSTFS_DATA_MOVE: &str = "X-Rustfs-Internal-data-mov";
// pub const X_RUSTFS_INLINE_DATA: &str = "x-rustfs-inline-data";
pub const VERSION_PURGE_STATUS_KEY: &str = "X-Rustfs-Internal-purgestatus";
pub const X_RUSTFS_HEALING: &str = "X-Rustfs-Internal-healing";
pub const X_RUSTFS_DATA_MOV: &str = "X-Rustfs-Internal-data-mov";
pub const AMZ_OBJECT_TAGGING: &str = "X-Amz-Tagging";
pub const AMZ_BUCKET_REPLICATION_STATUS: &str = "X-Amz-Replication-Status";
pub const AMZ_DECODED_CONTENT_LENGTH: &str = "X-Amz-Decoded-Content-Length";
pub const RUSTFS_DATA_MOVE: &str = "X-Rustfs-Internal-data-mov";
+14
View File
@@ -0,0 +1,14 @@
mod error;
mod fileinfo;
mod filemeta;
mod filemeta_inline;
pub mod headers;
mod metacache;
pub mod test_data;
pub use error::*;
pub use fileinfo::*;
pub use filemeta::*;
pub use filemeta_inline::*;
pub use metacache::*;
+874
View File
@@ -0,0 +1,874 @@
use crate::error::{Error, Result};
use crate::{FileInfo, FileInfoVersions, FileMeta, FileMetaShallowVersion, VersionType, merge_file_meta_versions};
use rmp::Marker;
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::str::from_utf8;
use std::{
fmt::Debug,
future::Future,
pin::Pin,
ptr,
sync::{
Arc,
atomic::{AtomicPtr, AtomicU64, Ordering as AtomicOrdering},
},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use time::OffsetDateTime;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::spawn;
use tokio::sync::Mutex;
use tracing::warn;
const SLASH_SEPARATOR: &str = "/";
#[derive(Clone, Debug, Default)]
pub struct MetadataResolutionParams {
pub dir_quorum: usize,
pub obj_quorum: usize,
pub requested_versions: usize,
pub bucket: String,
pub strict: bool,
pub candidates: Vec<Vec<FileMetaShallowVersion>>,
}
#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
pub struct MetaCacheEntry {
/// name is the full name of the object including prefixes
pub name: String,
/// Metadata. If none is present it is not an object but only a prefix.
/// Entries without metadata will only be present in non-recursive scans.
pub metadata: Vec<u8>,
/// cached contains the metadata if decoded.
#[serde(skip)]
pub cached: Option<FileMeta>,
/// Indicates the entry can be reused and only one reference to metadata is expected.
pub reusable: bool,
}
impl MetaCacheEntry {
pub fn marshal_msg(&self) -> Result<Vec<u8>> {
let mut wr = Vec::new();
rmp::encode::write_bool(&mut wr, true)?;
rmp::encode::write_str(&mut wr, &self.name)?;
rmp::encode::write_bin(&mut wr, &self.metadata)?;
Ok(wr)
}
pub fn is_dir(&self) -> bool {
self.metadata.is_empty() && self.name.ends_with('/')
}
pub fn is_in_dir(&self, dir: &str, separator: &str) -> bool {
if dir.is_empty() {
let idx = self.name.find(separator);
return idx.is_none() || idx.unwrap() == self.name.len() - separator.len();
}
let ext = self.name.trim_start_matches(dir);
if ext.len() != self.name.len() {
let idx = ext.find(separator);
return idx.is_none() || idx.unwrap() == ext.len() - separator.len();
}
false
}
pub fn is_object(&self) -> bool {
!self.metadata.is_empty()
}
pub fn is_object_dir(&self) -> bool {
!self.metadata.is_empty() && self.name.ends_with(SLASH_SEPARATOR)
}
pub fn is_latest_delete_marker(&mut self) -> bool {
if let Some(cached) = &self.cached {
if cached.versions.is_empty() {
return true;
}
return cached.versions[0].header.version_type == VersionType::Delete;
}
if !FileMeta::is_xl2_v1_format(&self.metadata) {
return false;
}
match FileMeta::check_xl2_v1(&self.metadata) {
Ok((meta, _, _)) => {
if !meta.is_empty() {
return FileMeta::is_latest_delete_marker(meta);
}
}
Err(_) => return true,
}
match self.xl_meta() {
Ok(res) => {
if res.versions.is_empty() {
return true;
}
res.versions[0].header.version_type == VersionType::Delete
}
Err(_) => true,
}
}
#[tracing::instrument(level = "debug", skip(self))]
pub fn to_fileinfo(&self, bucket: &str) -> Result<FileInfo> {
if self.is_dir() {
return Ok(FileInfo {
volume: bucket.to_owned(),
name: self.name.clone(),
..Default::default()
});
}
if self.cached.is_some() {
let fm = self.cached.as_ref().unwrap();
if fm.versions.is_empty() {
return Ok(FileInfo {
volume: bucket.to_owned(),
name: self.name.clone(),
deleted: true,
is_latest: true,
mod_time: Some(OffsetDateTime::UNIX_EPOCH),
..Default::default()
});
}
let fi = fm.into_fileinfo(bucket, self.name.as_str(), "", false, false)?;
return Ok(fi);
}
let mut fm = FileMeta::new();
fm.unmarshal_msg(&self.metadata)?;
let fi = fm.into_fileinfo(bucket, self.name.as_str(), "", false, false)?;
Ok(fi)
}
pub fn file_info_versions(&self, bucket: &str) -> Result<FileInfoVersions> {
if self.is_dir() {
return Ok(FileInfoVersions {
volume: bucket.to_string(),
name: self.name.clone(),
versions: vec![FileInfo {
volume: bucket.to_string(),
name: self.name.clone(),
..Default::default()
}],
..Default::default()
});
}
let mut fm = FileMeta::new();
fm.unmarshal_msg(&self.metadata)?;
fm.into_file_info_versions(bucket, self.name.as_str(), false)
}
pub fn matches(&self, other: Option<&MetaCacheEntry>, strict: bool) -> (Option<MetaCacheEntry>, bool) {
if other.is_none() {
return (None, false);
}
let other = other.unwrap();
if self.name != other.name {
if self.name < other.name {
return (Some(self.clone()), false);
}
return (Some(other.clone()), false);
}
if other.is_dir() || self.is_dir() {
if self.is_dir() {
return (Some(self.clone()), other.is_dir() == self.is_dir());
}
return (Some(other.clone()), other.is_dir() == self.is_dir());
}
let self_vers = match &self.cached {
Some(file_meta) => file_meta.clone(),
None => match FileMeta::load(&self.metadata) {
Ok(meta) => meta,
Err(_) => return (None, false),
},
};
let other_vers = match &other.cached {
Some(file_meta) => file_meta.clone(),
None => match FileMeta::load(&other.metadata) {
Ok(meta) => meta,
Err(_) => return (None, false),
},
};
if self_vers.versions.len() != other_vers.versions.len() {
match self_vers.lastest_mod_time().cmp(&other_vers.lastest_mod_time()) {
Ordering::Greater => return (Some(self.clone()), false),
Ordering::Less => return (Some(other.clone()), false),
_ => {}
}
if self_vers.versions.len() > other_vers.versions.len() {
return (Some(self.clone()), false);
}
return (Some(other.clone()), false);
}
let mut prefer = None;
for (s_version, o_version) in self_vers.versions.iter().zip(other_vers.versions.iter()) {
if s_version.header != o_version.header {
if s_version.header.has_ec() != o_version.header.has_ec() {
// One version has EC and the other doesn't - may have been written later.
// Compare without considering EC.
let (mut a, mut b) = (s_version.header.clone(), o_version.header.clone());
(a.ec_n, a.ec_m, b.ec_n, b.ec_m) = (0, 0, 0, 0);
if a == b {
continue;
}
}
if !strict && s_version.header.matches_not_strict(&o_version.header) {
if prefer.is_none() {
if s_version.header.sorts_before(&o_version.header) {
prefer = Some(self.clone());
} else {
prefer = Some(other.clone());
}
}
continue;
}
if prefer.is_some() {
return (prefer, false);
}
if s_version.header.sorts_before(&o_version.header) {
return (Some(self.clone()), false);
}
return (Some(other.clone()), false);
}
}
if prefer.is_none() {
prefer = Some(self.clone());
}
(prefer, true)
}
pub fn xl_meta(&mut self) -> Result<FileMeta> {
if self.is_dir() {
return Err(Error::FileNotFound);
}
if let Some(meta) = &self.cached {
Ok(meta.clone())
} else {
if self.metadata.is_empty() {
return Err(Error::FileNotFound);
}
let meta = FileMeta::load(&self.metadata)?;
self.cached = Some(meta.clone());
Ok(meta)
}
}
}
#[derive(Debug, Default)]
pub struct MetaCacheEntries(pub Vec<Option<MetaCacheEntry>>);
impl MetaCacheEntries {
#[allow(clippy::should_implement_trait)]
pub fn as_ref(&self) -> &[Option<MetaCacheEntry>] {
&self.0
}
pub fn resolve(&self, mut params: MetadataResolutionParams) -> Option<MetaCacheEntry> {
if self.0.is_empty() {
warn!("decommission_pool: entries resolve empty");
return None;
}
let mut dir_exists = 0;
let mut selected = None;
params.candidates.clear();
let mut objs_agree = 0;
let mut objs_valid = 0;
for entry in self.0.iter().flatten() {
let mut entry = entry.clone();
warn!("decommission_pool: entries resolve entry {:?}", entry.name);
if entry.name.is_empty() {
continue;
}
if entry.is_dir() {
dir_exists += 1;
selected = Some(entry.clone());
warn!("decommission_pool: entries resolve entry dir {:?}", entry.name);
continue;
}
let xl = match entry.xl_meta() {
Ok(xl) => xl,
Err(e) => {
warn!("decommission_pool: entries resolve entry xl_meta {:?}", e);
continue;
}
};
objs_valid += 1;
params.candidates.push(xl.versions.clone());
if selected.is_none() {
selected = Some(entry.clone());
objs_agree = 1;
warn!("decommission_pool: entries resolve entry selected {:?}", entry.name);
continue;
}
if let (prefer, true) = entry.matches(selected.as_ref(), params.strict) {
selected = prefer;
objs_agree += 1;
warn!("decommission_pool: entries resolve entry prefer {:?}", entry.name);
continue;
}
}
let Some(selected) = selected else {
warn!("decommission_pool: entries resolve entry no selected");
return None;
};
if selected.is_dir() && dir_exists >= params.dir_quorum {
warn!("decommission_pool: entries resolve entry dir selected {:?}", selected.name);
return Some(selected);
}
// If we would never be able to reach read quorum.
if objs_valid < params.obj_quorum {
warn!(
"decommission_pool: entries resolve entry not enough objects {} < {}",
objs_valid, params.obj_quorum
);
return None;
}
if objs_agree == objs_valid {
warn!("decommission_pool: entries resolve entry all agree {} == {}", objs_agree, objs_valid);
return Some(selected);
}
let Some(cached) = selected.cached else {
warn!("decommission_pool: entries resolve entry no cached");
return None;
};
let versions = merge_file_meta_versions(params.obj_quorum, params.strict, params.requested_versions, &params.candidates);
if versions.is_empty() {
warn!("decommission_pool: entries resolve entry no versions");
return None;
}
let metadata = match cached.marshal_msg() {
Ok(meta) => meta,
Err(e) => {
warn!("decommission_pool: entries resolve entry marshal_msg {:?}", e);
return None;
}
};
// Merge if we have disagreement.
// Create a new merged result.
let new_selected = MetaCacheEntry {
name: selected.name.clone(),
cached: Some(FileMeta {
meta_ver: cached.meta_ver,
versions,
..Default::default()
}),
reusable: true,
metadata,
};
warn!("decommission_pool: entries resolve entry selected {:?}", new_selected.name);
Some(new_selected)
}
pub fn first_found(&self) -> (Option<MetaCacheEntry>, usize) {
(self.0.iter().find(|x| x.is_some()).cloned().unwrap_or_default(), self.0.len())
}
}
#[derive(Debug, Default)]
pub struct MetaCacheEntriesSortedResult {
pub entries: Option<MetaCacheEntriesSorted>,
pub err: Option<Error>,
}
#[derive(Debug, Default)]
pub struct MetaCacheEntriesSorted {
pub o: MetaCacheEntries,
pub list_id: Option<String>,
pub reuse: bool,
pub last_skipped_entry: Option<String>,
}
impl MetaCacheEntriesSorted {
pub fn entries(&self) -> Vec<&MetaCacheEntry> {
let entries: Vec<&MetaCacheEntry> = self.o.0.iter().flatten().collect();
entries
}
pub fn forward_past(&mut self, marker: Option<String>) {
if let Some(val) = marker {
if let Some(idx) = self.o.0.iter().flatten().position(|v| v.name > val) {
self.o.0 = self.o.0.split_off(idx);
}
}
}
}
const METACACHE_STREAM_VERSION: u8 = 2;
#[derive(Debug)]
pub struct MetacacheWriter<W> {
wr: W,
created: bool,
buf: Vec<u8>,
}
impl<W: AsyncWrite + Unpin> MetacacheWriter<W> {
pub fn new(wr: W) -> Self {
Self {
wr,
created: false,
buf: Vec::new(),
}
}
pub async fn flush(&mut self) -> Result<()> {
self.wr.write_all(&self.buf).await?;
self.buf.clear();
Ok(())
}
pub async fn init(&mut self) -> Result<()> {
if !self.created {
rmp::encode::write_u8(&mut self.buf, METACACHE_STREAM_VERSION).map_err(|e| Error::other(format!("{:?}", e)))?;
self.flush().await?;
self.created = true;
}
Ok(())
}
pub async fn write(&mut self, objs: &[MetaCacheEntry]) -> Result<()> {
if objs.is_empty() {
return Ok(());
}
self.init().await?;
for obj in objs.iter() {
if obj.name.is_empty() {
return Err(Error::other("metacacheWriter: no name"));
}
self.write_obj(obj).await?;
}
Ok(())
}
pub async fn write_obj(&mut self, obj: &MetaCacheEntry) -> Result<()> {
self.init().await?;
rmp::encode::write_bool(&mut self.buf, true).map_err(|e| Error::other(format!("{:?}", e)))?;
rmp::encode::write_str(&mut self.buf, &obj.name).map_err(|e| Error::other(format!("{:?}", e)))?;
rmp::encode::write_bin(&mut self.buf, &obj.metadata).map_err(|e| Error::other(format!("{:?}", e)))?;
self.flush().await?;
Ok(())
}
pub async fn close(&mut self) -> Result<()> {
rmp::encode::write_bool(&mut self.buf, false).map_err(|e| Error::other(format!("{:?}", e)))?;
self.flush().await?;
Ok(())
}
}
pub struct MetacacheReader<R> {
rd: R,
init: bool,
err: Option<Error>,
buf: Vec<u8>,
offset: usize,
current: Option<MetaCacheEntry>,
}
impl<R: AsyncRead + Unpin> MetacacheReader<R> {
pub fn new(rd: R) -> Self {
Self {
rd,
init: false,
err: None,
buf: Vec::new(),
offset: 0,
current: None,
}
}
pub async fn read_more(&mut self, read_size: usize) -> Result<&[u8]> {
let ext_size = read_size + self.offset;
let extra = ext_size - self.offset;
if self.buf.capacity() >= ext_size {
// Extend the buffer if we have enough space.
self.buf.resize(ext_size, 0);
} else {
self.buf.extend(vec![0u8; extra]);
}
let pref = self.offset;
self.rd.read_exact(&mut self.buf[pref..ext_size]).await?;
self.offset += read_size;
let data = &self.buf[pref..ext_size];
Ok(data)
}
fn reset(&mut self) {
self.buf.clear();
self.offset = 0;
}
async fn check_init(&mut self) -> Result<()> {
if !self.init {
let ver = match rmp::decode::read_u8(&mut self.read_more(2).await?) {
Ok(res) => res,
Err(err) => {
self.err = Some(Error::other(format!("{:?}", err)));
0
}
};
match ver {
1 | 2 => (),
_ => {
self.err = Some(Error::other("invalid version"));
}
}
self.init = true;
}
Ok(())
}
async fn read_str_len(&mut self) -> Result<u32> {
let mark = match rmp::decode::read_marker(&mut self.read_more(1).await?) {
Ok(res) => res,
Err(err) => {
let err: Error = err.into();
self.err = Some(err.clone());
return Err(err);
}
};
match mark {
Marker::FixStr(size) => Ok(u32::from(size)),
Marker::Str8 => Ok(u32::from(self.read_u8().await?)),
Marker::Str16 => Ok(u32::from(self.read_u16().await?)),
Marker::Str32 => Ok(self.read_u32().await?),
_marker => Err(Error::other("str marker err")),
}
}
async fn read_bin_len(&mut self) -> Result<u32> {
let mark = match rmp::decode::read_marker(&mut self.read_more(1).await?) {
Ok(res) => res,
Err(err) => {
let err: Error = err.into();
self.err = Some(err.clone());
return Err(err);
}
};
match mark {
Marker::Bin8 => Ok(u32::from(self.read_u8().await?)),
Marker::Bin16 => Ok(u32::from(self.read_u16().await?)),
Marker::Bin32 => Ok(self.read_u32().await?),
_ => Err(Error::other("bin marker err")),
}
}
async fn read_u8(&mut self) -> Result<u8> {
let buf = self.read_more(1).await?;
Ok(u8::from_be_bytes(buf.try_into().expect("Slice with incorrect length")))
}
async fn read_u16(&mut self) -> Result<u16> {
let buf = self.read_more(2).await?;
Ok(u16::from_be_bytes(buf.try_into().expect("Slice with incorrect length")))
}
async fn read_u32(&mut self) -> Result<u32> {
let buf = self.read_more(4).await?;
Ok(u32::from_be_bytes(buf.try_into().expect("Slice with incorrect length")))
}
pub async fn skip(&mut self, size: usize) -> Result<()> {
self.check_init().await?;
if let Some(err) = &self.err {
return Err(err.clone());
}
let mut n = size;
if self.current.is_some() {
n -= 1;
self.current = None;
}
while n > 0 {
match rmp::decode::read_bool(&mut self.read_more(1).await?) {
Ok(res) => {
if !res {
return Ok(());
}
}
Err(err) => {
let err: Error = err.into();
self.err = Some(err.clone());
return Err(err);
}
};
let l = self.read_str_len().await?;
let _ = self.read_more(l as usize).await?;
let l = self.read_bin_len().await?;
let _ = self.read_more(l as usize).await?;
n -= 1;
}
Ok(())
}
pub async fn peek(&mut self) -> Result<Option<MetaCacheEntry>> {
self.check_init().await?;
if let Some(err) = &self.err {
return Err(err.clone());
}
match rmp::decode::read_bool(&mut self.read_more(1).await?) {
Ok(res) => {
if !res {
return Ok(None);
}
}
Err(err) => {
let err: Error = err.into();
self.err = Some(err.clone());
return Err(err);
}
};
let l = self.read_str_len().await?;
let buf = self.read_more(l as usize).await?;
let name_buf = buf.to_vec();
let name = match from_utf8(&name_buf) {
Ok(decoded) => decoded.to_owned(),
Err(err) => {
self.err = Some(Error::other(err.to_string()));
return Err(Error::other(err.to_string()));
}
};
let l = self.read_bin_len().await?;
let buf = self.read_more(l as usize).await?;
let metadata = buf.to_vec();
self.reset();
let entry = Some(MetaCacheEntry {
name,
metadata,
cached: None,
reusable: false,
});
self.current = entry.clone();
Ok(entry)
}
pub async fn read_all(&mut self) -> Result<Vec<MetaCacheEntry>> {
let mut ret = Vec::new();
loop {
if let Some(entry) = self.peek().await? {
ret.push(entry);
continue;
}
break;
}
Ok(ret)
}
}
pub type UpdateFn<T> = Box<dyn Fn() -> Pin<Box<dyn Future<Output = std::io::Result<T>> + Send>> + Send + Sync + 'static>;
#[derive(Clone, Debug, Default)]
pub struct Opts {
pub return_last_good: bool,
pub no_wait: bool,
}
pub struct Cache<T: Clone + Debug + Send> {
update_fn: UpdateFn<T>,
ttl: Duration,
opts: Opts,
val: AtomicPtr<T>,
last_update_ms: AtomicU64,
updating: Arc<Mutex<bool>>,
}
impl<T: Clone + Debug + Send + 'static> Cache<T> {
pub fn new(update_fn: UpdateFn<T>, ttl: Duration, opts: Opts) -> Self {
let val = AtomicPtr::new(ptr::null_mut());
Self {
update_fn,
ttl,
opts,
val,
last_update_ms: AtomicU64::new(0),
updating: Arc::new(Mutex::new(false)),
}
}
#[allow(unsafe_code)]
pub async fn get(self: Arc<Self>) -> std::io::Result<T> {
let v_ptr = self.val.load(AtomicOrdering::SeqCst);
let v = if v_ptr.is_null() {
None
} else {
Some(unsafe { (*v_ptr).clone() })
};
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("Time went backwards")
.as_secs();
if now - self.last_update_ms.load(AtomicOrdering::SeqCst) < self.ttl.as_secs() {
if let Some(v) = v {
return Ok(v);
}
}
if self.opts.no_wait && v.is_some() && now - self.last_update_ms.load(AtomicOrdering::SeqCst) < self.ttl.as_secs() * 2 {
if self.updating.try_lock().is_ok() {
let this = Arc::clone(&self);
spawn(async move {
let _ = this.update().await;
});
}
return Ok(v.unwrap());
}
let _ = self.updating.lock().await;
if let Ok(duration) =
SystemTime::now().duration_since(UNIX_EPOCH + Duration::from_secs(self.last_update_ms.load(AtomicOrdering::SeqCst)))
{
if duration < self.ttl {
return Ok(v.unwrap());
}
}
match self.update().await {
Ok(_) => {
let v_ptr = self.val.load(AtomicOrdering::SeqCst);
let v = if v_ptr.is_null() {
None
} else {
Some(unsafe { (*v_ptr).clone() })
};
Ok(v.unwrap())
}
Err(err) => Err(err),
}
}
async fn update(&self) -> std::io::Result<()> {
match (self.update_fn)().await {
Ok(val) => {
self.val.store(Box::into_raw(Box::new(val)), AtomicOrdering::SeqCst);
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("Time went backwards")
.as_secs();
self.last_update_ms.store(now, AtomicOrdering::SeqCst);
Ok(())
}
Err(err) => {
let v_ptr = self.val.load(AtomicOrdering::SeqCst);
if self.opts.return_last_good && !v_ptr.is_null() {
return Ok(());
}
Err(err)
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
#[tokio::test]
async fn test_writer() {
let mut f = Cursor::new(Vec::new());
let mut w = MetacacheWriter::new(&mut f);
let mut objs = Vec::new();
for i in 0..10 {
let info = MetaCacheEntry {
name: format!("item{}", i),
metadata: vec![0u8, 10],
cached: None,
reusable: false,
};
objs.push(info);
}
w.write(&objs).await.unwrap();
w.close().await.unwrap();
let data = f.into_inner();
let nf = Cursor::new(data);
let mut r = MetacacheReader::new(nf);
let nobjs = r.read_all().await.unwrap();
assert_eq!(objs, nobjs);
}
}
+292
View File
@@ -0,0 +1,292 @@
use crate::error::Result;
use crate::filemeta::*;
use std::collections::HashMap;
use time::OffsetDateTime;
use uuid::Uuid;
/// 创建一个真实的 xl.meta 文件数据用于测试
pub fn create_real_xlmeta() -> Result<Vec<u8>> {
let mut fm = FileMeta::new();
// 创建一个真实的对象版本
let version_id = Uuid::parse_str("01234567-89ab-cdef-0123-456789abcdef")?;
let data_dir = Uuid::parse_str("fedcba98-7654-3210-fedc-ba9876543210")?;
let mut metadata = HashMap::new();
metadata.insert("Content-Type".to_string(), "text/plain".to_string());
metadata.insert("X-Amz-Meta-Author".to_string(), "test-user".to_string());
metadata.insert("X-Amz-Meta-Created".to_string(), "2024-01-15T10:30:00Z".to_string());
let object_version = MetaObject {
version_id: Some(version_id),
data_dir: Some(data_dir),
erasure_algorithm: crate::fileinfo::ErasureAlgo::ReedSolomon,
erasure_m: 4,
erasure_n: 2,
erasure_block_size: 1024 * 1024, // 1MB
erasure_index: 1,
erasure_dist: vec![0, 1, 2, 3, 4, 5],
bitrot_checksum_algo: ChecksumAlgo::HighwayHash,
part_numbers: vec![1],
part_etags: vec!["d41d8cd98f00b204e9800998ecf8427e".to_string()],
part_sizes: vec![1024],
part_actual_sizes: vec![1024],
part_indices: Vec::new(),
size: 1024,
mod_time: Some(OffsetDateTime::from_unix_timestamp(1705312200)?), // 2024-01-15 10:30:00 UTC
meta_sys: HashMap::new(),
meta_user: metadata,
};
let file_version = FileMetaVersion {
version_type: VersionType::Object,
object: Some(object_version),
delete_marker: None,
write_version: 1,
};
let shallow_version = FileMetaShallowVersion::try_from(file_version)?;
fm.versions.push(shallow_version);
// 添加一个删除标记版本
let delete_version_id = Uuid::parse_str("11111111-2222-3333-4444-555555555555")?;
let delete_marker = MetaDeleteMarker {
version_id: Some(delete_version_id),
mod_time: Some(OffsetDateTime::from_unix_timestamp(1705312260)?), // 1分钟后
meta_sys: None,
};
let delete_file_version = FileMetaVersion {
version_type: VersionType::Delete,
object: None,
delete_marker: Some(delete_marker),
write_version: 2,
};
let delete_shallow_version = FileMetaShallowVersion::try_from(delete_file_version)?;
fm.versions.push(delete_shallow_version);
// 添加一个 Legacy 版本用于测试
let legacy_version_id = Uuid::parse_str("aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee")?;
let legacy_version = FileMetaVersion {
version_type: VersionType::Legacy,
object: None,
delete_marker: None,
write_version: 3,
};
let mut legacy_shallow = FileMetaShallowVersion::try_from(legacy_version)?;
legacy_shallow.header.version_id = Some(legacy_version_id);
legacy_shallow.header.mod_time = Some(OffsetDateTime::from_unix_timestamp(1705312140)?); // 更早的时间
fm.versions.push(legacy_shallow);
// 按修改时间排序(最新的在前)
fm.versions.sort_by(|a, b| b.header.mod_time.cmp(&a.header.mod_time));
fm.marshal_msg()
}
/// 创建一个包含多个版本的复杂 xl.meta 文件
pub fn create_complex_xlmeta() -> Result<Vec<u8>> {
let mut fm = FileMeta::new();
// 创建10个版本的对象
for i in 0i64..10i64 {
let version_id = Uuid::new_v4();
let data_dir = if i % 3 == 0 { Some(Uuid::new_v4()) } else { None };
let mut metadata = HashMap::new();
metadata.insert("Content-Type".to_string(), "application/octet-stream".to_string());
metadata.insert("X-Amz-Meta-Version".to_string(), i.to_string());
metadata.insert("X-Amz-Meta-Test".to_string(), format!("test-value-{}", i));
let object_version = MetaObject {
version_id: Some(version_id),
data_dir,
erasure_algorithm: crate::fileinfo::ErasureAlgo::ReedSolomon,
erasure_m: 4,
erasure_n: 2,
erasure_block_size: 1024 * 1024,
erasure_index: (i % 6) as usize,
erasure_dist: vec![0, 1, 2, 3, 4, 5],
bitrot_checksum_algo: ChecksumAlgo::HighwayHash,
part_numbers: vec![1],
part_etags: vec![format!("etag-{:08x}", i)],
part_sizes: vec![1024 * (i + 1) as usize],
part_actual_sizes: vec![1024 * (i + 1)],
part_indices: Vec::new(),
size: 1024 * (i + 1),
mod_time: Some(OffsetDateTime::from_unix_timestamp(1705312200 + i * 60)?),
meta_sys: HashMap::new(),
meta_user: metadata,
};
let file_version = FileMetaVersion {
version_type: VersionType::Object,
object: Some(object_version),
delete_marker: None,
write_version: (i + 1) as u64,
};
let shallow_version = FileMetaShallowVersion::try_from(file_version)?;
fm.versions.push(shallow_version);
// 每隔3个版本添加一个删除标记
if i % 3 == 2 {
let delete_version_id = Uuid::new_v4();
let delete_marker = MetaDeleteMarker {
version_id: Some(delete_version_id),
mod_time: Some(OffsetDateTime::from_unix_timestamp(1705312200 + i * 60 + 30)?),
meta_sys: None,
};
let delete_file_version = FileMetaVersion {
version_type: VersionType::Delete,
object: None,
delete_marker: Some(delete_marker),
write_version: (i + 100) as u64,
};
let delete_shallow_version = FileMetaShallowVersion::try_from(delete_file_version)?;
fm.versions.push(delete_shallow_version);
}
}
// 按修改时间排序(最新的在前)
fm.versions.sort_by(|a, b| b.header.mod_time.cmp(&a.header.mod_time));
fm.marshal_msg()
}
/// 创建一个损坏的 xl.meta 文件用于错误处理测试
pub fn create_corrupted_xlmeta() -> Vec<u8> {
let mut data = vec![
// 正确的文件头
b'X', b'L', b'2', b' ', // 版本号
1, 0, 3, 0, // 版本号
0xc6, 0x00, 0x00, 0x00, 0x10, // 正确的 bin32 长度标记,但数据长度不匹配
];
// 添加不足的数据(少于声明的长度)
data.extend_from_slice(&[0x42; 8]); // 只有8字节,但声明了16字节
data
}
/// 创建一个空的 xl.meta 文件
pub fn create_empty_xlmeta() -> Result<Vec<u8>> {
let fm = FileMeta::new();
fm.marshal_msg()
}
/// 验证解析结果的辅助函数
pub fn verify_parsed_metadata(fm: &FileMeta, expected_versions: usize) -> Result<()> {
assert_eq!(fm.versions.len(), expected_versions, "版本数量不匹配");
assert_eq!(fm.meta_ver, crate::filemeta::XL_META_VERSION, "元数据版本不匹配");
// 验证版本是否按修改时间排序
for i in 1..fm.versions.len() {
let prev_time = fm.versions[i - 1].header.mod_time;
let curr_time = fm.versions[i].header.mod_time;
if let (Some(prev), Some(curr)) = (prev_time, curr_time) {
assert!(prev >= curr, "版本未按修改时间正确排序");
}
}
Ok(())
}
/// 创建一个包含内联数据的 xl.meta 文件
pub fn create_xlmeta_with_inline_data() -> Result<Vec<u8>> {
let mut fm = FileMeta::new();
// 添加内联数据
let inline_data = b"This is inline data for testing purposes";
let version_id = Uuid::new_v4();
fm.data.replace(&version_id.to_string(), inline_data.to_vec())?;
let object_version = MetaObject {
version_id: Some(version_id),
data_dir: None,
erasure_algorithm: crate::fileinfo::ErasureAlgo::ReedSolomon,
erasure_m: 1,
erasure_n: 1,
erasure_block_size: 64 * 1024,
erasure_index: 0,
erasure_dist: vec![0, 1],
bitrot_checksum_algo: ChecksumAlgo::HighwayHash,
part_numbers: vec![1],
part_etags: Vec::new(),
part_sizes: vec![inline_data.len()],
part_actual_sizes: Vec::new(),
part_indices: Vec::new(),
size: inline_data.len() as i64,
mod_time: Some(OffsetDateTime::now_utc()),
meta_sys: HashMap::new(),
meta_user: HashMap::new(),
};
let file_version = FileMetaVersion {
version_type: VersionType::Object,
object: Some(object_version),
delete_marker: None,
write_version: 1,
};
let shallow_version = FileMetaShallowVersion::try_from(file_version)?;
fm.versions.push(shallow_version);
fm.marshal_msg()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_create_real_xlmeta() {
let data = create_real_xlmeta().expect("创建测试数据失败");
assert!(!data.is_empty(), "生成的数据不应为空");
// 验证文件头
assert_eq!(&data[0..4], b"XL2 ", "文件头不正确");
// 尝试解析
let fm = FileMeta::load(&data).expect("解析失败");
verify_parsed_metadata(&fm, 3).expect("验证失败");
}
#[test]
fn test_create_complex_xlmeta() {
let data = create_complex_xlmeta().expect("创建复杂测试数据失败");
assert!(!data.is_empty(), "生成的数据不应为空");
let fm = FileMeta::load(&data).expect("解析失败");
assert!(fm.versions.len() >= 10, "应该有至少10个版本");
}
#[test]
fn test_create_xlmeta_with_inline_data() {
let data = create_xlmeta_with_inline_data().expect("创建内联数据测试失败");
assert!(!data.is_empty(), "生成的数据不应为空");
let fm = FileMeta::load(&data).expect("解析失败");
assert_eq!(fm.versions.len(), 1, "应该有1个版本");
assert!(!fm.data.as_slice().is_empty(), "应该包含内联数据");
}
#[test]
fn test_corrupted_xlmeta_handling() {
let data = create_corrupted_xlmeta();
let result = FileMeta::load(&data);
assert!(result.is_err(), "损坏的数据应该解析失败");
}
#[test]
fn test_empty_xlmeta() {
let data = create_empty_xlmeta().expect("创建空测试数据失败");
let fm = FileMeta::load(&data).expect("解析空数据失败");
assert_eq!(fm.versions.len(), 0, "空文件应该没有版本");
}
}
+1 -1
View File
@@ -1,7 +1,7 @@
use crate::Error;
use serde::{Deserialize, Serialize};
use serde_with::{DeserializeFromStr, SerializeDisplay};
use smallvec::{smallvec, SmallVec};
use smallvec::{SmallVec, smallvec};
use std::borrow::Cow;
use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};
+1 -1
View File
@@ -1,5 +1,5 @@
use opentelemetry::global;
use rustfs_obs::{get_logger, init_obs, log_info, BaseLogEntry, ServerLogEntry, SystemObserver};
use rustfs_obs::{BaseLogEntry, ServerLogEntry, SystemObserver, get_logger, init_obs, log_info};
use std::collections::HashMap;
use std::time::{Duration, SystemTime};
use tracing::{error, info, instrument};
+4 -4
View File
@@ -1,6 +1,6 @@
use crate::logger::InitLogStatus;
use crate::telemetry::{init_telemetry, OtelGuard};
use crate::{get_global_logger, init_global_logger, AppConfig, Logger};
use crate::telemetry::{OtelGuard, init_telemetry};
use crate::{AppConfig, Logger, get_global_logger, init_global_logger};
use std::sync::{Arc, Mutex};
use tokio::sync::{OnceCell, SetError};
use tracing::{error, info};
@@ -102,8 +102,8 @@ pub fn get_logger() -> &'static Arc<tokio::sync::Mutex<Logger>> {
/// ```rust
/// use rustfs_obs::{ init_obs, set_global_guard};
///
/// fn init() -> Result<(), Box<dyn std::error::Error>> {
/// let guard = init_obs(None);
/// async fn init() -> Result<(), Box<dyn std::error::Error>> {
/// let (_, guard) = init_obs(None).await;
/// set_global_guard(guard)?;
/// Ok(())
/// }
+1 -1
View File
@@ -1,6 +1,6 @@
use crate::sinks::Sink;
use crate::{
sinks, AppConfig, AuditLogEntry, BaseLogEntry, ConsoleLogEntry, GlobalError, OtelConfig, ServerLogEntry, UnifiedLogEntry,
AppConfig, AuditLogEntry, BaseLogEntry, ConsoleLogEntry, GlobalError, OtelConfig, ServerLogEntry, UnifiedLogEntry, sinks,
};
use rustfs_config::{APP_NAME, ENVIRONMENT, SERVICE_VERSION};
use std::sync::Arc;
+1 -1
View File
@@ -1,5 +1,5 @@
use crate::sinks::Sink;
use crate::UnifiedLogEntry;
use crate::sinks::Sink;
use async_trait::async_trait;
/// Webhook Sink Implementation
+3 -3
View File
@@ -1,11 +1,11 @@
use crate::GlobalError;
use crate::system::attributes::ProcessAttributes;
use crate::system::gpu::GpuCollector;
use crate::system::metrics::{Metrics, DIRECTION, INTERFACE, STATUS};
use crate::GlobalError;
use crate::system::metrics::{DIRECTION, INTERFACE, Metrics, STATUS};
use opentelemetry::KeyValue;
use std::time::SystemTime;
use sysinfo::{Networks, Pid, ProcessStatus, System};
use tokio::time::{sleep, Duration};
use tokio::time::{Duration, sleep};
/// Collector is responsible for collecting system metrics and attributes.
/// It uses the sysinfo crate to gather information about the system and processes.
+3 -3
View File
@@ -1,14 +1,14 @@
#[cfg(feature = "gpu")]
use crate::GlobalError;
#[cfg(feature = "gpu")]
use crate::system::attributes::ProcessAttributes;
#[cfg(feature = "gpu")]
use crate::system::metrics::Metrics;
#[cfg(feature = "gpu")]
use crate::GlobalError;
use nvml_wrapper::Nvml;
#[cfg(feature = "gpu")]
use nvml_wrapper::enums::device::UsedGpuMemory;
#[cfg(feature = "gpu")]
use nvml_wrapper::Nvml;
#[cfg(feature = "gpu")]
use sysinfo::Pid;
#[cfg(feature = "gpu")]
use tracing::warn;
+5 -5
View File
@@ -1,19 +1,19 @@
use crate::OtelConfig;
use flexi_logger::{style, Age, Cleanup, Criterion, DeferredNow, FileSpec, LogSpecification, Naming, Record, WriteMode};
use flexi_logger::{Age, Cleanup, Criterion, DeferredNow, FileSpec, LogSpecification, Naming, Record, WriteMode, style};
use nu_ansi_term::Color;
use opentelemetry::trace::TracerProvider;
use opentelemetry::{global, KeyValue};
use opentelemetry::{KeyValue, global};
use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::logs::SdkLoggerProvider;
use opentelemetry_sdk::{
Resource,
metrics::{MeterProviderBuilder, PeriodicReader, SdkMeterProvider},
trace::{RandomIdGenerator, Sampler, SdkTracerProvider},
Resource,
};
use opentelemetry_semantic_conventions::{
attribute::{DEPLOYMENT_ENVIRONMENT_NAME, NETWORK_LOCAL_ADDRESS, SERVICE_VERSION as OTEL_SERVICE_VERSION},
SCHEMA_URL,
attribute::{DEPLOYMENT_ENVIRONMENT_NAME, NETWORK_LOCAL_ADDRESS, SERVICE_VERSION as OTEL_SERVICE_VERSION},
};
use rustfs_config::{
APP_NAME, DEFAULT_LOG_DIR, DEFAULT_LOG_KEEP_FILES, DEFAULT_LOG_LEVEL, ENVIRONMENT, METER_INTERVAL, SAMPLE_RATIO,
@@ -27,7 +27,7 @@ use tracing::info;
use tracing_error::ErrorLayer;
use tracing_opentelemetry::{MetricsLayer, OpenTelemetryLayer};
use tracing_subscriber::fmt::format::FmtSpan;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer};
use tracing_subscriber::{EnvFilter, Layer, layer::SubscriberExt, util::SubscriberInitExt};
/// A guard object that manages the lifecycle of OpenTelemetry components.
///
+1 -1
View File
@@ -1,4 +1,4 @@
use crate::{sinks::Sink, UnifiedLogEntry};
use crate::{UnifiedLogEntry, sinks::Sink};
use std::sync::Arc;
use tokio::sync::mpsc::Receiver;
+34
View File
@@ -0,0 +1,34 @@
[package]
name = "rustfs-rio"
edition.workspace = true
license.workspace = true
repository.workspace = true
rust-version.workspace = true
version.workspace = true
[lints]
workspace = true
[dependencies]
tokio = { workspace = true, features = ["full"] }
rand = { workspace = true }
md-5 = { workspace = true }
http.workspace = true
aes-gcm = "0.10.3"
crc32fast = "1.4.2"
pin-project-lite.workspace = true
async-trait.workspace = true
base64-simd = "0.8.0"
hex-simd = "0.8.0"
serde = { workspace = true }
bytes.workspace = true
reqwest.workspace = true
tokio-util.workspace = true
futures.workspace = true
rustfs-utils = {workspace = true, features= ["io","hash","compress"]}
byteorder.workspace = true
serde_json.workspace = true
[dev-dependencies]
criterion = { version = "0.5.1", features = ["async", "async_tokio", "tokio"] }
tokio-test = "0.4"
+672
View File
@@ -0,0 +1,672 @@
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use std::io::{self, Read, Seek, SeekFrom};
const S2_INDEX_HEADER: &[u8] = b"s2idx\x00";
const S2_INDEX_TRAILER: &[u8] = b"\x00xdi2s";
const MAX_INDEX_ENTRIES: usize = 1 << 16;
const MIN_INDEX_DIST: i64 = 1 << 20;
// const MIN_INDEX_DIST: i64 = 0;
pub trait TryGetIndex {
fn try_get_index(&self) -> Option<&Index> {
None
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Index {
pub total_uncompressed: i64,
pub total_compressed: i64,
info: Vec<IndexInfo>,
est_block_uncomp: i64,
}
impl Default for Index {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexInfo {
pub compressed_offset: i64,
pub uncompressed_offset: i64,
}
#[allow(dead_code)]
impl Index {
pub fn new() -> Self {
Self {
total_uncompressed: -1,
total_compressed: -1,
info: Vec::new(),
est_block_uncomp: 0,
}
}
#[allow(dead_code)]
fn reset(&mut self, max_block: usize) {
self.est_block_uncomp = max_block as i64;
self.total_compressed = -1;
self.total_uncompressed = -1;
self.info.clear();
}
pub fn len(&self) -> usize {
self.info.len()
}
fn alloc_infos(&mut self, n: usize) {
if n > MAX_INDEX_ENTRIES {
panic!("n > MAX_INDEX_ENTRIES");
}
self.info = Vec::with_capacity(n);
}
pub fn add(&mut self, compressed_offset: i64, uncompressed_offset: i64) -> io::Result<()> {
if self.info.is_empty() {
self.info.push(IndexInfo {
compressed_offset,
uncompressed_offset,
});
return Ok(());
}
let last_idx = self.info.len() - 1;
let latest = &mut self.info[last_idx];
if latest.uncompressed_offset == uncompressed_offset {
latest.compressed_offset = compressed_offset;
return Ok(());
}
if latest.uncompressed_offset > uncompressed_offset {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"internal error: Earlier uncompressed received ({} > {})",
latest.uncompressed_offset, uncompressed_offset
),
));
}
if latest.compressed_offset > compressed_offset {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"internal error: Earlier compressed received ({} > {})",
latest.uncompressed_offset, uncompressed_offset
),
));
}
if latest.uncompressed_offset + MIN_INDEX_DIST > uncompressed_offset {
return Ok(());
}
self.info.push(IndexInfo {
compressed_offset,
uncompressed_offset,
});
self.total_compressed = compressed_offset;
self.total_uncompressed = uncompressed_offset;
Ok(())
}
pub fn find(&self, offset: i64) -> io::Result<(i64, i64)> {
if self.total_uncompressed < 0 {
return Err(io::Error::other("corrupt index"));
}
let mut offset = offset;
if offset < 0 {
offset += self.total_uncompressed;
if offset < 0 {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "offset out of bounds"));
}
}
if offset > self.total_uncompressed {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "offset out of bounds"));
}
if self.info.is_empty() {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "empty index"));
}
if self.info.len() > 200 {
let n = self
.info
.binary_search_by(|info| {
if info.uncompressed_offset > offset {
std::cmp::Ordering::Greater
} else {
std::cmp::Ordering::Less
}
})
.unwrap_or_else(|i| i);
if n == 0 {
return Ok((self.info[0].compressed_offset, self.info[0].uncompressed_offset));
}
return Ok((self.info[n - 1].compressed_offset, self.info[n - 1].uncompressed_offset));
}
let mut compressed_off = 0;
let mut uncompressed_off = 0;
for info in &self.info {
if info.uncompressed_offset > offset {
break;
}
compressed_off = info.compressed_offset;
uncompressed_off = info.uncompressed_offset;
}
Ok((compressed_off, uncompressed_off))
}
fn reduce(&mut self) {
if self.info.len() < MAX_INDEX_ENTRIES && self.est_block_uncomp >= MIN_INDEX_DIST {
return;
}
let mut remove_n = (self.info.len() + 1) / MAX_INDEX_ENTRIES;
let src = self.info.clone();
let mut j = 0;
while self.est_block_uncomp * (remove_n as i64 + 1) < MIN_INDEX_DIST && self.info.len() / (remove_n + 1) > 1000 {
remove_n += 1;
}
let mut idx = 0;
while idx < src.len() {
self.info[j] = src[idx].clone();
j += 1;
idx += remove_n + 1;
}
self.info.truncate(j);
self.est_block_uncomp += self.est_block_uncomp * remove_n as i64;
}
pub fn into_vec(mut self) -> Bytes {
let mut b = Vec::new();
self.append_to(&mut b, self.total_uncompressed, self.total_compressed);
Bytes::from(b)
}
pub fn append_to(&mut self, b: &mut Vec<u8>, uncomp_total: i64, comp_total: i64) {
self.reduce();
let init_size = b.len();
// Add skippable header
b.extend_from_slice(&[0x50, 0x2A, 0x4D, 0x18]); // ChunkTypeIndex
b.extend_from_slice(&[0, 0, 0]); // Placeholder for chunk length
// Add header
b.extend_from_slice(S2_INDEX_HEADER);
// Add total sizes
let mut tmp = [0u8; 8];
let n = write_varint(&mut tmp, uncomp_total);
b.extend_from_slice(&tmp[..n]);
let n = write_varint(&mut tmp, comp_total);
b.extend_from_slice(&tmp[..n]);
let n = write_varint(&mut tmp, self.est_block_uncomp);
b.extend_from_slice(&tmp[..n]);
let n = write_varint(&mut tmp, self.info.len() as i64);
b.extend_from_slice(&tmp[..n]);
// Check if we should add uncompressed offsets
let mut has_uncompressed = 0u8;
for (idx, info) in self.info.iter().enumerate() {
if idx == 0 {
if info.uncompressed_offset != 0 {
has_uncompressed = 1;
break;
}
continue;
}
if info.uncompressed_offset != self.info[idx - 1].uncompressed_offset + self.est_block_uncomp {
has_uncompressed = 1;
break;
}
}
b.push(has_uncompressed);
// Add uncompressed offsets if needed
if has_uncompressed == 1 {
for (idx, info) in self.info.iter().enumerate() {
let mut u_off = info.uncompressed_offset;
if idx > 0 {
let prev = &self.info[idx - 1];
u_off -= prev.uncompressed_offset + self.est_block_uncomp;
}
let n = write_varint(&mut tmp, u_off);
b.extend_from_slice(&tmp[..n]);
}
}
// Add compressed offsets
let mut c_predict = self.est_block_uncomp / 2;
for (idx, info) in self.info.iter().enumerate() {
let mut c_off = info.compressed_offset;
if idx > 0 {
let prev = &self.info[idx - 1];
c_off -= prev.compressed_offset + c_predict;
c_predict += c_off / 2;
}
let n = write_varint(&mut tmp, c_off);
b.extend_from_slice(&tmp[..n]);
}
// Add total size and trailer
let total_size = (b.len() - init_size + 4 + S2_INDEX_TRAILER.len()) as u32;
b.extend_from_slice(&total_size.to_le_bytes());
b.extend_from_slice(S2_INDEX_TRAILER);
// Update chunk length
let chunk_len = b.len() - init_size - 4;
b[init_size + 1] = chunk_len as u8;
b[init_size + 2] = (chunk_len >> 8) as u8;
b[init_size + 3] = (chunk_len >> 16) as u8;
}
pub fn load<'a>(&mut self, mut b: &'a [u8]) -> io::Result<&'a [u8]> {
if b.len() <= 4 + S2_INDEX_HEADER.len() + S2_INDEX_TRAILER.len() {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buffer too small"));
}
if b[0] != 0x50 || b[1] != 0x2A || b[2] != 0x4D || b[3] != 0x18 {
return Err(io::Error::other("invalid chunk type"));
}
let chunk_len = (b[1] as usize) | ((b[2] as usize) << 8) | ((b[3] as usize) << 16);
b = &b[4..];
if b.len() < chunk_len {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buffer too small"));
}
if !b.starts_with(S2_INDEX_HEADER) {
return Err(io::Error::other("invalid header"));
}
b = &b[S2_INDEX_HEADER.len()..];
// Read total uncompressed
let (v, n) = read_varint(b)?;
if v < 0 {
return Err(io::Error::other("invalid uncompressed size"));
}
self.total_uncompressed = v;
b = &b[n..];
// Read total compressed
let (v, n) = read_varint(b)?;
if v < 0 {
return Err(io::Error::other("invalid compressed size"));
}
self.total_compressed = v;
b = &b[n..];
// Read est block uncomp
let (v, n) = read_varint(b)?;
if v < 0 {
return Err(io::Error::other("invalid block size"));
}
self.est_block_uncomp = v;
b = &b[n..];
// Read number of entries
let (v, n) = read_varint(b)?;
if v < 0 || v > MAX_INDEX_ENTRIES as i64 {
return Err(io::Error::other("invalid number of entries"));
}
let entries = v as usize;
b = &b[n..];
self.alloc_infos(entries);
if b.is_empty() {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buffer too small"));
}
let has_uncompressed = b[0];
b = &b[1..];
if has_uncompressed & 1 != has_uncompressed {
return Err(io::Error::other("invalid uncompressed flag"));
}
// Read uncompressed offsets
for idx in 0..entries {
let mut u_off = 0i64;
if has_uncompressed != 0 {
let (v, n) = read_varint(b)?;
u_off = v;
b = &b[n..];
}
if idx > 0 {
let prev = self.info[idx - 1].uncompressed_offset;
u_off += prev + self.est_block_uncomp;
if u_off <= prev {
return Err(io::Error::other("invalid offset"));
}
}
if u_off < 0 {
return Err(io::Error::other("negative offset"));
}
self.info[idx].uncompressed_offset = u_off;
}
// Read compressed offsets
let mut c_predict = self.est_block_uncomp / 2;
for idx in 0..entries {
let (v, n) = read_varint(b)?;
let mut c_off = v;
b = &b[n..];
if idx > 0 {
c_predict += c_off / 2;
let prev = self.info[idx - 1].compressed_offset;
c_off += prev + c_predict;
if c_off <= prev {
return Err(io::Error::other("invalid offset"));
}
}
if c_off < 0 {
return Err(io::Error::other("negative offset"));
}
self.info[idx].compressed_offset = c_off;
}
if b.len() < 4 + S2_INDEX_TRAILER.len() {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buffer too small"));
}
// Skip size
b = &b[4..];
// Check trailer
if !b.starts_with(S2_INDEX_TRAILER) {
return Err(io::Error::other("invalid trailer"));
}
Ok(&b[S2_INDEX_TRAILER.len()..])
}
pub fn load_stream<R: Read + Seek>(&mut self, mut rs: R) -> io::Result<()> {
// Go to end
rs.seek(SeekFrom::End(-10))?;
let mut tmp = [0u8; 10];
rs.read_exact(&mut tmp)?;
// Check trailer
if &tmp[4..4 + S2_INDEX_TRAILER.len()] != S2_INDEX_TRAILER {
return Err(io::Error::other("invalid trailer"));
}
let sz = u32::from_le_bytes(tmp[..4].try_into().unwrap());
if sz > 0x7fffffff {
return Err(io::Error::other("size too large"));
}
rs.seek(SeekFrom::End(-(sz as i64)))?;
let mut buf = vec![0u8; sz as usize];
rs.read_exact(&mut buf)?;
self.load(&buf)?;
Ok(())
}
pub fn to_json(&self) -> serde_json::Result<Vec<u8>> {
#[derive(Serialize)]
struct Offset {
compressed: i64,
uncompressed: i64,
}
#[derive(Serialize)]
struct IndexJson {
total_uncompressed: i64,
total_compressed: i64,
offsets: Vec<Offset>,
est_block_uncompressed: i64,
}
let json = IndexJson {
total_uncompressed: self.total_uncompressed,
total_compressed: self.total_compressed,
offsets: self
.info
.iter()
.map(|info| Offset {
compressed: info.compressed_offset,
uncompressed: info.uncompressed_offset,
})
.collect(),
est_block_uncompressed: self.est_block_uncomp,
};
serde_json::to_vec_pretty(&json)
}
}
// Helper functions for varint encoding/decoding
fn write_varint(buf: &mut [u8], mut v: i64) -> usize {
let mut n = 0;
while v >= 0x80 {
buf[n] = (v as u8) | 0x80;
v >>= 7;
n += 1;
}
buf[n] = v as u8;
n + 1
}
fn read_varint(buf: &[u8]) -> io::Result<(i64, usize)> {
let mut result = 0i64;
let mut shift = 0;
let mut n = 0;
while n < buf.len() {
let byte = buf[n];
n += 1;
result |= ((byte & 0x7F) as i64) << shift;
if byte < 0x80 {
return Ok((result, n));
}
shift += 7;
}
Err(io::Error::new(io::ErrorKind::UnexpectedEof, "unexpected EOF"))
}
// Helper functions for index header manipulation
#[allow(dead_code)]
pub fn remove_index_headers(b: &[u8]) -> Option<&[u8]> {
if b.len() < 4 + S2_INDEX_TRAILER.len() {
return None;
}
// Skip size
let b = &b[4..];
// Check trailer
if !b.starts_with(S2_INDEX_TRAILER) {
return None;
}
Some(&b[S2_INDEX_TRAILER.len()..])
}
#[allow(dead_code)]
pub fn restore_index_headers(in_data: &[u8]) -> Vec<u8> {
if in_data.is_empty() {
return Vec::new();
}
let mut b = Vec::with_capacity(4 + S2_INDEX_HEADER.len() + in_data.len() + S2_INDEX_TRAILER.len() + 4);
b.extend_from_slice(&[0x50, 0x2A, 0x4D, 0x18]);
b.extend_from_slice(S2_INDEX_HEADER);
b.extend_from_slice(in_data);
let total_size = (b.len() + 4 + S2_INDEX_TRAILER.len()) as u32;
b.extend_from_slice(&total_size.to_le_bytes());
b.extend_from_slice(S2_INDEX_TRAILER);
let chunk_len = b.len() - 4;
b[1] = chunk_len as u8;
b[2] = (chunk_len >> 8) as u8;
b[3] = (chunk_len >> 16) as u8;
b
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_index_new() {
let index = Index::new();
assert_eq!(index.total_uncompressed, -1);
assert_eq!(index.total_compressed, -1);
assert!(index.info.is_empty());
assert_eq!(index.est_block_uncomp, 0);
}
#[test]
fn test_index_add() -> io::Result<()> {
let mut index = Index::new();
// 测试添加第一个索引
index.add(100, 1000)?;
assert_eq!(index.info.len(), 1);
assert_eq!(index.info[0].compressed_offset, 100);
assert_eq!(index.info[0].uncompressed_offset, 1000);
// 测试添加相同未压缩偏移量的索引
index.add(200, 1000)?;
assert_eq!(index.info.len(), 1);
assert_eq!(index.info[0].compressed_offset, 200);
assert_eq!(index.info[0].uncompressed_offset, 1000);
// 测试添加新的索引(确保距离足够大)
index.add(300, 2000 + MIN_INDEX_DIST)?;
assert_eq!(index.info.len(), 2);
assert_eq!(index.info[1].compressed_offset, 300);
assert_eq!(index.info[1].uncompressed_offset, 2000 + MIN_INDEX_DIST);
Ok(())
}
#[test]
fn test_index_add_errors() {
let mut index = Index::new();
// 添加初始索引
index.add(100, 1000).unwrap();
// 测试添加更小的未压缩偏移量
let err = index.add(200, 500).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
// 测试添加更小的压缩偏移量
let err = index.add(50, 2000).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn test_index_find() -> io::Result<()> {
let mut index = Index::new();
index.total_uncompressed = 1000 + MIN_INDEX_DIST * 3;
index.total_compressed = 5000;
// 添加一些测试数据,确保索引间距满足 MIN_INDEX_DIST 要求
index.add(100, 1000)?;
index.add(300, 1000 + MIN_INDEX_DIST)?;
index.add(500, 1000 + MIN_INDEX_DIST * 2)?;
// 测试查找存在的偏移量
let (comp, uncomp) = index.find(1500)?;
assert_eq!(comp, 100);
assert_eq!(uncomp, 1000);
// 测试查找边界值
let (comp, uncomp) = index.find(1000 + MIN_INDEX_DIST)?;
assert_eq!(comp, 300);
assert_eq!(uncomp, 1000 + MIN_INDEX_DIST);
// 测试查找最后一个索引
let (comp, uncomp) = index.find(1000 + MIN_INDEX_DIST * 2)?;
assert_eq!(comp, 500);
assert_eq!(uncomp, 1000 + MIN_INDEX_DIST * 2);
Ok(())
}
#[test]
fn test_index_find_errors() {
let mut index = Index::new();
index.total_uncompressed = 10000;
index.total_compressed = 5000;
// 测试未初始化的索引
let uninit_index = Index::new();
let err = uninit_index.find(1000).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::Other);
// 测试超出范围的偏移量
let err = index.find(15000).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
// 测试负数偏移量
let err = match index.find(-1000) {
Ok(_) => panic!("should be error"),
Err(e) => e,
};
assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
}
#[test]
fn test_index_reduce() {
let mut index = Index::new();
index.est_block_uncomp = MIN_INDEX_DIST;
// 添加超过最大索引数量的条目,确保间距满足 MIN_INDEX_DIST 要求
for i in 0..MAX_INDEX_ENTRIES + 100 {
index.add(i as i64 * 100, i as i64 * MIN_INDEX_DIST).unwrap();
}
// 手动调用 reduce 方法
index.reduce();
// 验证索引数量是否被正确减少
assert!(index.info.len() <= MAX_INDEX_ENTRIES);
}
#[test]
fn test_index_json() -> io::Result<()> {
let mut index = Index::new();
// 添加一些测试数据
index.add(100, 1000)?;
index.add(300, 2000 + MIN_INDEX_DIST)?;
// 测试 JSON 序列化
let json = index.to_json().unwrap();
let json_str = String::from_utf8(json).unwrap();
println!("json_str: {}", json_str);
// 验证 JSON 内容
assert!(json_str.contains("\"compressed\": 100"));
assert!(json_str.contains("\"uncompressed\": 1000"));
assert!(json_str.contains("\"est_block_uncompressed\": 0"));
Ok(())
}
}
+510
View File
@@ -0,0 +1,510 @@
use crate::compress_index::{Index, TryGetIndex};
use crate::{EtagResolvable, HashReaderDetector};
use crate::{HashReaderMut, Reader};
use pin_project_lite::pin_project;
use rustfs_utils::compress::{CompressionAlgorithm, compress_block, decompress_block};
use rustfs_utils::{put_uvarint, uvarint};
use std::cmp::min;
use std::io::{self};
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
// use tracing::error;
const COMPRESS_TYPE_COMPRESSED: u8 = 0x00;
const COMPRESS_TYPE_UNCOMPRESSED: u8 = 0x01;
const COMPRESS_TYPE_END: u8 = 0xFF;
const DEFAULT_BLOCK_SIZE: usize = 1 << 20; // 1MB
const HEADER_LEN: usize = 8;
pin_project! {
#[derive(Debug)]
/// A reader wrapper that compresses data on the fly using DEFLATE algorithm.
pub struct CompressReader<R> {
#[pin]
pub inner: R,
buffer: Vec<u8>,
pos: usize,
done: bool,
block_size: usize,
compression_algorithm: CompressionAlgorithm,
index: Index,
written: usize,
uncomp_written: usize,
temp_buffer: Vec<u8>,
temp_pos: usize,
}
}
impl<R> CompressReader<R>
where
R: Reader,
{
pub fn new(inner: R, compression_algorithm: CompressionAlgorithm) -> Self {
Self {
inner,
buffer: Vec::new(),
pos: 0,
done: false,
compression_algorithm,
block_size: DEFAULT_BLOCK_SIZE,
index: Index::new(),
written: 0,
uncomp_written: 0,
temp_buffer: Vec::with_capacity(DEFAULT_BLOCK_SIZE), // Pre-allocate capacity
temp_pos: 0,
}
}
/// Optional: allow users to customize block_size
pub fn with_block_size(inner: R, block_size: usize, compression_algorithm: CompressionAlgorithm) -> Self {
Self {
inner,
buffer: Vec::new(),
pos: 0,
done: false,
compression_algorithm,
block_size,
index: Index::new(),
written: 0,
uncomp_written: 0,
temp_buffer: Vec::with_capacity(block_size),
temp_pos: 0,
}
}
}
impl<R> TryGetIndex for CompressReader<R>
where
R: Reader,
{
fn try_get_index(&self) -> Option<&Index> {
Some(&self.index)
}
}
impl<R> AsyncRead for CompressReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
let mut this = self.project();
// Copy from buffer first if available
if *this.pos < this.buffer.len() {
let to_copy = min(buf.remaining(), this.buffer.len() - *this.pos);
buf.put_slice(&this.buffer[*this.pos..*this.pos + to_copy]);
*this.pos += to_copy;
if *this.pos == this.buffer.len() {
this.buffer.clear();
*this.pos = 0;
}
return Poll::Ready(Ok(()));
}
if *this.done {
return Poll::Ready(Ok(()));
}
// Fill temporary buffer
while this.temp_buffer.len() < *this.block_size {
let remaining = *this.block_size - this.temp_buffer.len();
let mut temp = vec![0u8; remaining];
let mut temp_buf = ReadBuf::new(&mut temp);
match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
Poll::Pending => {
if this.temp_buffer.is_empty() {
return Poll::Pending;
}
break;
}
Poll::Ready(Ok(())) => {
let n = temp_buf.filled().len();
if n == 0 {
if this.temp_buffer.is_empty() {
return Poll::Ready(Ok(()));
}
break;
}
this.temp_buffer.extend_from_slice(&temp[..n]);
}
Poll::Ready(Err(e)) => {
// error!("CompressReader poll_read: read inner error: {e}");
return Poll::Ready(Err(e));
}
}
}
// Process accumulated data
if !this.temp_buffer.is_empty() {
let uncompressed_data = &this.temp_buffer;
let out = build_compressed_block(uncompressed_data, *this.compression_algorithm);
*this.written += out.len();
*this.uncomp_written += uncompressed_data.len();
if let Err(e) = this.index.add(*this.written as i64, *this.uncomp_written as i64) {
// error!("CompressReader index add error: {e}");
return Poll::Ready(Err(e));
}
*this.buffer = out;
*this.pos = 0;
this.temp_buffer.truncate(0); // More efficient way to clear
let to_copy = min(buf.remaining(), this.buffer.len());
buf.put_slice(&this.buffer[..to_copy]);
*this.pos += to_copy;
if *this.pos == this.buffer.len() {
this.buffer.clear();
*this.pos = 0;
}
Poll::Ready(Ok(()))
} else {
Poll::Pending
}
}
}
impl<R> EtagResolvable for CompressReader<R>
where
R: EtagResolvable,
{
fn try_resolve_etag(&mut self) -> Option<String> {
self.inner.try_resolve_etag()
}
}
impl<R> HashReaderDetector for CompressReader<R>
where
R: HashReaderDetector,
{
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
pin_project! {
/// A reader wrapper that decompresses data on the fly using DEFLATE algorithm.
/// Header format:
/// - First byte: compression type (00 = compressed, 01 = uncompressed, FF = end)
/// - Bytes 1-3: length of compressed data (little-endian)
/// - Bytes 4-7: CRC32 checksum of uncompressed data (little-endian)
#[derive(Debug)]
pub struct DecompressReader<R> {
#[pin]
pub inner: R,
buffer: Vec<u8>,
buffer_pos: usize,
finished: bool,
// Fields for saving header read progress across polls
header_buf: [u8; 8],
header_read: usize,
header_done: bool,
// Fields for saving compressed block read progress across polls
compressed_buf: Option<Vec<u8>>,
compressed_read: usize,
compressed_len: usize,
compression_algorithm: CompressionAlgorithm,
}
}
impl<R> DecompressReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
pub fn new(inner: R, compression_algorithm: CompressionAlgorithm) -> Self {
Self {
inner,
buffer: Vec::new(),
buffer_pos: 0,
finished: false,
header_buf: [0u8; 8],
header_read: 0,
header_done: false,
compressed_buf: None,
compressed_read: 0,
compressed_len: 0,
compression_algorithm,
}
}
}
impl<R> AsyncRead for DecompressReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
let mut this = self.project();
// Copy from buffer first if available
if *this.buffer_pos < this.buffer.len() {
let to_copy = min(buf.remaining(), this.buffer.len() - *this.buffer_pos);
buf.put_slice(&this.buffer[*this.buffer_pos..*this.buffer_pos + to_copy]);
*this.buffer_pos += to_copy;
if *this.buffer_pos == this.buffer.len() {
this.buffer.clear();
*this.buffer_pos = 0;
}
return Poll::Ready(Ok(()));
}
if *this.finished {
return Poll::Ready(Ok(()));
}
// Read header
while !*this.header_done && *this.header_read < HEADER_LEN {
let mut temp = [0u8; HEADER_LEN];
let mut temp_buf = ReadBuf::new(&mut temp[0..HEADER_LEN - *this.header_read]);
match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
Poll::Pending => return Poll::Pending,
Poll::Ready(Ok(())) => {
let n = temp_buf.filled().len();
if n == 0 {
break;
}
this.header_buf[*this.header_read..*this.header_read + n].copy_from_slice(&temp_buf.filled()[..n]);
*this.header_read += n;
}
Poll::Ready(Err(e)) => {
// error!("DecompressReader poll_read: read header error: {e}");
return Poll::Ready(Err(e));
}
}
if *this.header_read < HEADER_LEN {
return Poll::Pending;
}
}
if !*this.header_done && *this.header_read == 0 {
return Poll::Ready(Ok(()));
}
let typ = this.header_buf[0];
let len = (this.header_buf[1] as usize) | ((this.header_buf[2] as usize) << 8) | ((this.header_buf[3] as usize) << 16);
let crc = (this.header_buf[4] as u32)
| ((this.header_buf[5] as u32) << 8)
| ((this.header_buf[6] as u32) << 16)
| ((this.header_buf[7] as u32) << 24);
*this.header_read = 0;
*this.header_done = true;
if this.compressed_buf.is_none() {
*this.compressed_len = len;
*this.compressed_buf = Some(vec![0u8; *this.compressed_len]);
*this.compressed_read = 0;
}
let compressed_buf = this.compressed_buf.as_mut().unwrap();
while *this.compressed_read < *this.compressed_len {
let mut temp_buf = ReadBuf::new(&mut compressed_buf[*this.compressed_read..]);
match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
Poll::Pending => return Poll::Pending,
Poll::Ready(Ok(())) => {
let n = temp_buf.filled().len();
if n == 0 {
break;
}
*this.compressed_read += n;
}
Poll::Ready(Err(e)) => {
// error!("DecompressReader poll_read: read compressed block error: {e}");
this.compressed_buf.take();
*this.compressed_read = 0;
*this.compressed_len = 0;
return Poll::Ready(Err(e));
}
}
}
let (uncompress_len, uvarint) = uvarint(&compressed_buf[0..16]);
let compressed_data = &compressed_buf[uvarint as usize..];
let decompressed = if typ == COMPRESS_TYPE_COMPRESSED {
match decompress_block(compressed_data, *this.compression_algorithm) {
Ok(out) => out,
Err(e) => {
// error!("DecompressReader decompress_block error: {e}");
this.compressed_buf.take();
*this.compressed_read = 0;
*this.compressed_len = 0;
return Poll::Ready(Err(e));
}
}
} else if typ == COMPRESS_TYPE_UNCOMPRESSED {
compressed_data.to_vec()
} else if typ == COMPRESS_TYPE_END {
this.compressed_buf.take();
*this.compressed_read = 0;
*this.compressed_len = 0;
*this.finished = true;
return Poll::Ready(Ok(()));
} else {
// error!("DecompressReader unknown compression type: {typ}");
this.compressed_buf.take();
*this.compressed_read = 0;
*this.compressed_len = 0;
return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "Unknown compression type")));
};
if decompressed.len() != uncompress_len as usize {
// error!("DecompressReader decompressed length mismatch: {} != {}", decompressed.len(), uncompress_len);
this.compressed_buf.take();
*this.compressed_read = 0;
*this.compressed_len = 0;
return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "Decompressed length mismatch")));
}
let actual_crc = crc32fast::hash(&decompressed);
if actual_crc != crc {
// error!("DecompressReader CRC32 mismatch: actual {actual_crc} != expected {crc}");
this.compressed_buf.take();
*this.compressed_read = 0;
*this.compressed_len = 0;
return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "CRC32 mismatch")));
}
*this.buffer = decompressed;
*this.buffer_pos = 0;
this.compressed_buf.take();
*this.compressed_read = 0;
*this.compressed_len = 0;
*this.header_done = false;
let to_copy = min(buf.remaining(), this.buffer.len());
buf.put_slice(&this.buffer[..to_copy]);
*this.buffer_pos += to_copy;
if *this.buffer_pos == this.buffer.len() {
this.buffer.clear();
*this.buffer_pos = 0;
}
Poll::Ready(Ok(()))
}
}
impl<R> EtagResolvable for DecompressReader<R>
where
R: EtagResolvable,
{
fn try_resolve_etag(&mut self) -> Option<String> {
self.inner.try_resolve_etag()
}
}
impl<R> HashReaderDetector for DecompressReader<R>
where
R: HashReaderDetector,
{
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
/// Build compressed block with header + uvarint + compressed data
fn build_compressed_block(uncompressed_data: &[u8], compression_algorithm: CompressionAlgorithm) -> Vec<u8> {
let crc = crc32fast::hash(uncompressed_data);
let compressed_data = compress_block(uncompressed_data, compression_algorithm);
let uncompressed_len = uncompressed_data.len();
let mut uncompressed_len_buf = [0u8; 10];
let int_len = put_uvarint(&mut uncompressed_len_buf[..], uncompressed_len as u64);
let len = compressed_data.len() + int_len;
let mut header = [0u8; HEADER_LEN];
header[0] = COMPRESS_TYPE_COMPRESSED;
header[1] = (len & 0xFF) as u8;
header[2] = ((len >> 8) & 0xFF) as u8;
header[3] = ((len >> 16) & 0xFF) as u8;
header[4] = (crc & 0xFF) as u8;
header[5] = ((crc >> 8) & 0xFF) as u8;
header[6] = ((crc >> 16) & 0xFF) as u8;
header[7] = ((crc >> 24) & 0xFF) as u8;
let mut out = Vec::with_capacity(len + HEADER_LEN);
out.extend_from_slice(&header);
out.extend_from_slice(&uncompressed_len_buf[..int_len]);
out.extend_from_slice(&compressed_data);
out
}
#[cfg(test)]
mod tests {
use crate::WarpReader;
use super::*;
use std::io::Cursor;
use tokio::io::{AsyncReadExt, BufReader};
#[tokio::test]
async fn test_compress_reader_basic() {
let data = b"hello world, hello world, hello world!";
let reader = Cursor::new(&data[..]);
let mut compress_reader = CompressReader::new(WarpReader::new(reader), CompressionAlgorithm::Gzip);
let mut compressed = Vec::new();
compress_reader.read_to_end(&mut compressed).await.unwrap();
// DecompressReader解包
let mut decompress_reader = DecompressReader::new(Cursor::new(compressed.clone()), CompressionAlgorithm::Gzip);
let mut decompressed = Vec::new();
decompress_reader.read_to_end(&mut decompressed).await.unwrap();
assert_eq!(&decompressed, data);
}
#[tokio::test]
async fn test_compress_reader_basic_deflate() {
let data = b"hello world, hello world, hello world!";
let reader = BufReader::new(&data[..]);
let mut compress_reader = CompressReader::new(WarpReader::new(reader), CompressionAlgorithm::Deflate);
let mut compressed = Vec::new();
compress_reader.read_to_end(&mut compressed).await.unwrap();
// DecompressReader解包
let mut decompress_reader = DecompressReader::new(Cursor::new(compressed.clone()), CompressionAlgorithm::Deflate);
let mut decompressed = Vec::new();
decompress_reader.read_to_end(&mut decompressed).await.unwrap();
assert_eq!(&decompressed, data);
}
#[tokio::test]
async fn test_compress_reader_empty() {
let data = b"";
let reader = BufReader::new(&data[..]);
let mut compress_reader = CompressReader::new(WarpReader::new(reader), CompressionAlgorithm::Gzip);
let mut compressed = Vec::new();
compress_reader.read_to_end(&mut compressed).await.unwrap();
let mut decompress_reader = DecompressReader::new(Cursor::new(compressed.clone()), CompressionAlgorithm::Gzip);
let mut decompressed = Vec::new();
decompress_reader.read_to_end(&mut decompressed).await.unwrap();
assert_eq!(&decompressed, data);
}
#[tokio::test]
async fn test_compress_reader_large() {
use rand::Rng;
// Generate 1MB of random bytes
let mut data = vec![0u8; 1024 * 1024];
rand::rng().fill(&mut data[..]);
let reader = Cursor::new(data.clone());
let mut compress_reader = CompressReader::new(WarpReader::new(reader), CompressionAlgorithm::Gzip);
let mut compressed = Vec::new();
compress_reader.read_to_end(&mut compressed).await.unwrap();
let mut decompress_reader = DecompressReader::new(Cursor::new(compressed.clone()), CompressionAlgorithm::Gzip);
let mut decompressed = Vec::new();
decompress_reader.read_to_end(&mut decompressed).await.unwrap();
assert_eq!(&decompressed, &data);
}
#[tokio::test]
async fn test_compress_reader_large_deflate() {
use rand::Rng;
// Generate 1MB of random bytes
let mut data = vec![0u8; 1024 * 1024 * 3 + 512];
rand::rng().fill(&mut data[..]);
let reader = Cursor::new(data.clone());
let mut compress_reader = CompressReader::new(WarpReader::new(reader), CompressionAlgorithm::default());
let mut compressed = Vec::new();
compress_reader.read_to_end(&mut compressed).await.unwrap();
let mut decompress_reader = DecompressReader::new(Cursor::new(compressed.clone()), CompressionAlgorithm::default());
let mut decompressed = Vec::new();
decompress_reader.read_to_end(&mut decompressed).await.unwrap();
assert_eq!(&decompressed, &data);
}
}
+436
View File
@@ -0,0 +1,436 @@
use crate::HashReaderDetector;
use crate::HashReaderMut;
use crate::compress_index::{Index, TryGetIndex};
use crate::{EtagResolvable, Reader};
use aes_gcm::aead::Aead;
use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
use pin_project_lite::pin_project;
use rustfs_utils::{put_uvarint, put_uvarint_len};
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
pin_project! {
/// A reader wrapper that encrypts data on the fly using AES-256-GCM.
/// This is a demonstration. For production, use a secure and audited crypto library.
#[derive(Debug)]
pub struct EncryptReader<R> {
#[pin]
pub inner: R,
key: [u8; 32], // AES-256-GCM key
nonce: [u8; 12], // 96-bit nonce for GCM
buffer: Vec<u8>,
buffer_pos: usize,
finished: bool,
}
}
impl<R> EncryptReader<R>
where
R: Reader,
{
pub fn new(inner: R, key: [u8; 32], nonce: [u8; 12]) -> Self {
Self {
inner,
key,
nonce,
buffer: Vec::new(),
buffer_pos: 0,
finished: false,
}
}
}
impl<R> AsyncRead for EncryptReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
let mut this = self.project();
// Serve from buffer if any
if *this.buffer_pos < this.buffer.len() {
let to_copy = std::cmp::min(buf.remaining(), this.buffer.len() - *this.buffer_pos);
buf.put_slice(&this.buffer[*this.buffer_pos..*this.buffer_pos + to_copy]);
*this.buffer_pos += to_copy;
if *this.buffer_pos == this.buffer.len() {
this.buffer.clear();
*this.buffer_pos = 0;
}
return Poll::Ready(Ok(()));
}
if *this.finished {
return Poll::Ready(Ok(()));
}
// Read a fixed block size from inner
let block_size = 8 * 1024;
let mut temp = vec![0u8; block_size];
let mut temp_buf = ReadBuf::new(&mut temp);
match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
Poll::Pending => Poll::Pending,
Poll::Ready(Ok(())) => {
let n = temp_buf.filled().len();
if n == 0 {
// EOF, write end header
let mut header = [0u8; 8];
header[0] = 0xFF; // type: end
*this.buffer = header.to_vec();
*this.buffer_pos = 0;
*this.finished = true;
let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
buf.put_slice(&this.buffer[..to_copy]);
*this.buffer_pos += to_copy;
Poll::Ready(Ok(()))
} else {
// Encrypt the chunk
let cipher = Aes256Gcm::new_from_slice(this.key).expect("key");
let nonce = Nonce::from_slice(this.nonce);
let plaintext = &temp_buf.filled()[..n];
let plaintext_len = plaintext.len();
let crc = crc32fast::hash(plaintext);
let ciphertext = cipher
.encrypt(nonce, plaintext)
.map_err(|e| std::io::Error::other(format!("encrypt error: {e}")))?;
let int_len = put_uvarint_len(plaintext_len as u64);
let clen = int_len + ciphertext.len() + 4;
// Header: 8 bytes
// 0: type (0 = encrypted, 0xFF = end)
// 1-3: length (little endian u24, ciphertext length)
// 4-7: CRC32 of ciphertext (little endian u32)
let mut header = [0u8; 8];
header[0] = 0x00; // 0 = encrypted
header[1] = (clen & 0xFF) as u8;
header[2] = ((clen >> 8) & 0xFF) as u8;
header[3] = ((clen >> 16) & 0xFF) as u8;
header[4] = (crc & 0xFF) as u8;
header[5] = ((crc >> 8) & 0xFF) as u8;
header[6] = ((crc >> 16) & 0xFF) as u8;
header[7] = ((crc >> 24) & 0xFF) as u8;
let mut out = Vec::with_capacity(8 + int_len + ciphertext.len());
out.extend_from_slice(&header);
let mut plaintext_len_buf = vec![0u8; int_len];
put_uvarint(&mut plaintext_len_buf, plaintext_len as u64);
out.extend_from_slice(&plaintext_len_buf);
out.extend_from_slice(&ciphertext);
*this.buffer = out;
*this.buffer_pos = 0;
let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
buf.put_slice(&this.buffer[..to_copy]);
*this.buffer_pos += to_copy;
Poll::Ready(Ok(()))
}
}
Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
}
}
}
impl<R> EtagResolvable for EncryptReader<R>
where
R: EtagResolvable,
{
fn try_resolve_etag(&mut self) -> Option<String> {
self.inner.try_resolve_etag()
}
}
impl<R> HashReaderDetector for EncryptReader<R>
where
R: EtagResolvable + HashReaderDetector,
{
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
impl<R> TryGetIndex for EncryptReader<R>
where
R: TryGetIndex,
{
fn try_get_index(&self) -> Option<&Index> {
self.inner.try_get_index()
}
}
pin_project! {
/// A reader wrapper that decrypts data on the fly using AES-256-GCM.
/// This is a demonstration. For production, use a secure and audited crypto library.
#[derive(Debug)]
pub struct DecryptReader<R> {
#[pin]
pub inner: R,
key: [u8; 32], // AES-256-GCM key
nonce: [u8; 12], // 96-bit nonce for GCM
buffer: Vec<u8>,
buffer_pos: usize,
finished: bool,
// For block framing
header_buf: [u8; 8],
header_read: usize,
header_done: bool,
ciphertext_buf: Option<Vec<u8>>,
ciphertext_read: usize,
ciphertext_len: usize,
}
}
impl<R> DecryptReader<R>
where
R: Reader,
{
pub fn new(inner: R, key: [u8; 32], nonce: [u8; 12]) -> Self {
Self {
inner,
key,
nonce,
buffer: Vec::new(),
buffer_pos: 0,
finished: false,
header_buf: [0u8; 8],
header_read: 0,
header_done: false,
ciphertext_buf: None,
ciphertext_read: 0,
ciphertext_len: 0,
}
}
}
impl<R> AsyncRead for DecryptReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
let mut this = self.project();
// Serve from buffer if any
if *this.buffer_pos < this.buffer.len() {
let to_copy = std::cmp::min(buf.remaining(), this.buffer.len() - *this.buffer_pos);
buf.put_slice(&this.buffer[*this.buffer_pos..*this.buffer_pos + to_copy]);
*this.buffer_pos += to_copy;
if *this.buffer_pos == this.buffer.len() {
this.buffer.clear();
*this.buffer_pos = 0;
}
return Poll::Ready(Ok(()));
}
if *this.finished {
return Poll::Ready(Ok(()));
}
// Read header (8 bytes), support partial header read
while !*this.header_done && *this.header_read < 8 {
let mut temp = [0u8; 8];
let mut temp_buf = ReadBuf::new(&mut temp[0..8 - *this.header_read]);
match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
Poll::Pending => return Poll::Pending,
Poll::Ready(Ok(())) => {
let n = temp_buf.filled().len();
if n == 0 {
break;
}
this.header_buf[*this.header_read..*this.header_read + n].copy_from_slice(&temp_buf.filled()[..n]);
*this.header_read += n;
}
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
}
if *this.header_read < 8 {
return Poll::Pending;
}
}
if !*this.header_done && *this.header_read == 8 {
*this.header_done = true;
}
if !*this.header_done {
return Poll::Pending;
}
let typ = this.header_buf[0];
let len = (this.header_buf[1] as usize) | ((this.header_buf[2] as usize) << 8) | ((this.header_buf[3] as usize) << 16);
let crc = (this.header_buf[4] as u32)
| ((this.header_buf[5] as u32) << 8)
| ((this.header_buf[6] as u32) << 16)
| ((this.header_buf[7] as u32) << 24);
*this.header_read = 0;
*this.header_done = false;
if typ == 0xFF {
*this.finished = true;
return Poll::Ready(Ok(()));
}
// Read ciphertext block (len bytes), support partial read
if this.ciphertext_buf.is_none() {
*this.ciphertext_len = len - 4; // 4 bytes for CRC32
*this.ciphertext_buf = Some(vec![0u8; *this.ciphertext_len]);
*this.ciphertext_read = 0;
}
let ciphertext_buf = this.ciphertext_buf.as_mut().unwrap();
while *this.ciphertext_read < *this.ciphertext_len {
let mut temp_buf = ReadBuf::new(&mut ciphertext_buf[*this.ciphertext_read..]);
match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
Poll::Pending => return Poll::Pending,
Poll::Ready(Ok(())) => {
let n = temp_buf.filled().len();
if n == 0 {
break;
}
*this.ciphertext_read += n;
}
Poll::Ready(Err(e)) => {
this.ciphertext_buf.take();
*this.ciphertext_read = 0;
*this.ciphertext_len = 0;
return Poll::Ready(Err(e));
}
}
}
if *this.ciphertext_read < *this.ciphertext_len {
return Poll::Pending;
}
// Parse uvarint for plaintext length
let (plaintext_len, uvarint_len) = rustfs_utils::uvarint(&ciphertext_buf[0..16]);
let ciphertext = &ciphertext_buf[uvarint_len as usize..];
// Decrypt
let cipher = Aes256Gcm::new_from_slice(this.key).expect("key");
let nonce = Nonce::from_slice(this.nonce);
let plaintext = cipher
.decrypt(nonce, ciphertext)
.map_err(|e| std::io::Error::other(format!("decrypt error: {e}")))?;
if plaintext.len() != plaintext_len as usize {
this.ciphertext_buf.take();
*this.ciphertext_read = 0;
*this.ciphertext_len = 0;
return Poll::Ready(Err(std::io::Error::other("Plaintext length mismatch")));
}
// CRC32 check
let actual_crc = crc32fast::hash(&plaintext);
if actual_crc != crc {
this.ciphertext_buf.take();
*this.ciphertext_read = 0;
*this.ciphertext_len = 0;
return Poll::Ready(Err(std::io::Error::other("CRC32 mismatch")));
}
*this.buffer = plaintext;
*this.buffer_pos = 0;
// Clear block state for next block
this.ciphertext_buf.take();
*this.ciphertext_read = 0;
*this.ciphertext_len = 0;
let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
buf.put_slice(&this.buffer[..to_copy]);
*this.buffer_pos += to_copy;
Poll::Ready(Ok(()))
}
}
impl<R> EtagResolvable for DecryptReader<R>
where
R: EtagResolvable,
{
fn try_resolve_etag(&mut self) -> Option<String> {
self.inner.try_resolve_etag()
}
}
impl<R> HashReaderDetector for DecryptReader<R>
where
R: EtagResolvable + HashReaderDetector,
{
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
#[cfg(test)]
mod tests {
use std::io::Cursor;
use crate::WarpReader;
use super::*;
use rand::RngCore;
use tokio::io::{AsyncReadExt, BufReader};
#[tokio::test]
async fn test_encrypt_decrypt_reader_aes256gcm() {
let data = b"hello sse encrypt";
let mut key = [0u8; 32];
let mut nonce = [0u8; 12];
rand::rng().fill_bytes(&mut key);
rand::rng().fill_bytes(&mut nonce);
let reader = BufReader::new(&data[..]);
let encrypt_reader = EncryptReader::new(WarpReader::new(reader), key, nonce);
// Encrypt
let mut encrypt_reader = encrypt_reader;
let mut encrypted = Vec::new();
encrypt_reader.read_to_end(&mut encrypted).await.unwrap();
// Decrypt using DecryptReader
let reader = Cursor::new(encrypted.clone());
let decrypt_reader = DecryptReader::new(WarpReader::new(reader), key, nonce);
let mut decrypt_reader = decrypt_reader;
let mut decrypted = Vec::new();
decrypt_reader.read_to_end(&mut decrypted).await.unwrap();
assert_eq!(&decrypted, data);
}
#[tokio::test]
async fn test_decrypt_reader_only() {
// Encrypt some data first
let data = b"test decrypt only";
let mut key = [0u8; 32];
let mut nonce = [0u8; 12];
rand::rng().fill_bytes(&mut key);
rand::rng().fill_bytes(&mut nonce);
// Encrypt
let reader = BufReader::new(&data[..]);
let encrypt_reader = EncryptReader::new(WarpReader::new(reader), key, nonce);
let mut encrypt_reader = encrypt_reader;
let mut encrypted = Vec::new();
encrypt_reader.read_to_end(&mut encrypted).await.unwrap();
// Now test DecryptReader
let reader = Cursor::new(encrypted.clone());
let decrypt_reader = DecryptReader::new(WarpReader::new(reader), key, nonce);
let mut decrypt_reader = decrypt_reader;
let mut decrypted = Vec::new();
decrypt_reader.read_to_end(&mut decrypted).await.unwrap();
assert_eq!(&decrypted, data);
}
#[tokio::test]
async fn test_encrypt_decrypt_reader_large() {
use rand::Rng;
let size = 1024 * 1024;
let mut data = vec![0u8; size];
rand::rng().fill(&mut data[..]);
let mut key = [0u8; 32];
let mut nonce = [0u8; 12];
rand::rng().fill_bytes(&mut key);
rand::rng().fill_bytes(&mut nonce);
let reader = std::io::Cursor::new(data.clone());
let encrypt_reader = EncryptReader::new(WarpReader::new(reader), key, nonce);
let mut encrypt_reader = encrypt_reader;
let mut encrypted = Vec::new();
encrypt_reader.read_to_end(&mut encrypted).await.unwrap();
let reader = std::io::Cursor::new(encrypted.clone());
let decrypt_reader = DecryptReader::new(WarpReader::new(reader), key, nonce);
let mut decrypt_reader = decrypt_reader;
let mut decrypted = Vec::new();
decrypt_reader.read_to_end(&mut decrypted).await.unwrap();
assert_eq!(&decrypted, &data);
}
}
+248
View File
@@ -0,0 +1,248 @@
/*!
# AsyncRead Wrapper Types with ETag Support
This module demonstrates a pattern for handling wrapped AsyncRead types where:
- Reader types contain the actual ETag capability
- Wrapper types need to be recursively unwrapped
- The system can handle arbitrary nesting like `CompressReader<EncryptReader<EtagReader<R>>>`
## Key Components
### Trait-Based Approach
The `EtagResolvable` trait provides a clean way to handle recursive unwrapping:
- Reader types implement it by returning their ETag directly
- Wrapper types implement it by delegating to their inner type
## Usage Examples
```rust
use rustfs_rio::{CompressReader, EtagReader, resolve_etag_generic};
use rustfs_rio::WarpReader;
use rustfs_utils::compress::CompressionAlgorithm;
use tokio::io::BufReader;
use std::io::Cursor;
// Direct usage with trait-based approach
let data = b"test data";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
let etag_reader = EtagReader::new(reader, Some("test_etag".to_string()));
let mut reader = CompressReader::new(etag_reader, CompressionAlgorithm::Gzip);
let etag = resolve_etag_generic(&mut reader);
```
*/
#[cfg(test)]
mod tests {
use crate::{CompressReader, EncryptReader, EtagReader, HashReader};
use crate::{WarpReader, resolve_etag_generic};
use rustfs_utils::compress::CompressionAlgorithm;
use std::io::Cursor;
use tokio::io::BufReader;
#[test]
fn test_etag_reader_resolution() {
let data = b"test data";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
let mut etag_reader = EtagReader::new(reader, Some("test_etag".to_string()));
// Test direct ETag resolution
assert_eq!(resolve_etag_generic(&mut etag_reader), Some("test_etag".to_string()));
}
#[test]
fn test_hash_reader_resolution() {
let data = b"test data";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
let mut hash_reader =
HashReader::new(reader, data.len() as i64, data.len() as i64, Some("hash_etag".to_string()), false).unwrap();
// Test HashReader ETag resolution
assert_eq!(resolve_etag_generic(&mut hash_reader), Some("hash_etag".to_string()));
}
#[test]
fn test_compress_reader_delegation() {
let data = b"test data for compression";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
let etag_reader = EtagReader::new(reader, Some("compress_etag".to_string()));
let mut compress_reader = CompressReader::new(etag_reader, CompressionAlgorithm::Gzip);
// Test that CompressReader delegates to inner EtagReader
assert_eq!(resolve_etag_generic(&mut compress_reader), Some("compress_etag".to_string()));
}
#[test]
fn test_encrypt_reader_delegation() {
let data = b"test data for encryption";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
let etag_reader = EtagReader::new(reader, Some("encrypt_etag".to_string()));
let key = [0u8; 32];
let nonce = [0u8; 12];
let mut encrypt_reader = EncryptReader::new(etag_reader, key, nonce);
// Test that EncryptReader delegates to inner EtagReader
assert_eq!(resolve_etag_generic(&mut encrypt_reader), Some("encrypt_etag".to_string()));
}
#[test]
fn test_complex_nesting() {
let data = b"test data for complex nesting";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
// Create a complex nested structure: CompressReader<EncryptReader<EtagReader<BufReader<Cursor>>>>
let etag_reader = EtagReader::new(reader, Some("nested_etag".to_string()));
let key = [0u8; 32];
let nonce = [0u8; 12];
let encrypt_reader = EncryptReader::new(etag_reader, key, nonce);
let mut compress_reader = CompressReader::new(encrypt_reader, CompressionAlgorithm::Gzip);
// Test that nested structure can resolve ETag
assert_eq!(resolve_etag_generic(&mut compress_reader), Some("nested_etag".to_string()));
}
#[test]
fn test_hash_reader_in_nested_structure() {
let data = b"test data for hash reader nesting";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
// Create nested structure: CompressReader<HashReader<BufReader<Cursor>>>
let hash_reader =
HashReader::new(reader, data.len() as i64, data.len() as i64, Some("hash_nested_etag".to_string()), false).unwrap();
let mut compress_reader = CompressReader::new(hash_reader, CompressionAlgorithm::Deflate);
// Test that nested HashReader can be resolved
assert_eq!(resolve_etag_generic(&mut compress_reader), Some("hash_nested_etag".to_string()));
}
#[test]
fn test_comprehensive_etag_extraction() {
println!("🔍 Testing comprehensive ETag extraction with real reader types...");
// Test 1: Simple EtagReader
let data1 = b"simple test";
let reader1 = BufReader::new(Cursor::new(&data1[..]));
let reader1 = Box::new(WarpReader::new(reader1));
let mut etag_reader = EtagReader::new(reader1, Some("simple_etag".to_string()));
assert_eq!(resolve_etag_generic(&mut etag_reader), Some("simple_etag".to_string()));
// Test 2: HashReader with ETag
let data2 = b"hash test";
let reader2 = BufReader::new(Cursor::new(&data2[..]));
let reader2 = Box::new(WarpReader::new(reader2));
let mut hash_reader =
HashReader::new(reader2, data2.len() as i64, data2.len() as i64, Some("hash_etag".to_string()), false).unwrap();
assert_eq!(resolve_etag_generic(&mut hash_reader), Some("hash_etag".to_string()));
// Test 3: Single wrapper - CompressReader<EtagReader>
let data3 = b"compress test";
let reader3 = BufReader::new(Cursor::new(&data3[..]));
let reader3 = Box::new(WarpReader::new(reader3));
let etag_reader3 = EtagReader::new(reader3, Some("compress_wrapped_etag".to_string()));
let mut compress_reader = CompressReader::new(etag_reader3, CompressionAlgorithm::Zstd);
assert_eq!(resolve_etag_generic(&mut compress_reader), Some("compress_wrapped_etag".to_string()));
// Test 4: Double wrapper - CompressReader<EncryptReader<EtagReader>>
let data4 = b"double wrap test";
let reader4 = BufReader::new(Cursor::new(&data4[..]));
let reader4 = Box::new(WarpReader::new(reader4));
let etag_reader4 = EtagReader::new(reader4, Some("double_wrapped_etag".to_string()));
let key = [1u8; 32];
let nonce = [1u8; 12];
let encrypt_reader4 = EncryptReader::new(etag_reader4, key, nonce);
let mut compress_reader4 = CompressReader::new(encrypt_reader4, CompressionAlgorithm::Gzip);
assert_eq!(resolve_etag_generic(&mut compress_reader4), Some("double_wrapped_etag".to_string()));
println!("✅ All ETag extraction methods work correctly!");
println!("✅ Trait-based approach handles recursive unwrapping!");
println!("✅ Complex nesting patterns with real reader types are supported!");
}
#[test]
fn test_real_world_scenario() {
println!("🔍 Testing real-world ETag extraction scenario with actual reader types...");
// Simulate a real-world scenario where we have nested AsyncRead wrappers
// and need to extract ETag information from deeply nested structures
let data = b"Real world test data that might be compressed and encrypted";
let base_reader = BufReader::new(Cursor::new(&data[..]));
let base_reader = Box::new(WarpReader::new(base_reader));
// Create a complex nested structure that might occur in practice:
// CompressReader<EncryptReader<HashReader<BufReader<Cursor>>>>
let hash_reader = HashReader::new(
base_reader,
data.len() as i64,
data.len() as i64,
Some("real_world_etag".to_string()),
false,
)
.unwrap();
let key = [42u8; 32];
let nonce = [24u8; 12];
let encrypt_reader = EncryptReader::new(hash_reader, key, nonce);
let mut compress_reader = CompressReader::new(encrypt_reader, CompressionAlgorithm::Deflate);
// Extract ETag using our generic system
let extracted_etag = resolve_etag_generic(&mut compress_reader);
println!("📋 Extracted ETag: {:?}", extracted_etag);
assert_eq!(extracted_etag, Some("real_world_etag".to_string()));
// Test another complex nesting with EtagReader at the core
let data2 = b"Another real world scenario";
let base_reader2 = BufReader::new(Cursor::new(&data2[..]));
let base_reader2 = Box::new(WarpReader::new(base_reader2));
let etag_reader = EtagReader::new(base_reader2, Some("core_etag".to_string()));
let key2 = [99u8; 32];
let nonce2 = [88u8; 12];
let encrypt_reader2 = EncryptReader::new(etag_reader, key2, nonce2);
let mut compress_reader2 = CompressReader::new(encrypt_reader2, CompressionAlgorithm::Zstd);
let trait_etag = resolve_etag_generic(&mut compress_reader2);
println!("📋 Trait-based ETag: {:?}", trait_etag);
assert_eq!(trait_etag, Some("core_etag".to_string()));
println!("✅ Real-world scenario test passed!");
println!(" - Successfully extracted ETag from nested CompressReader<EncryptReader<HashReader<AsyncRead>>>");
println!(" - Successfully extracted ETag from nested CompressReader<EncryptReader<EtagReader<AsyncRead>>>");
println!(" - Trait-based approach works with real reader types");
println!(" - System handles arbitrary nesting depths with actual implementations");
}
#[test]
fn test_no_etag_scenarios() {
println!("🔍 Testing scenarios where no ETag is available...");
// Test with HashReader that has no etag
let data = b"no etag test";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
let mut hash_reader_no_etag = HashReader::new(reader, data.len() as i64, data.len() as i64, None, false).unwrap();
assert_eq!(resolve_etag_generic(&mut hash_reader_no_etag), None);
// Test with EtagReader that has None etag
let data2 = b"no etag test 2";
let reader2 = BufReader::new(Cursor::new(&data2[..]));
let reader2 = Box::new(WarpReader::new(reader2));
let mut etag_reader_none = EtagReader::new(reader2, None);
assert_eq!(resolve_etag_generic(&mut etag_reader_none), None);
// Test nested structure with no ETag at the core
let data3 = b"nested no etag test";
let reader3 = BufReader::new(Cursor::new(&data3[..]));
let reader3 = Box::new(WarpReader::new(reader3));
let etag_reader3 = EtagReader::new(reader3, None);
let mut compress_reader3 = CompressReader::new(etag_reader3, CompressionAlgorithm::Gzip);
assert_eq!(resolve_etag_generic(&mut compress_reader3), None);
println!("✅ No ETag scenarios handled correctly!");
}
}
+229
View File
@@ -0,0 +1,229 @@
use crate::compress_index::{Index, TryGetIndex};
use crate::{EtagResolvable, HashReaderDetector, HashReaderMut, Reader};
use md5::{Digest, Md5};
use pin_project_lite::pin_project;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
pin_project! {
pub struct EtagReader {
#[pin]
pub inner: Box<dyn Reader>,
pub md5: Md5,
pub finished: bool,
pub checksum: Option<String>,
}
}
impl EtagReader {
pub fn new(inner: Box<dyn Reader>, checksum: Option<String>) -> Self {
Self {
inner,
md5: Md5::new(),
finished: false,
checksum,
}
}
/// Get the final md5 value (etag) as a hex string, only compute once.
/// Can be called multiple times, always returns the same result after finished.
pub fn get_etag(&mut self) -> String {
format!("{:x}", self.md5.clone().finalize())
}
}
impl AsyncRead for EtagReader {
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
let mut this = self.project();
let orig_filled = buf.filled().len();
let poll = this.inner.as_mut().poll_read(cx, buf);
if let Poll::Ready(Ok(())) = &poll {
let filled = &buf.filled()[orig_filled..];
if !filled.is_empty() {
this.md5.update(filled);
} else {
// EOF
*this.finished = true;
if let Some(checksum) = this.checksum {
let etag = format!("{:x}", this.md5.clone().finalize());
if *checksum != etag {
return Poll::Ready(Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "Checksum mismatch")));
}
}
}
}
poll
}
}
impl EtagResolvable for EtagReader {
fn is_etag_reader(&self) -> bool {
true
}
fn try_resolve_etag(&mut self) -> Option<String> {
// EtagReader provides its own etag, not delegating to inner
if let Some(checksum) = &self.checksum {
Some(checksum.clone())
} else if self.finished {
Some(self.get_etag())
} else {
None
}
}
}
impl HashReaderDetector for EtagReader {
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
impl TryGetIndex for EtagReader {
fn try_get_index(&self) -> Option<&Index> {
self.inner.try_get_index()
}
}
#[cfg(test)]
mod tests {
use crate::WarpReader;
use super::*;
use std::io::Cursor;
use tokio::io::{AsyncReadExt, BufReader};
#[tokio::test]
async fn test_etag_reader_basic() {
let data = b"hello world";
let mut hasher = Md5::new();
hasher.update(data);
let expected = format!("{:x}", hasher.finalize());
let reader = BufReader::new(&data[..]);
let reader = Box::new(WarpReader::new(reader));
let mut etag_reader = EtagReader::new(reader, None);
let mut buf = Vec::new();
let n = etag_reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, data.len());
assert_eq!(&buf, data);
let etag = etag_reader.try_resolve_etag();
assert_eq!(etag, Some(expected));
}
#[tokio::test]
async fn test_etag_reader_empty() {
let data = b"";
let mut hasher = Md5::new();
hasher.update(data);
let expected = format!("{:x}", hasher.finalize());
let reader = BufReader::new(&data[..]);
let reader = Box::new(WarpReader::new(reader));
let mut etag_reader = EtagReader::new(reader, None);
let mut buf = Vec::new();
let n = etag_reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, 0);
assert!(buf.is_empty());
let etag = etag_reader.try_resolve_etag();
assert_eq!(etag, Some(expected));
}
#[tokio::test]
async fn test_etag_reader_multiple_get() {
let data = b"abc123";
let mut hasher = Md5::new();
hasher.update(data);
let expected = format!("{:x}", hasher.finalize());
let reader = BufReader::new(&data[..]);
let reader = Box::new(WarpReader::new(reader));
let mut etag_reader = EtagReader::new(reader, None);
let mut buf = Vec::new();
let _ = etag_reader.read_to_end(&mut buf).await.unwrap();
// Call etag multiple times, should always return the same result
let etag1 = { etag_reader.try_resolve_etag() };
let etag2 = { etag_reader.try_resolve_etag() };
assert_eq!(etag1, Some(expected.clone()));
assert_eq!(etag2, Some(expected.clone()));
}
#[tokio::test]
async fn test_etag_reader_not_finished() {
let data = b"abc123";
let reader = BufReader::new(&data[..]);
let reader = Box::new(WarpReader::new(reader));
let mut etag_reader = EtagReader::new(reader, None);
// Do not read to end, etag should be None
let mut buf = [0u8; 2];
let _ = etag_reader.read(&mut buf).await.unwrap();
assert_eq!(etag_reader.try_resolve_etag(), None);
}
#[tokio::test]
async fn test_etag_reader_large_data() {
use rand::Rng;
// Generate 3MB random data
let size = 3 * 1024 * 1024;
let mut data = vec![0u8; size];
rand::rng().fill(&mut data[..]);
let mut hasher = Md5::new();
hasher.update(&data);
let cloned_data = data.clone();
let expected = format!("{:x}", hasher.finalize());
let reader = Cursor::new(data.clone());
let reader = Box::new(WarpReader::new(reader));
let mut etag_reader = EtagReader::new(reader, None);
let mut buf = Vec::new();
let n = etag_reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, size);
assert_eq!(&buf, &cloned_data);
let etag = etag_reader.try_resolve_etag();
assert_eq!(etag, Some(expected));
}
#[tokio::test]
async fn test_etag_reader_checksum_match() {
let data = b"checksum test data";
let mut hasher = Md5::new();
hasher.update(data);
let expected = format!("{:x}", hasher.finalize());
let reader = BufReader::new(&data[..]);
let reader = Box::new(WarpReader::new(reader));
let mut etag_reader = EtagReader::new(reader, Some(expected.clone()));
let mut buf = Vec::new();
let n = etag_reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, data.len());
assert_eq!(&buf, data);
// 校验通过,etag应等于expected
assert_eq!(etag_reader.try_resolve_etag(), Some(expected));
}
#[tokio::test]
async fn test_etag_reader_checksum_mismatch() {
let data = b"checksum test data";
let wrong_checksum = "deadbeefdeadbeefdeadbeefdeadbeef".to_string();
let reader = BufReader::new(&data[..]);
let reader = Box::new(WarpReader::new(reader));
let mut etag_reader = EtagReader::new(reader, Some(wrong_checksum));
let mut buf = Vec::new();
// 校验失败,应该返回InvalidData错误
let err = etag_reader.read_to_end(&mut buf).await.unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
}
}
+141
View File
@@ -0,0 +1,141 @@
use crate::compress_index::{Index, TryGetIndex};
use crate::{EtagResolvable, HashReaderDetector, HashReaderMut, Reader};
use pin_project_lite::pin_project;
use std::io::{Error, Result};
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
pin_project! {
pub struct HardLimitReader {
#[pin]
pub inner: Box<dyn Reader>,
remaining: i64,
}
}
impl HardLimitReader {
pub fn new(inner: Box<dyn Reader>, limit: i64) -> Self {
HardLimitReader { inner, remaining: limit }
}
}
impl AsyncRead for HardLimitReader {
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<Result<()>> {
if self.remaining < 0 {
return Poll::Ready(Err(Error::other("input provided more bytes than specified")));
}
// Save the initial length
let before = buf.filled().len();
// Poll the inner reader
let this = self.as_mut().project();
let poll = this.inner.poll_read(cx, buf);
if let Poll::Ready(Ok(())) = &poll {
let after = buf.filled().len();
let read = (after - before) as i64;
self.remaining -= read;
if self.remaining < 0 {
return Poll::Ready(Err(Error::other("input provided more bytes than specified")));
}
}
poll
}
}
impl EtagResolvable for HardLimitReader {
fn try_resolve_etag(&mut self) -> Option<String> {
self.inner.try_resolve_etag()
}
}
impl HashReaderDetector for HardLimitReader {
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
impl TryGetIndex for HardLimitReader {
fn try_get_index(&self) -> Option<&Index> {
self.inner.try_get_index()
}
}
#[cfg(test)]
mod tests {
use std::vec;
use crate::WarpReader;
use super::*;
use rustfs_utils::read_full;
use tokio::io::{AsyncReadExt, BufReader};
#[tokio::test]
async fn test_hardlimit_reader_normal() {
let data = b"hello world";
let reader = BufReader::new(&data[..]);
let reader = Box::new(WarpReader::new(reader));
let hardlimit = HardLimitReader::new(reader, 20);
let mut r = hardlimit;
let mut buf = Vec::new();
let n = r.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, data.len());
assert_eq!(&buf, data);
}
#[tokio::test]
async fn test_hardlimit_reader_exact_limit() {
let data = b"1234567890";
let reader = BufReader::new(&data[..]);
let reader = Box::new(WarpReader::new(reader));
let hardlimit = HardLimitReader::new(reader, 10);
let mut r = hardlimit;
let mut buf = Vec::new();
let n = r.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, 10);
assert_eq!(&buf, data);
}
#[tokio::test]
async fn test_hardlimit_reader_exceed_limit() {
let data = b"abcdef";
let reader = BufReader::new(&data[..]);
let reader = Box::new(WarpReader::new(reader));
let hardlimit = HardLimitReader::new(reader, 3);
let mut r = hardlimit;
let mut buf = vec![0u8; 10];
// 读取超限,应该返回错误
let err = match read_full(&mut r, &mut buf).await {
Ok(n) => {
println!("Read {} bytes", n);
assert_eq!(n, 3);
assert_eq!(&buf[..n], b"abc");
None
}
Err(e) => Some(e),
};
assert!(err.is_some());
let err = err.unwrap();
assert_eq!(err.kind(), std::io::ErrorKind::Other);
}
#[tokio::test]
async fn test_hardlimit_reader_empty() {
let data = b"";
let reader = BufReader::new(&data[..]);
let reader = Box::new(WarpReader::new(reader));
let hardlimit = HardLimitReader::new(reader, 5);
let mut r = hardlimit;
let mut buf = Vec::new();
let n = r.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, 0);
assert_eq!(&buf, data);
}
}
+582
View File
@@ -0,0 +1,582 @@
//! HashReader implementation with generic support
//!
//! This module provides a generic `HashReader<R>` that can wrap any type implementing
//! `AsyncRead + Unpin + Send + Sync + 'static + EtagResolvable`.
//!
//! ## Migration from the original Reader enum
//!
//! The original `HashReader::new` method that worked with the `Reader` enum
//! has been replaced with a generic approach. To preserve the original logic:
//!
//! ### Original logic (before generics):
//! ```ignore
//! // Original code would do:
//! // 1. Check if inner is already a HashReader
//! // 2. If size > 0, wrap with HardLimitReader
//! // 3. If !diskable_md5, wrap with EtagReader
//! // 4. Create HashReader with the wrapped reader
//!
//! let reader = HashReader::new(inner, size, actual_size, etag, diskable_md5)?;
//! ```
//!
//! ### New generic approach:
//! ```rust
//! use rustfs_rio::{HashReader, HardLimitReader, EtagReader};
//! use tokio::io::BufReader;
//! use std::io::Cursor;
//! use rustfs_rio::WarpReader;
//!
//! # tokio_test::block_on(async {
//! let data = b"hello world";
//! let reader = BufReader::new(Cursor::new(&data[..]));
//! let reader = Box::new(WarpReader::new(reader));
//! let size = data.len() as i64;
//! let actual_size = size;
//! let etag = None;
//! let diskable_md5 = false;
//!
//! // Method 1: Simple creation (recommended for most cases)
//! let hash_reader = HashReader::new(reader, size, actual_size, etag.clone(), diskable_md5).unwrap();
//!
//! // Method 2: With manual wrapping to recreate original logic
//! let reader2 = BufReader::new(Cursor::new(&data[..]));
//! let reader2 = Box::new(WarpReader::new(reader2));
//! let wrapped_reader: Box<dyn rustfs_rio::Reader> = if size > 0 {
//! if !diskable_md5 {
//! // Wrap with both HardLimitReader and EtagReader
//! let hard_limit = HardLimitReader::new(reader2, size);
//! Box::new(EtagReader::new(Box::new(hard_limit), etag.clone()))
//! } else {
//! // Only wrap with HardLimitReader
//! Box::new(HardLimitReader::new(reader2, size))
//! }
//! } else if !diskable_md5 {
//! // Only wrap with EtagReader
//! Box::new(EtagReader::new(reader2, etag.clone()))
//! } else {
//! // No wrapping needed
//! reader2
//! };
//! let hash_reader2 = HashReader::new(wrapped_reader, size, actual_size, etag, diskable_md5).unwrap();
//! # });
//! ```
//!
//! ## HashReader Detection
//!
//! The `HashReaderDetector` trait allows detection of existing HashReader instances:
//!
//! ```rust
//! use rustfs_rio::{HashReader, HashReaderDetector};
//! use tokio::io::BufReader;
//! use std::io::Cursor;
//! use rustfs_rio::WarpReader;
//!
//! # tokio_test::block_on(async {
//! let data = b"test";
//! let reader = BufReader::new(Cursor::new(&data[..]));
//! let hash_reader = HashReader::new(Box::new(WarpReader::new(reader)), 4, 4, None, false).unwrap();
//!
//! // Check if a type is a HashReader
//! assert!(hash_reader.is_hash_reader());
//!
//! // Use new for compatibility (though it's simpler to use new() directly)
//! let reader2 = BufReader::new(Cursor::new(&data[..]));
//! let result = HashReader::new(Box::new(WarpReader::new(reader2)), 4, 4, None, false);
//! assert!(result.is_ok());
//! # });
//! ```
use pin_project_lite::pin_project;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
use crate::compress_index::{Index, TryGetIndex};
use crate::{EtagReader, EtagResolvable, HardLimitReader, HashReaderDetector, Reader};
/// Trait for mutable operations on HashReader
pub trait HashReaderMut {
fn bytes_read(&self) -> u64;
fn checksum(&self) -> &Option<String>;
fn set_checksum(&mut self, checksum: Option<String>);
fn size(&self) -> i64;
fn set_size(&mut self, size: i64);
fn actual_size(&self) -> i64;
fn set_actual_size(&mut self, actual_size: i64);
}
pin_project! {
pub struct HashReader {
#[pin]
pub inner: Box<dyn Reader>,
pub size: i64,
checksum: Option<String>,
pub actual_size: i64,
pub diskable_md5: bool,
bytes_read: u64,
// TODO: content_hash
}
}
impl HashReader {
pub fn new(
mut inner: Box<dyn Reader>,
size: i64,
actual_size: i64,
md5: Option<String>,
diskable_md5: bool,
) -> std::io::Result<Self> {
// Check if it's already a HashReader and update its parameters
if let Some(existing_hash_reader) = inner.as_hash_reader_mut() {
if existing_hash_reader.bytes_read() > 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"Cannot create HashReader from an already read HashReader",
));
}
if let Some(checksum) = existing_hash_reader.checksum() {
if let Some(ref md5) = md5 {
if checksum != md5 {
return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "HashReader checksum mismatch"));
}
}
}
if existing_hash_reader.size() > 0 && size > 0 && existing_hash_reader.size() != size {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("HashReader size mismatch: expected {}, got {}", existing_hash_reader.size(), size),
));
}
existing_hash_reader.set_checksum(md5.clone());
if existing_hash_reader.size() < 0 && size >= 0 {
existing_hash_reader.set_size(size);
}
if existing_hash_reader.actual_size() <= 0 && actual_size >= 0 {
existing_hash_reader.set_actual_size(actual_size);
}
return Ok(Self {
inner,
size,
checksum: md5,
actual_size,
diskable_md5,
bytes_read: 0,
});
}
if size > 0 {
let hr = HardLimitReader::new(inner, size);
inner = Box::new(hr);
if !diskable_md5 && !inner.is_hash_reader() {
let er = EtagReader::new(inner, md5.clone());
inner = Box::new(er);
}
} else if !diskable_md5 {
let er = EtagReader::new(inner, md5.clone());
inner = Box::new(er);
}
Ok(Self {
inner,
size,
checksum: md5,
actual_size,
diskable_md5,
bytes_read: 0,
})
}
/// Update HashReader parameters
pub fn update_params(&mut self, size: i64, actual_size: i64, etag: Option<String>) {
if self.size < 0 && size >= 0 {
self.size = size;
}
if self.actual_size <= 0 && actual_size > 0 {
self.actual_size = actual_size;
}
if etag.is_some() {
self.checksum = etag;
}
}
pub fn size(&self) -> i64 {
self.size
}
pub fn actual_size(&self) -> i64 {
self.actual_size
}
}
impl HashReaderMut for HashReader {
fn bytes_read(&self) -> u64 {
self.bytes_read
}
fn checksum(&self) -> &Option<String> {
&self.checksum
}
fn set_checksum(&mut self, checksum: Option<String>) {
self.checksum = checksum;
}
fn size(&self) -> i64 {
self.size
}
fn set_size(&mut self, size: i64) {
self.size = size;
}
fn actual_size(&self) -> i64 {
self.actual_size
}
fn set_actual_size(&mut self, actual_size: i64) {
self.actual_size = actual_size;
}
}
impl AsyncRead for HashReader {
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
let this = self.project();
let poll = this.inner.poll_read(cx, buf);
if let Poll::Ready(Ok(())) = &poll {
let filled = buf.filled().len();
*this.bytes_read += filled as u64;
if filled == 0 {
// EOF
// TODO: check content_hash
}
}
poll
}
}
impl EtagResolvable for HashReader {
fn try_resolve_etag(&mut self) -> Option<String> {
if self.diskable_md5 {
return None;
}
if let Some(etag) = self.inner.try_resolve_etag() {
return Some(etag);
}
// If no etag from inner and we have a stored checksum, return it
self.checksum.clone()
}
}
impl HashReaderDetector for HashReader {
fn is_hash_reader(&self) -> bool {
true
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
Some(self)
}
}
impl TryGetIndex for HashReader {
fn try_get_index(&self) -> Option<&Index> {
self.inner.try_get_index()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{DecryptReader, WarpReader, encrypt_reader};
use std::io::Cursor;
use tokio::io::{AsyncReadExt, BufReader};
#[tokio::test]
async fn test_hashreader_wrapping_logic() {
let data = b"hello world";
let size = data.len() as i64;
let actual_size = size;
let etag = None;
// Test 1: Simple creation
let reader1 = BufReader::new(Cursor::new(&data[..]));
let reader1 = Box::new(WarpReader::new(reader1));
let hash_reader1 = HashReader::new(reader1, size, actual_size, etag.clone(), false).unwrap();
assert_eq!(hash_reader1.size(), size);
assert_eq!(hash_reader1.actual_size(), actual_size);
// Test 2: With HardLimitReader wrapping
let reader2 = BufReader::new(Cursor::new(&data[..]));
let reader2 = Box::new(WarpReader::new(reader2));
let hard_limit = HardLimitReader::new(reader2, size);
let hard_limit = Box::new(hard_limit);
let hash_reader2 = HashReader::new(hard_limit, size, actual_size, etag.clone(), false).unwrap();
assert_eq!(hash_reader2.size(), size);
assert_eq!(hash_reader2.actual_size(), actual_size);
// Test 3: With EtagReader wrapping
let reader3 = BufReader::new(Cursor::new(&data[..]));
let reader3 = Box::new(WarpReader::new(reader3));
let etag_reader = EtagReader::new(reader3, etag.clone());
let etag_reader = Box::new(etag_reader);
let hash_reader3 = HashReader::new(etag_reader, size, actual_size, etag.clone(), false).unwrap();
assert_eq!(hash_reader3.size(), size);
assert_eq!(hash_reader3.actual_size(), actual_size);
}
#[tokio::test]
async fn test_hashreader_etag_basic() {
let data = b"hello hashreader";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
let mut hash_reader = HashReader::new(reader, data.len() as i64, data.len() as i64, None, false).unwrap();
let mut buf = Vec::new();
let _ = hash_reader.read_to_end(&mut buf).await.unwrap();
// Since we removed EtagReader integration, etag might be None
let _etag = hash_reader.try_resolve_etag();
// Just check that we can call etag() without error
assert_eq!(buf, data);
}
#[tokio::test]
async fn test_hashreader_diskable_md5() {
let data = b"no etag";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
let mut hash_reader = HashReader::new(reader, data.len() as i64, data.len() as i64, None, true).unwrap();
let mut buf = Vec::new();
let _ = hash_reader.read_to_end(&mut buf).await.unwrap();
// Etag should be None when diskable_md5 is true
let etag = hash_reader.try_resolve_etag();
assert!(etag.is_none());
assert_eq!(buf, data);
}
#[tokio::test]
async fn test_hashreader_new_logic() {
let data = b"test data";
let reader = BufReader::new(Cursor::new(&data[..]));
let reader = Box::new(WarpReader::new(reader));
// Create a HashReader first
let hash_reader =
HashReader::new(reader, data.len() as i64, data.len() as i64, Some("test_etag".to_string()), false).unwrap();
let hash_reader = Box::new(WarpReader::new(hash_reader));
// Now try to create another HashReader from the existing one using new
let result = HashReader::new(hash_reader, data.len() as i64, data.len() as i64, Some("test_etag".to_string()), false);
assert!(result.is_ok());
let final_reader = result.unwrap();
assert_eq!(final_reader.checksum, Some("test_etag".to_string()));
assert_eq!(final_reader.size(), data.len() as i64);
}
#[tokio::test]
async fn test_for_wrapping_readers() {
use crate::{CompressReader, DecompressReader};
use md5::{Digest, Md5};
use rand::Rng;
use rand::RngCore;
use rustfs_utils::compress::CompressionAlgorithm;
// Generate 1MB random data
let size = 1024 * 1024;
let mut data = vec![0u8; size];
rand::rng().fill(&mut data[..]);
let mut hasher = Md5::new();
hasher.update(&data);
let expected = format!("{:x}", hasher.finalize());
println!("expected: {}", expected);
let reader = Cursor::new(data.clone());
let reader = BufReader::new(reader);
// 启用压缩测试
let is_compress = true;
let size = data.len() as i64;
let actual_size = data.len() as i64;
let reader = Box::new(WarpReader::new(reader));
// 创建 HashReader
let mut hr = HashReader::new(reader, size, actual_size, Some(expected.clone()), false).unwrap();
// 如果启用压缩,先压缩数据
let compressed_data = if is_compress {
let mut compressed_buf = Vec::new();
let compress_reader = CompressReader::new(hr, CompressionAlgorithm::Gzip);
let mut compress_reader = compress_reader;
compress_reader.read_to_end(&mut compressed_buf).await.unwrap();
println!("Original size: {}, Compressed size: {}", data.len(), compressed_buf.len());
compressed_buf
} else {
// 如果不压缩,直接读取原始数据
let mut buf = Vec::new();
hr.read_to_end(&mut buf).await.unwrap();
buf
};
let mut key = [0u8; 32];
let mut nonce = [0u8; 12];
rand::rng().fill_bytes(&mut key);
rand::rng().fill_bytes(&mut nonce);
let is_encrypt = true;
if is_encrypt {
// 加密压缩后的数据
let encrypt_reader = encrypt_reader::EncryptReader::new(WarpReader::new(Cursor::new(compressed_data)), key, nonce);
let mut encrypted_data = Vec::new();
let mut encrypt_reader = encrypt_reader;
encrypt_reader.read_to_end(&mut encrypted_data).await.unwrap();
println!("Encrypted size: {}", encrypted_data.len());
// 解密数据
let decrypt_reader = DecryptReader::new(WarpReader::new(Cursor::new(encrypted_data)), key, nonce);
let mut decrypt_reader = decrypt_reader;
let mut decrypted_data = Vec::new();
decrypt_reader.read_to_end(&mut decrypted_data).await.unwrap();
if is_compress {
// 如果使用了压缩,需要解压缩
let decompress_reader =
DecompressReader::new(WarpReader::new(Cursor::new(decrypted_data)), CompressionAlgorithm::Gzip);
let mut decompress_reader = decompress_reader;
let mut final_data = Vec::new();
decompress_reader.read_to_end(&mut final_data).await.unwrap();
println!("Final decompressed size: {}", final_data.len());
assert_eq!(final_data.len() as i64, actual_size);
assert_eq!(&final_data, &data);
} else {
// 如果没有压缩,直接比较解密后的数据
assert_eq!(decrypted_data.len() as i64, actual_size);
assert_eq!(&decrypted_data, &data);
}
return;
}
// 如果不加密,直接处理压缩/解压缩
if is_compress {
let decompress_reader =
DecompressReader::new(WarpReader::new(Cursor::new(compressed_data)), CompressionAlgorithm::Gzip);
let mut decompress_reader = decompress_reader;
let mut decompressed = Vec::new();
decompress_reader.read_to_end(&mut decompressed).await.unwrap();
assert_eq!(decompressed.len() as i64, actual_size);
assert_eq!(&decompressed, &data);
} else {
assert_eq!(compressed_data.len() as i64, actual_size);
assert_eq!(&compressed_data, &data);
}
// 验证 etag(注意:压缩会改变数据,所以这里的 etag 验证可能需要调整)
println!(
"Test completed successfully with compression: {}, encryption: {}",
is_compress, is_encrypt
);
}
#[tokio::test]
async fn test_compression_with_compressible_data() {
use crate::{CompressReader, DecompressReader};
use rustfs_utils::compress::CompressionAlgorithm;
// Create highly compressible data (repeated pattern)
let pattern = b"Hello, World! This is a test pattern that should compress well. ";
let repeat_count = 16384; // 16K repetitions
let mut data = Vec::new();
for _ in 0..repeat_count {
data.extend_from_slice(pattern);
}
println!("Original data size: {} bytes", data.len());
let reader = BufReader::new(Cursor::new(data.clone()));
let reader = Box::new(WarpReader::new(reader));
let hash_reader = HashReader::new(reader, data.len() as i64, data.len() as i64, None, false).unwrap();
// Test compression
let compress_reader = CompressReader::new(hash_reader, CompressionAlgorithm::Gzip);
let mut compressed_data = Vec::new();
let mut compress_reader = compress_reader;
compress_reader.read_to_end(&mut compressed_data).await.unwrap();
println!("Compressed data size: {} bytes", compressed_data.len());
println!("Compression ratio: {:.2}%", (compressed_data.len() as f64 / data.len() as f64) * 100.0);
// Verify compression actually reduced size for this compressible data
assert!(compressed_data.len() < data.len(), "Compression should reduce size for repetitive data");
// Test decompression
let decompress_reader = DecompressReader::new(Cursor::new(compressed_data), CompressionAlgorithm::Gzip);
let mut decompressed_data = Vec::new();
let mut decompress_reader = decompress_reader;
decompress_reader.read_to_end(&mut decompressed_data).await.unwrap();
// Verify decompressed data matches original
assert_eq!(decompressed_data.len(), data.len());
assert_eq!(&decompressed_data, &data);
println!("Compression/decompression test passed successfully!");
}
#[tokio::test]
async fn test_compression_algorithms() {
use crate::{CompressReader, DecompressReader};
use rustfs_utils::compress::CompressionAlgorithm;
let data = b"This is test data for compression algorithm testing. ".repeat(1000);
println!("Testing with {} bytes of data", data.len());
let algorithms = vec![
CompressionAlgorithm::Gzip,
CompressionAlgorithm::Deflate,
CompressionAlgorithm::Zstd,
];
for algorithm in algorithms {
println!("\nTesting algorithm: {:?}", algorithm);
let reader = BufReader::new(Cursor::new(data.clone()));
let reader = Box::new(WarpReader::new(reader));
let hash_reader = HashReader::new(reader, data.len() as i64, data.len() as i64, None, false).unwrap();
// Compress
let compress_reader = CompressReader::new(hash_reader, algorithm);
let mut compressed_data = Vec::new();
let mut compress_reader = compress_reader;
compress_reader.read_to_end(&mut compressed_data).await.unwrap();
println!(
" Compressed size: {} bytes (ratio: {:.2}%)",
compressed_data.len(),
(compressed_data.len() as f64 / data.len() as f64) * 100.0
);
// Decompress
let decompress_reader = DecompressReader::new(Cursor::new(compressed_data), algorithm);
let mut decompressed_data = Vec::new();
let mut decompress_reader = decompress_reader;
decompress_reader.read_to_end(&mut decompressed_data).await.unwrap();
// Verify
assert_eq!(decompressed_data.len(), data.len());
assert_eq!(&decompressed_data, &data);
println!(" ✓ Algorithm {:?} test passed", algorithm);
}
}
}
+423
View File
@@ -0,0 +1,423 @@
use bytes::Bytes;
use futures::{Stream, TryStreamExt as _};
use http::HeaderMap;
use pin_project_lite::pin_project;
use reqwest::{Client, Method, RequestBuilder};
use std::error::Error as _;
use std::io::{self, Error};
use std::ops::Not as _;
use std::pin::Pin;
use std::sync::LazyLock;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::sync::mpsc;
use tokio_util::io::StreamReader;
use crate::{EtagResolvable, HashReaderDetector, HashReaderMut};
fn get_http_client() -> Client {
// Reuse the HTTP connection pool in the global `reqwest::Client` instance
// TODO: interact with load balancing?
static CLIENT: LazyLock<Client> = LazyLock::new(Client::new);
CLIENT.clone()
}
static HTTP_DEBUG_LOG: bool = false;
#[inline(always)]
fn http_debug_log(args: std::fmt::Arguments) {
if HTTP_DEBUG_LOG {
println!("{}", args);
}
}
macro_rules! http_log {
($($arg:tt)*) => {
http_debug_log(format_args!($($arg)*));
};
}
pin_project! {
pub struct HttpReader {
url:String,
method: Method,
headers: HeaderMap,
inner: StreamReader<Pin<Box<dyn Stream<Item=std::io::Result<Bytes>>+Send+Sync>>, Bytes>,
}
}
impl HttpReader {
pub async fn new(url: String, method: Method, headers: HeaderMap, body: Option<Vec<u8>>) -> io::Result<Self> {
// http_log!("[HttpReader::new] url: {url}, method: {method:?}, headers: {headers:?}");
Self::with_capacity(url, method, headers, body, 0).await
}
/// Create a new HttpReader from a URL. The request is performed immediately.
pub async fn with_capacity(
url: String,
method: Method,
headers: HeaderMap,
body: Option<Vec<u8>>,
_read_buf_size: usize,
) -> io::Result<Self> {
// http_log!(
// "[HttpReader::with_capacity] url: {url}, method: {method:?}, headers: {headers:?}, buf_size: {}",
// _read_buf_size
// );
// First, check if the connection is available (HEAD)
let client = get_http_client();
let head_resp = client.head(&url).headers(headers.clone()).send().await;
match head_resp {
Ok(resp) => {
http_log!("[HttpReader::new] HEAD status: {}", resp.status());
if !resp.status().is_success() {
return Err(Error::other(format!("HEAD failed: url: {}, status {}", url, resp.status())));
}
}
Err(e) => {
http_log!("[HttpReader::new] HEAD error: {e}");
return Err(Error::other(e.source().map(|s| s.to_string()).unwrap_or_else(|| e.to_string())));
}
}
let client = get_http_client();
let mut request: RequestBuilder = client.request(method.clone(), url.clone()).headers(headers.clone());
if let Some(body) = body {
request = request.body(body);
}
let resp = request
.send()
.await
.map_err(|e| Error::other(format!("HttpReader HTTP request error: {}", e)))?;
if resp.status().is_success().not() {
return Err(Error::other(format!(
"HttpReader HTTP request failed with non-200 status {}",
resp.status()
)));
}
let stream = resp
.bytes_stream()
.map_err(|e| Error::other(format!("HttpReader stream error: {}", e)));
Ok(Self {
inner: StreamReader::new(Box::pin(stream)),
url,
method,
headers,
})
}
pub fn url(&self) -> &str {
&self.url
}
pub fn method(&self) -> &Method {
&self.method
}
pub fn headers(&self) -> &HeaderMap {
&self.headers
}
}
impl AsyncRead for HttpReader {
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
// http_log!(
// "[HttpReader::poll_read] url: {}, method: {:?}, buf.remaining: {}",
// self.url,
// self.method,
// buf.remaining()
// );
// Read from the inner stream
Pin::new(&mut self.inner).poll_read(cx, buf)
}
}
impl EtagResolvable for HttpReader {
fn is_etag_reader(&self) -> bool {
false
}
fn try_resolve_etag(&mut self) -> Option<String> {
None
}
}
impl HashReaderDetector for HttpReader {
fn is_hash_reader(&self) -> bool {
false
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
None
}
}
struct ReceiverStream {
receiver: mpsc::Receiver<Option<Bytes>>,
}
impl Stream for ReceiverStream {
type Item = Result<Bytes, std::io::Error>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let poll = Pin::new(&mut self.receiver).poll_recv(cx);
// match &poll {
// Poll::Ready(Some(Some(bytes))) => {
// // http_log!("[ReceiverStream] poll_next: got {} bytes", bytes.len());
// }
// Poll::Ready(Some(None)) => {
// // http_log!("[ReceiverStream] poll_next: sender shutdown");
// }
// Poll::Ready(None) => {
// // http_log!("[ReceiverStream] poll_next: channel closed");
// }
// Poll::Pending => {
// // http_log!("[ReceiverStream] poll_next: pending");
// }
// }
match poll {
Poll::Ready(Some(Some(bytes))) => Poll::Ready(Some(Ok(bytes))),
Poll::Ready(Some(None)) => Poll::Ready(None), // Sender shutdown
Poll::Ready(None) => Poll::Ready(None),
Poll::Pending => Poll::Pending,
}
}
}
pin_project! {
pub struct HttpWriter {
url:String,
method: Method,
headers: HeaderMap,
err_rx: tokio::sync::oneshot::Receiver<std::io::Error>,
sender: tokio::sync::mpsc::Sender<Option<Bytes>>,
handle: tokio::task::JoinHandle<std::io::Result<()>>,
finish:bool,
}
}
impl HttpWriter {
/// Create a new HttpWriter for the given URL. The HTTP request is performed in the background.
pub async fn new(url: String, method: Method, headers: HeaderMap) -> io::Result<Self> {
// http_log!("[HttpWriter::new] url: {url}, method: {method:?}, headers: {headers:?}");
let url_clone = url.clone();
let method_clone = method.clone();
let headers_clone = headers.clone();
// First, try to write empty data to check if writable
let client = get_http_client();
let resp = client.put(&url).headers(headers.clone()).body(Vec::new()).send().await;
match resp {
Ok(resp) => {
// http_log!("[HttpWriter::new] empty PUT status: {}", resp.status());
if !resp.status().is_success() {
return Err(Error::other(format!("Empty PUT failed: status {}", resp.status())));
}
}
Err(e) => {
// http_log!("[HttpWriter::new] empty PUT error: {e}");
return Err(Error::other(format!("Empty PUT failed: {e}")));
}
}
let (sender, receiver) = tokio::sync::mpsc::channel::<Option<Bytes>>(8);
let (err_tx, err_rx) = tokio::sync::oneshot::channel::<io::Error>();
let handle = tokio::spawn(async move {
let stream = ReceiverStream { receiver };
let body = reqwest::Body::wrap_stream(stream);
// http_log!(
// "[HttpWriter::spawn] sending HTTP request: url={url_clone}, method={method_clone:?}, headers={headers_clone:?}"
// );
let client = get_http_client();
let request = client
.request(method_clone, url_clone.clone())
.headers(headers_clone.clone())
.body(body);
// Hold the request until the shutdown signal is received
let response = request.send().await;
match response {
Ok(resp) => {
// http_log!("[HttpWriter::spawn] got response: status={}", resp.status());
if !resp.status().is_success() {
let _ = err_tx.send(Error::other(format!(
"HttpWriter HTTP request failed with non-200 status {}",
resp.status()
)));
return Err(Error::other(format!("HTTP request failed with non-200 status {}", resp.status())));
}
}
Err(e) => {
// http_log!("[HttpWriter::spawn] HTTP request error: {e}");
let _ = err_tx.send(Error::other(format!("HTTP request failed: {}", e)));
return Err(Error::other(format!("HTTP request failed: {}", e)));
}
}
// http_log!("[HttpWriter::spawn] HTTP request completed, exiting");
Ok(())
});
// http_log!("[HttpWriter::new] connection established successfully");
Ok(Self {
url,
method,
headers,
err_rx,
sender,
handle,
finish: false,
})
}
pub fn url(&self) -> &str {
&self.url
}
pub fn method(&self) -> &Method {
&self.method
}
pub fn headers(&self) -> &HeaderMap {
&self.headers
}
}
impl AsyncWrite for HttpWriter {
fn poll_write(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
// http_log!(
// "[HttpWriter::poll_write] url: {}, method: {:?}, buf.len: {}",
// self.url,
// self.method,
// buf.len()
// );
if let Ok(e) = Pin::new(&mut self.err_rx).try_recv() {
return Poll::Ready(Err(e));
}
self.sender
.try_send(Some(Bytes::copy_from_slice(buf)))
.map_err(|e| Error::other(format!("HttpWriter send error: {}", e)))?;
Poll::Ready(Ok(buf.len()))
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
// let url = self.url.clone();
// let method = self.method.clone();
if !self.finish {
// http_log!("[HttpWriter::poll_shutdown] url: {}, method: {:?}", url, method);
self.sender
.try_send(None)
.map_err(|e| Error::other(format!("HttpWriter shutdown error: {}", e)))?;
// http_log!(
// "[HttpWriter::poll_shutdown] sent shutdown signal to HTTP request, url: {}, method: {:?}",
// url,
// method
// );
self.finish = true;
}
// Wait for the HTTP request to complete
use futures::FutureExt;
match Pin::new(&mut self.get_mut().handle).poll_unpin(_cx) {
Poll::Ready(Ok(_)) => {
// http_log!(
// "[HttpWriter::poll_shutdown] HTTP request finished successfully, url: {}, method: {:?}",
// url,
// method
// );
}
Poll::Ready(Err(e)) => {
// http_log!("[HttpWriter::poll_shutdown] HTTP request failed: {e}, url: {}, method: {:?}", url, method);
return Poll::Ready(Err(Error::other(format!("HTTP request failed: {}", e))));
}
Poll::Pending => {
// http_log!("[HttpWriter::poll_shutdown] HTTP request pending, url: {}, method: {:?}", url, method);
return Poll::Pending;
}
}
Poll::Ready(Ok(()))
}
}
// #[cfg(test)]
// mod tests {
// use super::*;
// use reqwest::Method;
// use std::vec;
// use tokio::io::{AsyncReadExt, AsyncWriteExt};
// #[tokio::test]
// async fn test_http_writer_err() {
// // Use a real local server for integration, or mockito for unit test
// // Here, we use the Go test server at 127.0.0.1:8081 (scripts/testfile.go)
// let url = "http://127.0.0.1:8081/testfile".to_string();
// let data = vec![42u8; 8];
// // Write
// // 添加 header X-Deny-Write = 1 模拟不可写入的情况
// let mut headers = HeaderMap::new();
// headers.insert("X-Deny-Write", "1".parse().unwrap());
// // 这里我们使用 PUT 方法
// let writer_result = HttpWriter::new(url.clone(), Method::PUT, headers).await;
// match writer_result {
// Ok(mut writer) => {
// // 如果能创建成功,写入应该报错
// let write_result = writer.write_all(&data).await;
// assert!(write_result.is_err(), "write_all should fail when server denies write");
// if let Err(e) = write_result {
// println!("write_all error: {e}");
// }
// let shutdown_result = writer.shutdown().await;
// if let Err(e) = shutdown_result {
// println!("shutdown error: {e}");
// }
// }
// Err(e) => {
// // 直接构造失败也可以
// println!("HttpWriter::new error: {e}");
// assert!(
// e.to_string().contains("Empty PUT failed") || e.to_string().contains("Forbidden"),
// "unexpected error: {e}"
// );
// return;
// }
// }
// // Should not reach here
// panic!("HttpWriter should not allow writing when server denies write");
// }
// #[tokio::test]
// async fn test_http_writer_and_reader_ok() {
// // 使用本地 Go 测试服务器
// let url = "http://127.0.0.1:8081/testfile".to_string();
// let data = vec![99u8; 512 * 1024]; // 512KB of data
// // Write (不加 X-Deny-Write)
// let headers = HeaderMap::new();
// let mut writer = HttpWriter::new(url.clone(), Method::PUT, headers).await.unwrap();
// writer.write_all(&data).await.unwrap();
// writer.shutdown().await.unwrap();
// http_log!("Wrote {} bytes to {} (ok case)", data.len(), url);
// // Read back
// let mut reader = HttpReader::with_capacity(url.clone(), Method::GET, HeaderMap::new(), 8192)
// .await
// .unwrap();
// let mut buf = Vec::new();
// reader.read_to_end(&mut buf).await.unwrap();
// assert_eq!(buf, data);
// // println!("Read {} bytes from {} (ok case)", buf.len(), url);
// // tokio::time::sleep(std::time::Duration::from_secs(2)).await; // Wait for server to process
// // println!("[test_http_writer_and_reader_ok] completed successfully");
// }
// }
+69
View File
@@ -0,0 +1,69 @@
mod limit_reader;
pub use limit_reader::LimitReader;
mod etag_reader;
pub use etag_reader::EtagReader;
mod compress_index;
mod compress_reader;
pub use compress_reader::{CompressReader, DecompressReader};
mod encrypt_reader;
pub use encrypt_reader::{DecryptReader, EncryptReader};
mod hardlimit_reader;
pub use hardlimit_reader::HardLimitReader;
mod hash_reader;
pub use hash_reader::*;
pub mod reader;
pub use reader::WarpReader;
mod writer;
pub use writer::*;
mod http_reader;
pub use http_reader::*;
pub use compress_index::TryGetIndex;
mod etag;
pub trait Reader: tokio::io::AsyncRead + Unpin + Send + Sync + EtagResolvable + HashReaderDetector + TryGetIndex {}
// Trait for types that can be recursively searched for etag capability
pub trait EtagResolvable {
fn is_etag_reader(&self) -> bool {
false
}
fn try_resolve_etag(&mut self) -> Option<String> {
None
}
}
// Generic function that can work with any EtagResolvable type
pub fn resolve_etag_generic<R>(reader: &mut R) -> Option<String>
where
R: EtagResolvable,
{
reader.try_resolve_etag()
}
/// Trait to detect and manipulate HashReader instances
pub trait HashReaderDetector {
fn is_hash_reader(&self) -> bool {
false
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
None
}
}
impl Reader for crate::HashReader {}
impl Reader for crate::HardLimitReader {}
impl Reader for crate::EtagReader {}
impl<R> Reader for crate::CompressReader<R> where R: Reader {}
impl<R> Reader for crate::EncryptReader<R> where R: Reader {}
+188
View File
@@ -0,0 +1,188 @@
//! LimitReader: a wrapper for AsyncRead that limits the total number of bytes read.
//!
//! # Example
//! ```
//! use tokio::io::{AsyncReadExt, BufReader};
//! use rustfs_rio::LimitReader;
//!
//! #[tokio::main]
//! async fn main() {
//! let data = b"hello world";
//! let reader = BufReader::new(&data[..]);
//! let mut limit_reader = LimitReader::new(reader, data.len());
//!
//! let mut buf = Vec::new();
//! let n = limit_reader.read_to_end(&mut buf).await.unwrap();
//! assert_eq!(n, data.len());
//! assert_eq!(&buf, data);
//! }
//! ```
use pin_project_lite::pin_project;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
use crate::{EtagResolvable, HashReaderDetector, HashReaderMut};
pin_project! {
#[derive(Debug)]
pub struct LimitReader<R> {
#[pin]
pub inner: R,
limit: usize,
read: usize,
}
}
/// A wrapper for AsyncRead that limits the total number of bytes read.
impl<R> LimitReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
/// Create a new LimitReader wrapping `inner`, with a total read limit of `limit` bytes.
pub fn new(inner: R, limit: usize) -> Self {
Self { inner, limit, read: 0 }
}
}
impl<R> AsyncRead for LimitReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
let mut this = self.project();
let remaining = this.limit.saturating_sub(*this.read);
if remaining == 0 {
return Poll::Ready(Ok(()));
}
let orig_remaining = buf.remaining();
let allowed = remaining.min(orig_remaining);
if allowed == 0 {
return Poll::Ready(Ok(()));
}
if allowed == orig_remaining {
let before_size = buf.filled().len();
let poll = this.inner.as_mut().poll_read(cx, buf);
if let Poll::Ready(Ok(())) = &poll {
let n = buf.filled().len() - before_size;
*this.read += n;
}
poll
} else {
let mut temp = vec![0u8; allowed];
let mut temp_buf = ReadBuf::new(&mut temp);
let poll = this.inner.as_mut().poll_read(cx, &mut temp_buf);
if let Poll::Ready(Ok(())) = &poll {
let n = temp_buf.filled().len();
buf.put_slice(temp_buf.filled());
*this.read += n;
}
poll
}
}
}
impl<R> EtagResolvable for LimitReader<R>
where
R: EtagResolvable,
{
fn try_resolve_etag(&mut self) -> Option<String> {
self.inner.try_resolve_etag()
}
}
impl<R> HashReaderDetector for LimitReader<R>
where
R: HashReaderDetector,
{
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
#[cfg(test)]
mod tests {
use std::io::Cursor;
use super::*;
use tokio::io::{AsyncReadExt, BufReader};
#[tokio::test]
async fn test_limit_reader_exact() {
let data = b"hello world";
let reader = BufReader::new(&data[..]);
let mut limit_reader = LimitReader::new(reader, data.len());
let mut buf = Vec::new();
let n = limit_reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, data.len());
assert_eq!(&buf, data);
}
#[tokio::test]
async fn test_limit_reader_less_than_data() {
let data = b"hello world";
let reader = BufReader::new(&data[..]);
let mut limit_reader = LimitReader::new(reader, 5);
let mut buf = Vec::new();
let n = limit_reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, 5);
assert_eq!(&buf, b"hello");
}
#[tokio::test]
async fn test_limit_reader_zero() {
let data = b"hello world";
let reader = BufReader::new(&data[..]);
let mut limit_reader = LimitReader::new(reader, 0);
let mut buf = Vec::new();
let n = limit_reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, 0);
assert!(buf.is_empty());
}
#[tokio::test]
async fn test_limit_reader_multiple_reads() {
let data = b"abcdefghij";
let reader = BufReader::new(&data[..]);
let mut limit_reader = LimitReader::new(reader, 7);
let mut buf1 = [0u8; 3];
let n1 = limit_reader.read(&mut buf1).await.unwrap();
assert_eq!(n1, 3);
assert_eq!(&buf1, b"abc");
let mut buf2 = [0u8; 5];
let n2 = limit_reader.read(&mut buf2).await.unwrap();
assert_eq!(n2, 4);
assert_eq!(&buf2[..n2], b"defg");
let mut buf3 = [0u8; 2];
let n3 = limit_reader.read(&mut buf3).await.unwrap();
assert_eq!(n3, 0);
}
#[tokio::test]
async fn test_limit_reader_large_file() {
use rand::Rng;
// Generate a 3MB random byte array for testing
let size = 3 * 1024 * 1024;
let mut data = vec![0u8; size];
rand::rng().fill(&mut data[..]);
let reader = Cursor::new(data.clone());
let mut limit_reader = LimitReader::new(reader, size);
// Read data into buffer
let mut buf = Vec::new();
let n = limit_reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(n, size);
assert_eq!(buf.len(), size);
assert_eq!(&buf, &data);
}
}
+30
View File
@@ -0,0 +1,30 @@
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
use crate::compress_index::TryGetIndex;
use crate::{EtagResolvable, HashReaderDetector, Reader};
pub struct WarpReader<R> {
inner: R,
}
impl<R: AsyncRead + Unpin + Send + Sync> WarpReader<R> {
pub fn new(inner: R) -> Self {
Self { inner }
}
}
impl<R: AsyncRead + Unpin + Send + Sync> AsyncRead for WarpReader<R> {
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.inner).poll_read(cx, buf)
}
}
impl<R: AsyncRead + Unpin + Send + Sync> HashReaderDetector for WarpReader<R> {}
impl<R: AsyncRead + Unpin + Send + Sync> EtagResolvable for WarpReader<R> {}
impl<R: AsyncRead + Unpin + Send + Sync> TryGetIndex for WarpReader<R> {}
impl<R: AsyncRead + Unpin + Send + Sync> Reader for WarpReader<R> {}
+92
View File
@@ -0,0 +1,92 @@
use std::io::Cursor;
use std::pin::Pin;
use tokio::io::AsyncWrite;
use crate::HttpWriter;
pub enum Writer {
Cursor(Cursor<Vec<u8>>),
Http(HttpWriter),
Other(Box<dyn AsyncWrite + Unpin + Send + Sync>),
}
impl Writer {
/// Create a Writer::Other from any AsyncWrite + Unpin + Send type.
pub fn from_tokio_writer<W>(w: W) -> Self
where
W: AsyncWrite + Unpin + Send + Sync + 'static,
{
Writer::Other(Box::new(w))
}
pub fn from_cursor(w: Cursor<Vec<u8>>) -> Self {
Writer::Cursor(w)
}
pub fn from_http(w: HttpWriter) -> Self {
Writer::Http(w)
}
pub fn into_cursor_inner(self) -> Option<Vec<u8>> {
match self {
Writer::Cursor(w) => Some(w.into_inner()),
_ => None,
}
}
pub fn as_cursor(&mut self) -> Option<&mut Cursor<Vec<u8>>> {
match self {
Writer::Cursor(w) => Some(w),
_ => None,
}
}
pub fn as_http(&mut self) -> Option<&mut HttpWriter> {
match self {
Writer::Http(w) => Some(w),
_ => None,
}
}
pub fn into_http(self) -> Option<HttpWriter> {
match self {
Writer::Http(w) => Some(w),
_ => None,
}
}
pub fn into_cursor(self) -> Option<Cursor<Vec<u8>>> {
match self {
Writer::Cursor(w) => Some(w),
_ => None,
}
}
}
impl AsyncWrite for Writer {
fn poll_write(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &[u8],
) -> std::task::Poll<std::io::Result<usize>> {
match self.get_mut() {
Writer::Cursor(w) => Pin::new(w).poll_write(cx, buf),
Writer::Http(w) => Pin::new(w).poll_write(cx, buf),
Writer::Other(w) => Pin::new(w.as_mut()).poll_write(cx, buf),
}
}
fn poll_flush(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<std::io::Result<()>> {
match self.get_mut() {
Writer::Cursor(w) => Pin::new(w).poll_flush(cx),
Writer::Http(w) => Pin::new(w).poll_flush(cx),
Writer::Other(w) => Pin::new(w.as_mut()).poll_flush(cx),
}
}
fn poll_shutdown(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<std::io::Result<()>> {
match self.get_mut() {
Writer::Cursor(w) => Pin::new(w).poll_shutdown(cx),
Writer::Http(w) => Pin::new(w).poll_shutdown(cx),
Writer::Other(w) => Pin::new(w.as_mut()).poll_shutdown(cx),
}
}
}
+34 -2
View File
@@ -7,15 +7,40 @@ rust-version.workspace = true
version.workspace = true
[dependencies]
base64-simd = { workspace = true, optional = true }
blake3 = { workspace = true, optional = true }
crc32fast.workspace = true
hex-simd = { workspace = true, optional = true }
highway = { workspace = true, optional = true }
lazy_static = { workspace = true, optional = true }
local-ip-address = { workspace = true, optional = true }
rustfs-config = { workspace = true, features = ["constants"] }
md-5 = { workspace = true, optional = true }
netif = { workspace = true, optional = true }
nix = { workspace = true, optional = true }
regex = { workspace = true, optional = true }
rustls = { workspace = true, optional = true }
rustls-pemfile = { workspace = true, optional = true }
rustls-pki-types = { workspace = true, optional = true }
serde = { workspace = true, optional = true }
sha2 = { workspace = true, optional = true }
siphasher = { workspace = true, optional = true }
tempfile = { workspace = true, optional = true }
tokio = { workspace = true, optional = true, features = ["io-util", "macros"] }
tracing = { workspace = true }
url = { workspace = true, optional = true }
flate2 = { workspace = true, optional = true }
brotli = { workspace = true, optional = true }
zstd = { workspace = true, optional = true }
snap = { workspace = true, optional = true }
lz4 = { workspace = true, optional = true }
[dev-dependencies]
tempfile = { workspace = true }
rand = { workspace = true }
[target.'cfg(windows)'.dependencies]
winapi = { workspace = true, optional = true, features = ["std", "fileapi", "minwindef", "ntdef", "winnt"] }
[lints]
workspace = true
@@ -24,6 +49,13 @@ workspace = true
default = ["ip"] # features that are enabled by default
ip = ["dep:local-ip-address"] # ip characteristics and their dependencies
tls = ["dep:rustls", "dep:rustls-pemfile", "dep:rustls-pki-types"] # tls characteristics and their dependencies
net = ["ip"] # empty network features
net = ["ip", "dep:url", "dep:netif", "dep:lazy_static"] # empty network features
io = ["dep:tokio"]
path = []
compress = ["dep:flate2", "dep:brotli", "dep:snap", "dep:lz4", "dep:zstd"]
string = ["dep:regex", "dep:lazy_static"]
crypto = ["dep:base64-simd", "dep:hex-simd"]
hash = ["dep:highway", "dep:md-5", "dep:sha2", "dep:blake3", "dep:serde", "dep:siphasher"]
os = ["dep:nix", "dep:tempfile", "winapi"] # operating system utilities
integration = [] # integration test features
full = ["ip", "tls", "net", "integration"] # all features
full = ["ip", "tls", "net", "io", "hash", "os", "integration", "path", "crypto", "string", "compress"] # all features
+12 -8
View File
@@ -396,10 +396,12 @@ mod tests {
// Should fail because no certificates found
let result = load_all_certs_from_directory(temp_dir.path().to_str().unwrap());
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("No valid certificate/private key pair found"));
assert!(
result
.unwrap_err()
.to_string()
.contains("No valid certificate/private key pair found")
);
}
#[test]
@@ -412,10 +414,12 @@ mod tests {
let result = load_all_certs_from_directory(unicode_dir.to_str().unwrap());
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("No valid certificate/private key pair found"));
assert!(
result
.unwrap_err()
.to_string()
.contains("No valid certificate/private key pair found")
);
}
#[test]
+318
View File
@@ -0,0 +1,318 @@
use std::io::Write;
use tokio::io;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum CompressionAlgorithm {
None,
Gzip,
Deflate,
Zstd,
#[default]
Lz4,
Brotli,
Snappy,
}
impl CompressionAlgorithm {
pub fn as_str(&self) -> &str {
match self {
CompressionAlgorithm::None => "none",
CompressionAlgorithm::Gzip => "gzip",
CompressionAlgorithm::Deflate => "deflate",
CompressionAlgorithm::Zstd => "zstd",
CompressionAlgorithm::Lz4 => "lz4",
CompressionAlgorithm::Brotli => "brotli",
CompressionAlgorithm::Snappy => "snappy",
}
}
}
impl std::fmt::Display for CompressionAlgorithm {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
impl std::str::FromStr for CompressionAlgorithm {
type Err = std::io::Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"gzip" => Ok(CompressionAlgorithm::Gzip),
"deflate" => Ok(CompressionAlgorithm::Deflate),
"zstd" => Ok(CompressionAlgorithm::Zstd),
"lz4" => Ok(CompressionAlgorithm::Lz4),
"brotli" => Ok(CompressionAlgorithm::Brotli),
"snappy" => Ok(CompressionAlgorithm::Snappy),
"none" => Ok(CompressionAlgorithm::None),
_ => Err(std::io::Error::other(format!("Unsupported compression algorithm: {}", s))),
}
}
}
pub fn compress_block(input: &[u8], algorithm: CompressionAlgorithm) -> Vec<u8> {
match algorithm {
CompressionAlgorithm::Gzip => {
let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
let _ = encoder.write_all(input);
let _ = encoder.flush();
encoder.finish().unwrap_or_default()
}
CompressionAlgorithm::Deflate => {
let mut encoder = flate2::write::DeflateEncoder::new(Vec::new(), flate2::Compression::default());
let _ = encoder.write_all(input);
let _ = encoder.flush();
encoder.finish().unwrap_or_default()
}
CompressionAlgorithm::Zstd => {
let mut encoder = zstd::Encoder::new(Vec::new(), 0).expect("zstd encoder");
let _ = encoder.write_all(input);
encoder.finish().unwrap_or_default()
}
CompressionAlgorithm::Lz4 => {
let mut encoder = lz4::EncoderBuilder::new().build(Vec::new()).expect("lz4 encoder");
let _ = encoder.write_all(input);
let (out, result) = encoder.finish();
result.expect("lz4 finish");
out
}
CompressionAlgorithm::Brotli => {
let mut out = Vec::new();
brotli::CompressorWriter::new(&mut out, 4096, 5, 22)
.write_all(input)
.expect("brotli compress");
out
}
CompressionAlgorithm::Snappy => {
let mut encoder = snap::write::FrameEncoder::new(Vec::new());
let _ = encoder.write_all(input);
encoder.into_inner().unwrap_or_default()
}
CompressionAlgorithm::None => input.to_vec(),
}
}
pub fn decompress_block(compressed: &[u8], algorithm: CompressionAlgorithm) -> io::Result<Vec<u8>> {
match algorithm {
CompressionAlgorithm::Gzip => {
let mut decoder = flate2::read::GzDecoder::new(std::io::Cursor::new(compressed));
let mut out = Vec::new();
std::io::Read::read_to_end(&mut decoder, &mut out)?;
Ok(out)
}
CompressionAlgorithm::Deflate => {
let mut decoder = flate2::read::DeflateDecoder::new(std::io::Cursor::new(compressed));
let mut out = Vec::new();
std::io::Read::read_to_end(&mut decoder, &mut out)?;
Ok(out)
}
CompressionAlgorithm::Zstd => {
let mut decoder = zstd::Decoder::new(std::io::Cursor::new(compressed))?;
let mut out = Vec::new();
std::io::Read::read_to_end(&mut decoder, &mut out)?;
Ok(out)
}
CompressionAlgorithm::Lz4 => {
let mut decoder = lz4::Decoder::new(std::io::Cursor::new(compressed)).expect("lz4 decoder");
let mut out = Vec::new();
std::io::Read::read_to_end(&mut decoder, &mut out)?;
Ok(out)
}
CompressionAlgorithm::Brotli => {
let mut out = Vec::new();
let mut decoder = brotli::Decompressor::new(std::io::Cursor::new(compressed), 4096);
std::io::Read::read_to_end(&mut decoder, &mut out)?;
Ok(out)
}
CompressionAlgorithm::Snappy => {
let mut decoder = snap::read::FrameDecoder::new(std::io::Cursor::new(compressed));
let mut out = Vec::new();
std::io::Read::read_to_end(&mut decoder, &mut out)?;
Ok(out)
}
CompressionAlgorithm::None => Ok(Vec::new()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::str::FromStr;
use std::time::Instant;
#[test]
fn test_compress_decompress_gzip() {
let data = b"hello gzip compress";
let compressed = compress_block(data, CompressionAlgorithm::Gzip);
let decompressed = decompress_block(&compressed, CompressionAlgorithm::Gzip).unwrap();
assert_eq!(decompressed, data);
}
#[test]
fn test_compress_decompress_deflate() {
let data = b"hello deflate compress";
let compressed = compress_block(data, CompressionAlgorithm::Deflate);
let decompressed = decompress_block(&compressed, CompressionAlgorithm::Deflate).unwrap();
assert_eq!(decompressed, data);
}
#[test]
fn test_compress_decompress_zstd() {
let data = b"hello zstd compress";
let compressed = compress_block(data, CompressionAlgorithm::Zstd);
let decompressed = decompress_block(&compressed, CompressionAlgorithm::Zstd).unwrap();
assert_eq!(decompressed, data);
}
#[test]
fn test_compress_decompress_lz4() {
let data = b"hello lz4 compress";
let compressed = compress_block(data, CompressionAlgorithm::Lz4);
let decompressed = decompress_block(&compressed, CompressionAlgorithm::Lz4).unwrap();
assert_eq!(decompressed, data);
}
#[test]
fn test_compress_decompress_brotli() {
let data = b"hello brotli compress";
let compressed = compress_block(data, CompressionAlgorithm::Brotli);
let decompressed = decompress_block(&compressed, CompressionAlgorithm::Brotli).unwrap();
assert_eq!(decompressed, data);
}
#[test]
fn test_compress_decompress_snappy() {
let data = b"hello snappy compress";
let compressed = compress_block(data, CompressionAlgorithm::Snappy);
let decompressed = decompress_block(&compressed, CompressionAlgorithm::Snappy).unwrap();
assert_eq!(decompressed, data);
}
#[test]
fn test_from_str() {
assert_eq!(CompressionAlgorithm::from_str("gzip").unwrap(), CompressionAlgorithm::Gzip);
assert_eq!(CompressionAlgorithm::from_str("deflate").unwrap(), CompressionAlgorithm::Deflate);
assert_eq!(CompressionAlgorithm::from_str("zstd").unwrap(), CompressionAlgorithm::Zstd);
assert_eq!(CompressionAlgorithm::from_str("lz4").unwrap(), CompressionAlgorithm::Lz4);
assert_eq!(CompressionAlgorithm::from_str("brotli").unwrap(), CompressionAlgorithm::Brotli);
assert_eq!(CompressionAlgorithm::from_str("snappy").unwrap(), CompressionAlgorithm::Snappy);
assert!(CompressionAlgorithm::from_str("unknown").is_err());
}
#[test]
fn test_compare_compression_algorithms() {
use std::time::Instant;
let data = vec![42u8; 1024 * 100]; // 100KB of repetitive data
// let mut data = vec![0u8; 1024 * 1024];
// rand::thread_rng().fill(&mut data[..]);
let start = Instant::now();
let mut times = Vec::new();
times.push(("original", start.elapsed(), data.len()));
let start = Instant::now();
let gzip = compress_block(&data, CompressionAlgorithm::Gzip);
let gzip_time = start.elapsed();
times.push(("gzip", gzip_time, gzip.len()));
let start = Instant::now();
let deflate = compress_block(&data, CompressionAlgorithm::Deflate);
let deflate_time = start.elapsed();
times.push(("deflate", deflate_time, deflate.len()));
let start = Instant::now();
let zstd = compress_block(&data, CompressionAlgorithm::Zstd);
let zstd_time = start.elapsed();
times.push(("zstd", zstd_time, zstd.len()));
let start = Instant::now();
let lz4 = compress_block(&data, CompressionAlgorithm::Lz4);
let lz4_time = start.elapsed();
times.push(("lz4", lz4_time, lz4.len()));
let start = Instant::now();
let brotli = compress_block(&data, CompressionAlgorithm::Brotli);
let brotli_time = start.elapsed();
times.push(("brotli", brotli_time, brotli.len()));
let start = Instant::now();
let snappy = compress_block(&data, CompressionAlgorithm::Snappy);
let snappy_time = start.elapsed();
times.push(("snappy", snappy_time, snappy.len()));
println!("Compression results:");
for (name, dur, size) in &times {
println!("{}: {} bytes, {:?}", name, size, dur);
}
// All should decompress to the original
assert_eq!(decompress_block(&gzip, CompressionAlgorithm::Gzip).unwrap(), data);
assert_eq!(decompress_block(&deflate, CompressionAlgorithm::Deflate).unwrap(), data);
assert_eq!(decompress_block(&zstd, CompressionAlgorithm::Zstd).unwrap(), data);
assert_eq!(decompress_block(&lz4, CompressionAlgorithm::Lz4).unwrap(), data);
assert_eq!(decompress_block(&brotli, CompressionAlgorithm::Brotli).unwrap(), data);
assert_eq!(decompress_block(&snappy, CompressionAlgorithm::Snappy).unwrap(), data);
// All compressed results should not be empty
assert!(
!gzip.is_empty()
&& !deflate.is_empty()
&& !zstd.is_empty()
&& !lz4.is_empty()
&& !brotli.is_empty()
&& !snappy.is_empty()
);
}
#[test]
fn test_compression_benchmark() {
let sizes = [128 * 1024, 512 * 1024, 1024 * 1024];
let algorithms = [
CompressionAlgorithm::Gzip,
CompressionAlgorithm::Deflate,
CompressionAlgorithm::Zstd,
CompressionAlgorithm::Lz4,
CompressionAlgorithm::Brotli,
CompressionAlgorithm::Snappy,
];
println!("\n压缩算法基准测试结果:");
println!(
"{:<10} {:<10} {:<15} {:<15} {:<15}",
"数据大小", "算法", "压缩时间(ms)", "压缩后大小", "压缩率"
);
for size in sizes {
// 生成可压缩的数据(重复的文本模式)
let pattern = b"Hello, this is a test pattern that will be repeated multiple times to create compressible data. ";
let data: Vec<u8> = pattern.iter().cycle().take(size).copied().collect();
for algo in algorithms {
// 压缩测试
let start = Instant::now();
let compressed = compress_block(&data, algo);
let compress_time = start.elapsed();
// 解压测试
let start = Instant::now();
let _decompressed = decompress_block(&compressed, algo).unwrap();
let _decompress_time = start.elapsed();
// 计算压缩率
let compression_ratio = (size as f64 / compressed.len() as f64) as f32;
println!(
"{:<10} {:<10} {:<15.2} {:<15} {:<15.2}x",
format!("{}KB", size / 1024),
algo.as_str(),
compress_time.as_secs_f64() * 1000.0,
compressed.len(),
compression_ratio
);
// 验证解压结果
assert_eq!(_decompressed, data);
}
println!(); // 添加空行分隔不同大小的结果
}
}
}
+39
View File
@@ -0,0 +1,39 @@
pub fn base64_encode(input: &[u8]) -> String {
base64_simd::URL_SAFE_NO_PAD.encode_to_string(input)
}
pub fn base64_decode(input: &[u8]) -> Result<Vec<u8>, base64_simd::Error> {
base64_simd::URL_SAFE_NO_PAD.decode_to_vec(input)
}
pub fn hex(data: impl AsRef<[u8]>) -> String {
hex_simd::encode_to_string(data, hex_simd::AsciiCase::Lower)
}
// #[cfg(windows)]
// pub fn sha256(data: &[u8]) -> impl AsRef<[u8; 32]> {
// use sha2::{Digest, Sha256};
// <Sha256 as Digest>::digest(data)
// }
// #[cfg(not(windows))]
// pub fn sha256(data: &[u8]) -> impl AsRef<[u8]> {
// use openssl::hash::{Hasher, MessageDigest};
// let mut h = Hasher::new(MessageDigest::sha256()).unwrap();
// h.update(data).unwrap();
// h.finish().unwrap()
// }
#[test]
fn test_base64_encoding_decoding() {
let original_uuid_timestamp = "c0194290-d911-45cb-8e12-79ec563f46a8x1735460504394878000";
let encoded_string = base64_encode(original_uuid_timestamp.as_bytes());
println!("Encoded: {}", &encoded_string);
let decoded_bytes = base64_decode(encoded_string.clone().as_bytes()).unwrap();
let decoded_string = String::from_utf8(decoded_bytes).unwrap();
assert_eq!(decoded_string, original_uuid_timestamp)
}
+206
View File
@@ -0,0 +1,206 @@
use highway::{HighwayHash, HighwayHasher, Key};
use md5::{Digest, Md5};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
/// The fixed key for HighwayHash256. DO NOT change for compatibility.
const HIGHWAY_HASH256_KEY: [u64; 4] = [3, 4, 2, 1];
#[derive(Serialize, Deserialize, Debug, PartialEq, Default, Clone, Eq, Hash)]
/// Supported hash algorithms for bitrot protection.
pub enum HashAlgorithm {
// SHA256 represents the SHA-256 hash function
SHA256,
// HighwayHash256 represents the HighwayHash-256 hash function
HighwayHash256,
// HighwayHash256S represents the Streaming HighwayHash-256 hash function
#[default]
HighwayHash256S,
// BLAKE2b512 represents the BLAKE2b-512 hash function
BLAKE2b512,
/// MD5 (128-bit)
Md5,
/// No hash (for testing or unprotected data)
None,
}
enum HashEncoded {
Md5([u8; 16]),
Sha256([u8; 32]),
HighwayHash256([u8; 32]),
HighwayHash256S([u8; 32]),
Blake2b512(blake3::Hash),
None,
}
impl AsRef<[u8]> for HashEncoded {
#[inline]
fn as_ref(&self) -> &[u8] {
match self {
HashEncoded::Md5(hash) => hash.as_ref(),
HashEncoded::Sha256(hash) => hash.as_ref(),
HashEncoded::HighwayHash256(hash) => hash.as_ref(),
HashEncoded::HighwayHash256S(hash) => hash.as_ref(),
HashEncoded::Blake2b512(hash) => hash.as_bytes(),
HashEncoded::None => &[],
}
}
}
#[inline]
fn u8x32_from_u64x4(input: [u64; 4]) -> [u8; 32] {
let mut output = [0u8; 32];
for (i, &n) in input.iter().enumerate() {
output[i * 8..(i + 1) * 8].copy_from_slice(&n.to_le_bytes());
}
output
}
impl HashAlgorithm {
/// Hash the input data and return the hash result as Vec<u8>.
pub fn hash_encode(&self, data: &[u8]) -> impl AsRef<[u8]> {
match self {
HashAlgorithm::Md5 => HashEncoded::Md5(Md5::digest(data).into()),
HashAlgorithm::HighwayHash256 => {
let mut hasher = HighwayHasher::new(Key(HIGHWAY_HASH256_KEY));
hasher.append(data);
HashEncoded::HighwayHash256(u8x32_from_u64x4(hasher.finalize256()))
}
HashAlgorithm::SHA256 => HashEncoded::Sha256(Sha256::digest(data).into()),
HashAlgorithm::HighwayHash256S => {
let mut hasher = HighwayHasher::new(Key(HIGHWAY_HASH256_KEY));
hasher.append(data);
HashEncoded::HighwayHash256S(u8x32_from_u64x4(hasher.finalize256()))
}
HashAlgorithm::BLAKE2b512 => HashEncoded::Blake2b512(blake3::hash(data)),
HashAlgorithm::None => HashEncoded::None,
}
}
/// Return the output size in bytes for the hash algorithm.
pub fn size(&self) -> usize {
match self {
HashAlgorithm::SHA256 => 32,
HashAlgorithm::HighwayHash256 => 32,
HashAlgorithm::HighwayHash256S => 32,
HashAlgorithm::BLAKE2b512 => 32, // blake3 outputs 32 bytes by default
HashAlgorithm::Md5 => 16,
HashAlgorithm::None => 0,
}
}
}
use crc32fast::Hasher;
use siphasher::sip::SipHasher;
pub fn sip_hash(key: &str, cardinality: usize, id: &[u8; 16]) -> usize {
// 你的密钥,必须是 16 字节
// 计算字符串的 SipHash 值
let result = SipHasher::new_with_key(id).hash(key.as_bytes());
result as usize % cardinality
}
pub fn crc_hash(key: &str, cardinality: usize) -> usize {
let mut hasher = Hasher::new(); // 创建一个新的哈希器
hasher.update(key.as_bytes()); // 更新哈希状态,添加数据
let checksum = hasher.finalize();
checksum as usize % cardinality
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hash_algorithm_sizes() {
assert_eq!(HashAlgorithm::Md5.size(), 16);
assert_eq!(HashAlgorithm::HighwayHash256.size(), 32);
assert_eq!(HashAlgorithm::HighwayHash256S.size(), 32);
assert_eq!(HashAlgorithm::SHA256.size(), 32);
assert_eq!(HashAlgorithm::BLAKE2b512.size(), 32);
assert_eq!(HashAlgorithm::None.size(), 0);
}
#[test]
fn test_hash_encode_none() {
let data = b"test data";
let hash = HashAlgorithm::None.hash_encode(data);
let hash = hash.as_ref();
assert_eq!(hash.len(), 0);
}
#[test]
fn test_hash_encode_md5() {
let data = b"test data";
let hash = HashAlgorithm::Md5.hash_encode(data);
let hash = hash.as_ref();
assert_eq!(hash.len(), 16);
// MD5 should be deterministic
let hash2 = HashAlgorithm::Md5.hash_encode(data);
let hash2 = hash2.as_ref();
assert_eq!(hash, hash2);
}
#[test]
fn test_hash_encode_highway() {
let data = b"test data";
let hash = HashAlgorithm::HighwayHash256.hash_encode(data);
let hash = hash.as_ref();
assert_eq!(hash.len(), 32);
// HighwayHash should be deterministic
let hash2 = HashAlgorithm::HighwayHash256.hash_encode(data);
let hash2 = hash2.as_ref();
assert_eq!(hash, hash2);
}
#[test]
fn test_hash_encode_sha256() {
let data = b"test data";
let hash = HashAlgorithm::SHA256.hash_encode(data);
let hash = hash.as_ref();
assert_eq!(hash.len(), 32);
// SHA256 should be deterministic
let hash2 = HashAlgorithm::SHA256.hash_encode(data);
let hash2 = hash2.as_ref();
assert_eq!(hash, hash2);
}
#[test]
fn test_hash_encode_blake2b512() {
let data = b"test data";
let hash = HashAlgorithm::BLAKE2b512.hash_encode(data);
let hash = hash.as_ref();
assert_eq!(hash.len(), 32); // blake3 outputs 32 bytes by default
// BLAKE2b512 should be deterministic
let hash2 = HashAlgorithm::BLAKE2b512.hash_encode(data);
let hash2 = hash2.as_ref();
assert_eq!(hash, hash2);
}
#[test]
fn test_different_data_different_hashes() {
let data1 = b"test data 1";
let data2 = b"test data 2";
let md5_hash1 = HashAlgorithm::Md5.hash_encode(data1);
let md5_hash2 = HashAlgorithm::Md5.hash_encode(data2);
assert_ne!(md5_hash1.as_ref(), md5_hash2.as_ref());
let highway_hash1 = HashAlgorithm::HighwayHash256.hash_encode(data1);
let highway_hash2 = HashAlgorithm::HighwayHash256.hash_encode(data2);
assert_ne!(highway_hash1.as_ref(), highway_hash2.as_ref());
let sha256_hash1 = HashAlgorithm::SHA256.hash_encode(data1);
let sha256_hash2 = HashAlgorithm::SHA256.hash_encode(data2);
assert_ne!(sha256_hash1.as_ref(), sha256_hash2.as_ref());
let blake_hash1 = HashAlgorithm::BLAKE2b512.hash_encode(data1);
let blake_hash2 = HashAlgorithm::BLAKE2b512.hash_encode(data2);
assert_ne!(blake_hash1.as_ref(), blake_hash2.as_ref());
}
}
+231
View File
@@ -0,0 +1,231 @@
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
/// Write all bytes from buf to writer, returning the total number of bytes written.
pub async fn write_all<W: AsyncWrite + Send + Sync + Unpin>(writer: &mut W, buf: &[u8]) -> std::io::Result<usize> {
let mut total = 0;
while total < buf.len() {
match writer.write(&buf[total..]).await {
Ok(0) => {
break;
}
Ok(n) => total += n,
Err(e) => return Err(e),
}
}
Ok(total)
}
/// Read exactly buf.len() bytes into buf, or return an error if EOF is reached before.
/// Like Go's io.ReadFull.
#[allow(dead_code)]
pub async fn read_full<R: AsyncRead + Send + Sync + Unpin>(mut reader: R, mut buf: &mut [u8]) -> std::io::Result<usize> {
let mut total = 0;
while !buf.is_empty() {
let n = match reader.read(buf).await {
Ok(n) => n,
Err(e) => {
if total == 0 {
return Err(e);
}
return Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
format!("read {} bytes, error: {}", total, e),
));
}
};
if n == 0 {
if total > 0 {
return Ok(total);
}
return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "early EOF"));
}
buf = &mut buf[n..];
total += n;
}
Ok(total)
}
/// Encodes a u64 into buf and returns the number of bytes written.
/// Panics if buf is too small.
pub fn put_uvarint(buf: &mut [u8], x: u64) -> usize {
let mut i = 0;
let mut x = x;
while x >= 0x80 {
buf[i] = (x as u8) | 0x80;
x >>= 7;
i += 1;
}
buf[i] = x as u8;
i + 1
}
pub fn put_uvarint_len(x: u64) -> usize {
let mut i = 0;
let mut x = x;
while x >= 0x80 {
x >>= 7;
i += 1;
}
i + 1
}
/// Decodes a u64 from buf and returns (value, number of bytes read).
/// If buf is too small, returns (0, 0).
/// If overflow, returns (0, -(n as isize)), where n is the number of bytes read.
pub fn uvarint(buf: &[u8]) -> (u64, isize) {
let mut x: u64 = 0;
let mut s: u32 = 0;
for (i, &b) in buf.iter().enumerate() {
if i == 10 {
// MaxVarintLen64 = 10
return (0, -((i + 1) as isize));
}
if b < 0x80 {
if i == 9 && b > 1 {
return (0, -((i + 1) as isize));
}
return (x | ((b as u64) << s), (i + 1) as isize);
}
x |= ((b & 0x7F) as u64) << s;
s += 7;
}
(0, 0)
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::BufReader;
#[tokio::test]
async fn test_read_full_exact() {
// let data = b"abcdef";
let data = b"channel async callback test data!";
let mut reader = BufReader::new(&data[..]);
let size = data.len();
let mut total = 0;
let mut rev = vec![0u8; size];
let mut count = 0;
while total < size {
let mut buf = [0u8; 8];
let n = read_full(&mut reader, &mut buf).await.unwrap();
total += n;
rev[total - n..total].copy_from_slice(&buf[..n]);
count += 1;
println!("count: {}, total: {}, n: {}", count, total, n);
}
assert_eq!(total, size);
assert_eq!(&rev, data);
}
#[tokio::test]
async fn test_read_full_short() {
let data = b"abc";
let mut reader = BufReader::new(&data[..]);
let mut buf = [0u8; 6];
let n = read_full(&mut reader, &mut buf).await.unwrap();
assert_eq!(n, 3);
assert_eq!(&buf[..n], data);
}
#[tokio::test]
async fn test_read_full_1m() {
let size = 1024 * 1024;
let data = vec![42u8; size];
let mut reader = BufReader::new(&data[..]);
let mut buf = vec![0u8; size / 3];
read_full(&mut reader, &mut buf).await.unwrap();
assert_eq!(buf, data[..size / 3]);
}
#[test]
fn test_put_uvarint_and_uvarint_zero() {
let mut buf = [0u8; 16];
let n = put_uvarint(&mut buf, 0);
let (decoded, m) = uvarint(&buf[..n]);
assert_eq!(decoded, 0);
assert_eq!(m as usize, n);
}
#[test]
fn test_put_uvarint_and_uvarint_max() {
let mut buf = [0u8; 16];
let n = put_uvarint(&mut buf, u64::MAX);
let (decoded, m) = uvarint(&buf[..n]);
assert_eq!(decoded, u64::MAX);
assert_eq!(m as usize, n);
}
#[test]
fn test_put_uvarint_and_uvarint_various() {
let mut buf = [0u8; 16];
for &v in &[1u64, 127, 128, 255, 300, 16384, u32::MAX as u64] {
let n = put_uvarint(&mut buf, v);
let (decoded, m) = uvarint(&buf[..n]);
assert_eq!(decoded, v, "decode mismatch for {}", v);
assert_eq!(m as usize, n, "length mismatch for {}", v);
}
}
#[test]
fn test_uvarint_incomplete() {
let buf = [0x80u8, 0x80, 0x80];
let (v, n) = uvarint(&buf);
assert_eq!(v, 0);
assert_eq!(n, 0);
}
#[test]
fn test_uvarint_overflow_case() {
let buf = [0xFFu8; 11];
let (v, n) = uvarint(&buf);
assert_eq!(v, 0);
assert!(n < 0);
}
#[tokio::test]
async fn test_write_all_basic() {
let data = b"hello world!";
let mut buf = Vec::new();
let n = write_all(&mut buf, data).await.unwrap();
assert_eq!(n, data.len());
assert_eq!(&buf, data);
}
#[tokio::test]
async fn test_write_all_partial() {
struct PartialWriter {
inner: Vec<u8>,
max_write: usize,
}
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::AsyncWrite;
impl AsyncWrite for PartialWriter {
fn poll_write(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll<std::io::Result<usize>> {
let n = buf.len().min(self.max_write);
self.inner.extend_from_slice(&buf[..n]);
Poll::Ready(Ok(n))
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Poll::Ready(Ok(()))
}
}
let data = b"abcdefghijklmnopqrstuvwxyz";
let mut writer = PartialWriter {
inner: Vec::new(),
max_write: 5,
};
let n = write_all(&mut writer, data).await.unwrap();
assert_eq!(n, data.len());
assert_eq!(&writer.inner, data);
}
}
+36 -3
View File
@@ -1,11 +1,44 @@
#[cfg(feature = "tls")]
mod certs;
pub mod certs;
#[cfg(feature = "ip")]
mod ip;
pub mod ip;
#[cfg(feature = "net")]
mod net;
pub mod net;
#[cfg(feature = "net")]
pub use net::*;
#[cfg(feature = "io")]
pub mod io;
#[cfg(feature = "hash")]
pub mod hash;
#[cfg(feature = "os")]
pub mod os;
#[cfg(feature = "path")]
pub mod path;
#[cfg(feature = "string")]
pub mod string;
#[cfg(feature = "crypto")]
pub mod crypto;
#[cfg(feature = "compress")]
pub mod compress;
#[cfg(feature = "tls")]
pub use certs::*;
#[cfg(feature = "hash")]
pub use hash::*;
#[cfg(feature = "io")]
pub use io::*;
#[cfg(feature = "ip")]
pub use ip::*;
#[cfg(feature = "crypto")]
pub use crypto::*;
#[cfg(feature = "compress")]
pub use compress::*;
+498
View File
@@ -1 +1,499 @@
use lazy_static::lazy_static;
use std::{
collections::HashSet,
fmt::Display,
net::{IpAddr, Ipv6Addr, SocketAddr, TcpListener, ToSocketAddrs},
};
use url::Host;
lazy_static! {
static ref LOCAL_IPS: Vec<IpAddr> = must_get_local_ips().unwrap();
}
/// helper for validating if the provided arg is an ip address.
pub fn is_socket_addr(addr: &str) -> bool {
// TODO IPv6 zone information?
addr.parse::<SocketAddr>().is_ok() || addr.parse::<IpAddr>().is_ok()
}
/// checks if server_addr is valid and local host.
pub fn check_local_server_addr(server_addr: &str) -> std::io::Result<SocketAddr> {
let addr: Vec<SocketAddr> = match server_addr.to_socket_addrs() {
Ok(addr) => addr.collect(),
Err(err) => return Err(std::io::Error::other(err)),
};
// 0.0.0.0 is a wildcard address and refers to local network
// addresses. I.e, 0.0.0.0:9000 like ":9000" refers to port
// 9000 on localhost.
for a in addr {
if a.ip().is_unspecified() {
return Ok(a);
}
let host = match a {
SocketAddr::V4(a) => Host::<&str>::Ipv4(*a.ip()),
SocketAddr::V6(a) => Host::Ipv6(*a.ip()),
};
if is_local_host(host, 0, 0)? {
return Ok(a);
}
}
Err(std::io::Error::other("host in server address should be this server"))
}
/// checks if the given parameter correspond to one of
/// the local IP of the current machine
pub fn is_local_host(host: Host<&str>, port: u16, local_port: u16) -> std::io::Result<bool> {
let local_set: HashSet<IpAddr> = LOCAL_IPS.iter().copied().collect();
let is_local_host = match host {
Host::Domain(domain) => {
let ips = match (domain, 0).to_socket_addrs().map(|v| v.map(|v| v.ip()).collect::<Vec<_>>()) {
Ok(ips) => ips,
Err(err) => return Err(std::io::Error::other(err)),
};
ips.iter().any(|ip| local_set.contains(ip))
}
Host::Ipv4(ip) => local_set.contains(&IpAddr::V4(ip)),
Host::Ipv6(ip) => local_set.contains(&IpAddr::V6(ip)),
};
if port > 0 {
return Ok(is_local_host && port == local_port);
}
Ok(is_local_host)
}
/// returns IP address of given host.
pub fn get_host_ip(host: Host<&str>) -> std::io::Result<HashSet<IpAddr>> {
match host {
Host::Domain(domain) => match (domain, 0)
.to_socket_addrs()
.map(|v| v.map(|v| v.ip()).collect::<HashSet<_>>())
{
Ok(ips) => Ok(ips),
Err(err) => Err(std::io::Error::other(err)),
},
Host::Ipv4(ip) => {
let mut set = HashSet::with_capacity(1);
set.insert(IpAddr::V4(ip));
Ok(set)
}
Host::Ipv6(ip) => {
let mut set = HashSet::with_capacity(1);
set.insert(IpAddr::V6(ip));
Ok(set)
}
}
}
pub fn get_available_port() -> u16 {
TcpListener::bind("0.0.0.0:0").unwrap().local_addr().unwrap().port()
}
/// returns IPs of local interface
pub fn must_get_local_ips() -> std::io::Result<Vec<IpAddr>> {
match netif::up() {
Ok(up) => Ok(up.map(|x| x.address().to_owned()).collect()),
Err(err) => Err(std::io::Error::other(format!("Unable to get IP addresses of this host: {}", err))),
}
}
#[derive(Debug, Clone)]
pub struct XHost {
pub name: String,
pub port: u16,
pub is_port_set: bool,
}
impl Display for XHost {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if !self.is_port_set {
write!(f, "{}", self.name)
} else if self.name.contains(':') {
write!(f, "[{}]:{}", self.name, self.port)
} else {
write!(f, "{}:{}", self.name, self.port)
}
}
}
impl TryFrom<String> for XHost {
type Error = std::io::Error;
fn try_from(value: String) -> std::result::Result<Self, Self::Error> {
if let Some(addr) = value.to_socket_addrs()?.next() {
Ok(Self {
name: addr.ip().to_string(),
port: addr.port(),
is_port_set: addr.port() > 0,
})
} else {
Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "value invalid"))
}
}
}
/// parses the address string, process the ":port" format for double-stack binding,
/// and resolve the host name or IP address. If the port is 0, an available port is assigned.
pub fn parse_and_resolve_address(addr_str: &str) -> std::io::Result<SocketAddr> {
let resolved_addr: SocketAddr = if let Some(port) = addr_str.strip_prefix(":") {
// Process the ":port" format for double stack binding
let port_str = port;
let port: u16 = port_str
.parse()
.map_err(|e| std::io::Error::other(format!("Invalid port format: {}, err:{:?}", addr_str, e)))?;
let final_port = if port == 0 {
get_available_port() // assume get_available_port is available here
} else {
port
};
// Using IPv6 without address specified [::], it should handle both IPv4 and IPv6
SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), final_port)
} else {
// Use existing logic to handle regular address formats
let mut addr = check_local_server_addr(addr_str)?; // assume check_local_server_addr is available here
if addr.port() == 0 {
addr.set_port(get_available_port());
}
addr
};
Ok(resolved_addr)
}
#[cfg(test)]
mod test {
use std::net::{Ipv4Addr, Ipv6Addr};
use super::*;
#[test]
fn test_is_socket_addr() {
let test_cases = [
// Valid IP addresses
("192.168.1.0", true),
("127.0.0.1", true),
("10.0.0.1", true),
("0.0.0.0", true),
("255.255.255.255", true),
// Valid IPv6 addresses
("2001:db8::1", true),
("::1", true),
("::", true),
("fe80::1", true),
// Valid socket addresses
("192.168.1.0:8080", true),
("127.0.0.1:9000", true),
("[2001:db8::1]:9000", true),
("[::1]:8080", true),
("0.0.0.0:0", true),
// Invalid addresses
("localhost", false),
("localhost:9000", false),
("example.com", false),
("example.com:8080", false),
("http://192.168.1.0", false),
("http://192.168.1.0:9000", false),
("256.256.256.256", false),
("192.168.1", false),
("192.168.1.0.1", false),
("", false),
(":", false),
(":::", false),
("invalid_ip", false),
];
for (addr, expected) in test_cases {
let result = is_socket_addr(addr);
assert_eq!(expected, result, "addr: '{}', expected: {}, got: {}", addr, expected, result);
}
}
#[test]
fn test_check_local_server_addr() {
// Test valid local addresses
let valid_cases = ["localhost:54321", "127.0.0.1:9000", "0.0.0.0:9000", "[::1]:8080", "::1:8080"];
for addr in valid_cases {
let result = check_local_server_addr(addr);
assert!(result.is_ok(), "Expected '{}' to be valid, but got error: {:?}", addr, result);
}
// Test invalid addresses
let invalid_cases = [
("localhost", "invalid socket address"),
("", "invalid socket address"),
("example.org:54321", "host in server address should be this server"),
("8.8.8.8:53", "host in server address should be this server"),
(":-10", "invalid port value"),
("invalid:port", "invalid port value"),
];
for (addr, expected_error_pattern) in invalid_cases {
let result = check_local_server_addr(addr);
assert!(result.is_err(), "Expected '{}' to be invalid, but it was accepted: {:?}", addr, result);
let error_msg = result.unwrap_err().to_string();
assert!(
error_msg.contains(expected_error_pattern) || error_msg.contains("invalid socket address"),
"Error message '{}' doesn't contain expected pattern '{}' for address '{}'",
error_msg,
expected_error_pattern,
addr
);
}
}
#[test]
fn test_is_local_host() {
// Test localhost domain
let localhost_host = Host::Domain("localhost");
assert!(is_local_host(localhost_host, 0, 0).unwrap());
// Test loopback IP addresses
let ipv4_loopback = Host::Ipv4(Ipv4Addr::new(127, 0, 0, 1));
assert!(is_local_host(ipv4_loopback, 0, 0).unwrap());
let ipv6_loopback = Host::Ipv6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1));
assert!(is_local_host(ipv6_loopback, 0, 0).unwrap());
// Test port matching
let localhost_with_port1 = Host::Domain("localhost");
assert!(is_local_host(localhost_with_port1, 8080, 8080).unwrap());
let localhost_with_port2 = Host::Domain("localhost");
assert!(!is_local_host(localhost_with_port2, 8080, 9000).unwrap());
// Test non-local host
let external_host = Host::Ipv4(Ipv4Addr::new(8, 8, 8, 8));
assert!(!is_local_host(external_host, 0, 0).unwrap());
// Test invalid domain should return error
let invalid_host = Host::Domain("invalid.nonexistent.domain.example");
assert!(is_local_host(invalid_host, 0, 0).is_err());
}
#[test]
fn test_get_host_ip() {
// Test IPv4 address
let ipv4_host = Host::Ipv4(Ipv4Addr::new(192, 168, 1, 1));
let ipv4_result = get_host_ip(ipv4_host).unwrap();
assert_eq!(ipv4_result.len(), 1);
assert!(ipv4_result.contains(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1))));
// Test IPv6 address
let ipv6_host = Host::Ipv6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1));
let ipv6_result = get_host_ip(ipv6_host).unwrap();
assert_eq!(ipv6_result.len(), 1);
assert!(ipv6_result.contains(&IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1))));
// Test localhost domain
let localhost_host = Host::Domain("localhost");
let localhost_result = get_host_ip(localhost_host).unwrap();
assert!(!localhost_result.is_empty());
// Should contain at least loopback address
assert!(
localhost_result.contains(&IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)))
|| localhost_result.contains(&IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)))
);
// Test invalid domain
let invalid_host = Host::Domain("invalid.nonexistent.domain.example");
assert!(get_host_ip(invalid_host).is_err());
}
#[test]
fn test_get_available_port() {
let port1 = get_available_port();
let port2 = get_available_port();
// Port should be in valid range (u16 max is always <= 65535)
assert!(port1 > 0);
assert!(port2 > 0);
// Different calls should typically return different ports
assert_ne!(port1, port2);
}
#[test]
fn test_must_get_local_ips() {
let local_ips = must_get_local_ips().unwrap();
let local_set: HashSet<IpAddr> = local_ips.into_iter().collect();
// Should contain loopback addresses
assert!(local_set.contains(&IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))));
// Should not be empty
assert!(!local_set.is_empty());
// All IPs should be valid
for ip in &local_set {
match ip {
IpAddr::V4(_) | IpAddr::V6(_) => {} // Valid
}
}
}
#[test]
fn test_xhost_display() {
// Test without port
let host_no_port = XHost {
name: "example.com".to_string(),
port: 0,
is_port_set: false,
};
assert_eq!(host_no_port.to_string(), "example.com");
// Test with port (IPv4-like name)
let host_with_port = XHost {
name: "192.168.1.1".to_string(),
port: 8080,
is_port_set: true,
};
assert_eq!(host_with_port.to_string(), "192.168.1.1:8080");
// Test with port (IPv6-like name)
let host_ipv6_with_port = XHost {
name: "2001:db8::1".to_string(),
port: 9000,
is_port_set: true,
};
assert_eq!(host_ipv6_with_port.to_string(), "[2001:db8::1]:9000");
// Test domain name with port
let host_domain_with_port = XHost {
name: "example.com".to_string(),
port: 443,
is_port_set: true,
};
assert_eq!(host_domain_with_port.to_string(), "example.com:443");
}
#[test]
fn test_xhost_try_from() {
// Test valid IPv4 address with port
let result = XHost::try_from("192.168.1.1:8080".to_string()).unwrap();
assert_eq!(result.name, "192.168.1.1");
assert_eq!(result.port, 8080);
assert!(result.is_port_set);
// Test valid IPv4 address without port
let result = XHost::try_from("192.168.1.1:0".to_string()).unwrap();
assert_eq!(result.name, "192.168.1.1");
assert_eq!(result.port, 0);
assert!(!result.is_port_set);
// Test valid IPv6 address with port
let result = XHost::try_from("[2001:db8::1]:9000".to_string()).unwrap();
assert_eq!(result.name, "2001:db8::1");
assert_eq!(result.port, 9000);
assert!(result.is_port_set);
// Test localhost with port (localhost may resolve to either IPv4 or IPv6)
let result = XHost::try_from("localhost:3000".to_string()).unwrap();
// localhost can resolve to either 127.0.0.1 or ::1 depending on system configuration
assert!(result.name == "127.0.0.1" || result.name == "::1");
assert_eq!(result.port, 3000);
assert!(result.is_port_set);
// Test invalid format
let result = XHost::try_from("invalid_format".to_string());
assert!(result.is_err());
// Test empty string
let result = XHost::try_from("".to_string());
assert!(result.is_err());
}
#[test]
fn test_parse_and_resolve_address() {
// Test port-only format
let result = parse_and_resolve_address(":8080").unwrap();
assert_eq!(result.ip(), IpAddr::V6(Ipv6Addr::UNSPECIFIED));
assert_eq!(result.port(), 8080);
// Test port-only format with port 0 (should get available port)
let result = parse_and_resolve_address(":0").unwrap();
assert_eq!(result.ip(), IpAddr::V6(Ipv6Addr::UNSPECIFIED));
assert!(result.port() > 0);
// Test localhost with port
let result = parse_and_resolve_address("localhost:9000").unwrap();
assert_eq!(result.port(), 9000);
// Test localhost with port 0 (should get available port)
let result = parse_and_resolve_address("localhost:0").unwrap();
assert!(result.port() > 0);
// Test 0.0.0.0 with port
let result = parse_and_resolve_address("0.0.0.0:7000").unwrap();
assert_eq!(result.ip(), IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)));
assert_eq!(result.port(), 7000);
// Test invalid port format
let result = parse_and_resolve_address(":invalid_port");
assert!(result.is_err());
// Test invalid address
let result = parse_and_resolve_address("example.org:8080");
assert!(result.is_err());
}
#[test]
fn test_edge_cases() {
// Test empty string for is_socket_addr
assert!(!is_socket_addr(""));
// Test single colon for is_socket_addr
assert!(!is_socket_addr(":"));
// Test malformed IPv6 for is_socket_addr
assert!(!is_socket_addr("[::]"));
assert!(!is_socket_addr("[::1"));
// Test very long strings
let long_string = "a".repeat(1000);
assert!(!is_socket_addr(&long_string));
// Test unicode characters
assert!(!is_socket_addr("测试.example.com"));
// Test special characters
assert!(!is_socket_addr("test@example.com:8080"));
assert!(!is_socket_addr("http://example.com:8080"));
}
#[test]
fn test_boundary_values() {
// Test port boundaries
assert!(is_socket_addr("127.0.0.1:0"));
assert!(is_socket_addr("127.0.0.1:65535"));
assert!(!is_socket_addr("127.0.0.1:65536"));
// Test IPv4 boundaries
assert!(is_socket_addr("0.0.0.0"));
assert!(is_socket_addr("255.255.255.255"));
assert!(!is_socket_addr("256.0.0.0"));
assert!(!is_socket_addr("0.0.0.256"));
// Test XHost with boundary ports
let host_max_port = XHost {
name: "example.com".to_string(),
port: 65535,
is_port_set: true,
};
assert_eq!(host_max_port.to_string(), "example.com:65535");
let host_zero_port = XHost {
name: "example.com".to_string(),
port: 0,
is_port_set: true,
};
assert_eq!(host_zero_port.to_string(), "example.com:0");
}
}
+176
View File
@@ -0,0 +1,176 @@
use nix::sys::stat::{self, stat};
use nix::sys::statfs::{self, FsType, statfs};
use std::fs::File;
use std::io::{self, BufRead, Error, ErrorKind};
use std::path::Path;
use super::{DiskInfo, IOStats};
/// Returns total and free bytes available in a directory, e.g. `/`.
pub fn get_info(p: impl AsRef<Path>) -> std::io::Result<DiskInfo> {
let stat_fs = statfs(p.as_ref())?;
let bsize = stat_fs.block_size() as u64;
let bfree = stat_fs.blocks_free() as u64;
let bavail = stat_fs.blocks_available() as u64;
let blocks = stat_fs.blocks() as u64;
let reserved = match bfree.checked_sub(bavail) {
Some(reserved) => reserved,
None => {
return Err(Error::other(format!(
"detected f_bavail space ({}) > f_bfree space ({}), fs corruption at ({}). please run 'fsck'",
bavail,
bfree,
p.as_ref().display()
)));
}
};
let total = match blocks.checked_sub(reserved) {
Some(total) => total * bsize,
None => {
return Err(Error::other(format!(
"detected reserved space ({}) > blocks space ({}), fs corruption at ({}). please run 'fsck'",
reserved,
blocks,
p.as_ref().display()
)));
}
};
let free = bavail * bsize;
let used = match total.checked_sub(free) {
Some(used) => used,
None => {
return Err(Error::other(format!(
"detected free space ({}) > total drive space ({}), fs corruption at ({}). please run 'fsck'",
free,
total,
p.as_ref().display()
)));
}
};
let st = stat(p.as_ref())?;
Ok(DiskInfo {
total,
free,
used,
files: stat_fs.files(),
ffree: stat_fs.files_free(),
fstype: get_fs_type(stat_fs.filesystem_type()).to_string(),
major: stat::major(st.st_dev),
minor: stat::minor(st.st_dev),
..Default::default()
})
}
/// Returns the filesystem type of the underlying mounted filesystem
///
/// TODO The following mapping could not find the corresponding constant in `nix`:
///
/// "137d" => "EXT",
/// "4244" => "HFS",
/// "5346544e" => "NTFS",
/// "61756673" => "AUFS",
/// "ef51" => "EXT2OLD",
/// "2fc12fc1" => "zfs",
/// "ff534d42" => "cifs",
/// "53464846" => "wslfs",
fn get_fs_type(fs_type: FsType) -> &'static str {
match fs_type {
statfs::TMPFS_MAGIC => "TMPFS",
statfs::MSDOS_SUPER_MAGIC => "MSDOS",
// statfs::XFS_SUPER_MAGIC => "XFS",
statfs::NFS_SUPER_MAGIC => "NFS",
statfs::EXT4_SUPER_MAGIC => "EXT4",
statfs::ECRYPTFS_SUPER_MAGIC => "ecryptfs",
statfs::OVERLAYFS_SUPER_MAGIC => "overlayfs",
statfs::REISERFS_SUPER_MAGIC => "REISERFS",
_ => "UNKNOWN",
}
}
pub fn same_disk(disk1: &str, disk2: &str) -> std::io::Result<bool> {
let stat1 = stat(disk1)?;
let stat2 = stat(disk2)?;
Ok(stat1.st_dev == stat2.st_dev)
}
pub fn get_drive_stats(major: u32, minor: u32) -> std::io::Result<IOStats> {
read_drive_stats(&format!("/sys/dev/block/{}:{}/stat", major, minor))
}
fn read_drive_stats(stats_file: &str) -> std::io::Result<IOStats> {
let stats = read_stat(stats_file)?;
if stats.len() < 11 {
return Err(Error::new(
ErrorKind::InvalidData,
format!("found invalid format while reading {}", stats_file),
));
}
let mut io_stats = IOStats {
read_ios: stats[0],
read_merges: stats[1],
read_sectors: stats[2],
read_ticks: stats[3],
write_ios: stats[4],
write_merges: stats[5],
write_sectors: stats[6],
write_ticks: stats[7],
current_ios: stats[8],
total_ticks: stats[9],
req_ticks: stats[10],
..Default::default()
};
if stats.len() > 14 {
io_stats.discard_ios = stats[11];
io_stats.discard_merges = stats[12];
io_stats.discard_sectors = stats[13];
io_stats.discard_ticks = stats[14];
}
Ok(io_stats)
}
fn read_stat(file_name: &str) -> std::io::Result<Vec<u64>> {
// Open file
let path = Path::new(file_name);
let file = File::open(path)?;
// Create a BufReader
let reader = io::BufReader::new(file);
// Read first line
let mut stats = Vec::new();
if let Some(line) = reader.lines().next() {
let line = line?;
// Split line and parse as u64
// https://rust-lang.github.io/rust-clippy/master/index.html#trim_split_whitespace
for token in line.split_whitespace() {
let ui64: u64 = token
.parse()
.map_err(|e| Error::new(ErrorKind::InvalidData, format!("failed to parse '{}' as u64: {}", token, e)))?;
stats.push(ui64);
}
}
Ok(stats)
}
#[cfg(test)]
mod test {
use super::get_drive_stats;
#[ignore] // FIXME: failed in github actions
#[test]
fn test_stats() {
let major = 7;
let minor = 11;
let s = get_drive_stats(major, minor).unwrap();
println!("{:?}", s);
}
}
+111
View File
@@ -0,0 +1,111 @@
#[cfg(target_os = "linux")]
mod linux;
#[cfg(all(unix, not(target_os = "linux")))]
mod unix;
#[cfg(target_os = "windows")]
mod windows;
#[cfg(target_os = "linux")]
pub use linux::{get_drive_stats, get_info, same_disk};
// pub use linux::same_disk;
#[cfg(all(unix, not(target_os = "linux")))]
pub use unix::{get_drive_stats, get_info, same_disk};
#[cfg(target_os = "windows")]
pub use windows::{get_drive_stats, get_info, same_disk};
#[derive(Debug, Default, PartialEq)]
pub struct IOStats {
pub read_ios: u64,
pub read_merges: u64,
pub read_sectors: u64,
pub read_ticks: u64,
pub write_ios: u64,
pub write_merges: u64,
pub write_sectors: u64,
pub write_ticks: u64,
pub current_ios: u64,
pub total_ticks: u64,
pub req_ticks: u64,
pub discard_ios: u64,
pub discard_merges: u64,
pub discard_sectors: u64,
pub discard_ticks: u64,
pub flush_ios: u64,
pub flush_ticks: u64,
}
#[derive(Debug, Default, PartialEq)]
pub struct DiskInfo {
pub total: u64,
pub free: u64,
pub used: u64,
pub files: u64,
pub ffree: u64,
pub fstype: String,
pub major: u64,
pub minor: u64,
pub name: String,
pub rotational: bool,
pub nrrequests: u64,
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
#[test]
fn test_get_info_valid_path() {
let temp_dir = tempfile::tempdir().unwrap();
let info = get_info(temp_dir.path()).unwrap();
println!("Disk Info: {:?}", info);
assert!(info.total > 0);
assert!(info.free > 0);
assert!(info.used > 0);
assert!(info.files > 0);
assert!(info.ffree > 0);
assert!(!info.fstype.is_empty());
}
#[test]
fn test_get_info_invalid_path() {
let invalid_path = PathBuf::from("/invalid/path");
let result = get_info(&invalid_path);
assert!(result.is_err());
}
#[test]
fn test_same_disk_same_path() {
let temp_dir = tempfile::tempdir().unwrap();
let path = temp_dir.path().to_str().unwrap();
let result = same_disk(path, path).unwrap();
assert!(result);
}
#[test]
fn test_same_disk_different_paths() {
let temp_dir1 = tempfile::tempdir().unwrap();
let temp_dir2 = tempfile::tempdir().unwrap();
let path1 = temp_dir1.path().to_str().unwrap();
let path2 = temp_dir2.path().to_str().unwrap();
let result = same_disk(path1, path2).unwrap();
// Since both temporary directories are created in the same file system,
// they should be on the same disk in most cases
println!("Path1: {}, Path2: {}, Same disk: {}", path1, path2, result);
// Test passes if the function doesn't panic - the actual result depends on test environment
}
#[ignore] // FIXME: failed in github actions
#[test]
fn test_get_drive_stats_default() {
let stats = get_drive_stats(0, 0).unwrap();
assert_eq!(stats, IOStats::default());
}
}
+72
View File
@@ -0,0 +1,72 @@
use super::{DiskInfo, IOStats};
use nix::sys::{stat::stat, statfs::statfs};
use std::io::Error;
use std::path::Path;
/// Returns total and free bytes available in a directory, e.g. `/`.
pub fn get_info(p: impl AsRef<Path>) -> std::io::Result<DiskInfo> {
let stat = statfs(p.as_ref())?;
let bsize = stat.block_size() as u64;
let bfree = stat.blocks_free();
let bavail = stat.blocks_available();
let blocks = stat.blocks();
let reserved = match bfree.checked_sub(bavail) {
Some(reserved) => reserved,
None => {
return Err(Error::other(format!(
"detected f_bavail space ({}) > f_bfree space ({}), fs corruption at ({}). please run 'fsck'",
bavail,
bfree,
p.as_ref().display()
)));
}
};
let total = match blocks.checked_sub(reserved) {
Some(total) => total * bsize,
None => {
return Err(Error::other(format!(
"detected reserved space ({}) > blocks space ({}), fs corruption at ({}). please run 'fsck'",
reserved,
blocks,
p.as_ref().display()
)));
}
};
let free = bavail * bsize;
let used = match total.checked_sub(free) {
Some(used) => used,
None => {
return Err(Error::other(format!(
"detected free space ({}) > total drive space ({}), fs corruption at ({}). please run 'fsck'",
free,
total,
p.as_ref().display()
)));
}
};
Ok(DiskInfo {
total,
free,
used,
files: stat.files(),
ffree: stat.files_free(),
fstype: stat.filesystem_type_name().to_string(),
..Default::default()
})
}
pub fn same_disk(disk1: &str, disk2: &str) -> std::io::Result<bool> {
let stat1 = stat(disk1)?;
let stat2 = stat(disk2)?;
Ok(stat1.st_dev == stat2.st_dev)
}
pub fn get_drive_stats(_major: u32, _minor: u32) -> std::io::Result<IOStats> {
Ok(IOStats::default())
}
+142
View File
@@ -0,0 +1,142 @@
#![allow(unsafe_code)] // TODO: audit unsafe code
use super::{DiskInfo, IOStats};
use std::io::{Error, ErrorKind};
use std::mem;
use std::os::windows::ffi::OsStrExt;
use std::path::Path;
use winapi::shared::minwindef::{DWORD, MAX_PATH};
use winapi::shared::ntdef::ULARGE_INTEGER;
use winapi::um::fileapi::{GetDiskFreeSpaceExW, GetDiskFreeSpaceW, GetVolumeInformationW, GetVolumePathNameW};
use winapi::um::winnt::{LPCWSTR, WCHAR};
/// Returns total and free bytes available in a directory, e.g. `C:\`.
pub fn get_info(p: impl AsRef<Path>) -> std::io::Result<DiskInfo> {
let path_wide: Vec<WCHAR> = p
.as_ref()
.canonicalize()?
.into_os_string()
.encode_wide()
.chain(std::iter::once(0)) // Null-terminate the string
.collect();
let mut lp_free_bytes_available: ULARGE_INTEGER = unsafe { mem::zeroed() };
let mut lp_total_number_of_bytes: ULARGE_INTEGER = unsafe { mem::zeroed() };
let mut lp_total_number_of_free_bytes: ULARGE_INTEGER = unsafe { mem::zeroed() };
let success = unsafe {
GetDiskFreeSpaceExW(
path_wide.as_ptr(),
&mut lp_free_bytes_available,
&mut lp_total_number_of_bytes,
&mut lp_total_number_of_free_bytes,
)
};
if success == 0 {
return Err(Error::last_os_error());
}
let total = unsafe { *lp_total_number_of_bytes.QuadPart() };
let free = unsafe { *lp_total_number_of_free_bytes.QuadPart() };
if free > total {
return Err(Error::new(
ErrorKind::Other,
format!(
"detected free space ({}) > total drive space ({}), fs corruption at ({}). please run 'fsck'",
free,
total,
p.as_ref().display()
),
));
}
let mut lp_sectors_per_cluster: DWORD = 0;
let mut lp_bytes_per_sector: DWORD = 0;
let mut lp_number_of_free_clusters: DWORD = 0;
let mut lp_total_number_of_clusters: DWORD = 0;
let success = unsafe {
GetDiskFreeSpaceW(
path_wide.as_ptr(),
&mut lp_sectors_per_cluster,
&mut lp_bytes_per_sector,
&mut lp_number_of_free_clusters,
&mut lp_total_number_of_clusters,
)
};
if success == 0 {
return Err(Error::last_os_error());
}
Ok(DiskInfo {
total,
free,
used: total - free,
files: lp_total_number_of_clusters as u64,
ffree: lp_number_of_free_clusters as u64,
fstype: get_fs_type(&path_wide)?,
..Default::default()
})
}
/// Returns leading volume name.
fn get_volume_name(v: &[WCHAR]) -> std::io::Result<LPCWSTR> {
let volume_name_size: DWORD = MAX_PATH as _;
let mut lp_volume_name_buffer: [WCHAR; MAX_PATH] = [0; MAX_PATH];
let success = unsafe { GetVolumePathNameW(v.as_ptr(), lp_volume_name_buffer.as_mut_ptr(), volume_name_size) };
if success == 0 {
return Err(Error::last_os_error());
}
Ok(lp_volume_name_buffer.as_ptr())
}
fn utf16_to_string(v: &[WCHAR]) -> String {
let len = v.iter().position(|&x| x == 0).unwrap_or(v.len());
String::from_utf16_lossy(&v[..len])
}
/// Returns the filesystem type of the underlying mounted filesystem
fn get_fs_type(p: &[WCHAR]) -> std::io::Result<String> {
let path = get_volume_name(p)?;
let volume_name_size: DWORD = MAX_PATH as _;
let n_file_system_name_size: DWORD = MAX_PATH as _;
let mut lp_volume_serial_number: DWORD = 0;
let mut lp_maximum_component_length: DWORD = 0;
let mut lp_file_system_flags: DWORD = 0;
let mut lp_volume_name_buffer: [WCHAR; MAX_PATH] = [0; MAX_PATH];
let mut lp_file_system_name_buffer: [WCHAR; MAX_PATH] = [0; MAX_PATH];
let success = unsafe {
GetVolumeInformationW(
path,
lp_volume_name_buffer.as_mut_ptr(),
volume_name_size,
&mut lp_volume_serial_number,
&mut lp_maximum_component_length,
&mut lp_file_system_flags,
lp_file_system_name_buffer.as_mut_ptr(),
n_file_system_name_size,
)
};
if success == 0 {
return Err(Error::last_os_error());
}
Ok(utf16_to_string(&lp_file_system_name_buffer))
}
pub fn same_disk(_disk1: &str, _disk2: &str) -> std::io::Result<bool> {
Ok(false)
}
pub fn get_drive_stats(_major: u32, _minor: u32) -> std::io::Result<IOStats> {
Ok(IOStats::default())
}
+308
View File
@@ -0,0 +1,308 @@
use std::path::Path;
use std::path::PathBuf;
pub const GLOBAL_DIR_SUFFIX: &str = "__XLDIR__";
pub const SLASH_SEPARATOR: &str = "/";
pub const GLOBAL_DIR_SUFFIX_WITH_SLASH: &str = "__XLDIR__/";
pub fn has_suffix(s: &str, suffix: &str) -> bool {
if cfg!(target_os = "windows") {
s.to_lowercase().ends_with(&suffix.to_lowercase())
} else {
s.ends_with(suffix)
}
}
pub fn encode_dir_object(object: &str) -> String {
if has_suffix(object, SLASH_SEPARATOR) {
format!("{}{}", object.trim_end_matches(SLASH_SEPARATOR), GLOBAL_DIR_SUFFIX)
} else {
object.to_string()
}
}
pub fn is_dir_object(object: &str) -> bool {
let obj = encode_dir_object(object);
obj.ends_with(GLOBAL_DIR_SUFFIX)
}
#[allow(dead_code)]
pub fn decode_dir_object(object: &str) -> String {
if has_suffix(object, GLOBAL_DIR_SUFFIX) {
format!("{}{}", object.trim_end_matches(GLOBAL_DIR_SUFFIX), SLASH_SEPARATOR)
} else {
object.to_string()
}
}
pub fn retain_slash(s: &str) -> String {
if s.is_empty() {
return s.to_string();
}
if s.ends_with(SLASH_SEPARATOR) {
s.to_string()
} else {
format!("{}{}", s, SLASH_SEPARATOR)
}
}
pub fn strings_has_prefix_fold(s: &str, prefix: &str) -> bool {
s.len() >= prefix.len() && (s[..prefix.len()] == *prefix || s[..prefix.len()].eq_ignore_ascii_case(prefix))
}
pub fn has_prefix(s: &str, prefix: &str) -> bool {
if cfg!(target_os = "windows") {
return strings_has_prefix_fold(s, prefix);
}
s.starts_with(prefix)
}
pub fn path_join(elem: &[PathBuf]) -> PathBuf {
let mut joined_path = PathBuf::new();
for path in elem {
joined_path.push(path);
}
joined_path
}
pub fn path_join_buf(elements: &[&str]) -> String {
let trailing_slash = !elements.is_empty() && elements.last().unwrap().ends_with(SLASH_SEPARATOR);
let mut dst = String::new();
let mut added = 0;
for e in elements {
if added > 0 || !e.is_empty() {
if added > 0 {
dst.push_str(SLASH_SEPARATOR);
}
dst.push_str(e);
added += e.len();
}
}
let result = dst.to_string();
let cpath = Path::new(&result).components().collect::<PathBuf>();
let clean_path = cpath.to_string_lossy();
if trailing_slash {
return format!("{}{}", clean_path, SLASH_SEPARATOR);
}
clean_path.to_string()
}
pub fn path_to_bucket_object_with_base_path(bash_path: &str, path: &str) -> (String, String) {
let path = path.trim_start_matches(bash_path).trim_start_matches(SLASH_SEPARATOR);
if let Some(m) = path.find(SLASH_SEPARATOR) {
return (path[..m].to_string(), path[m + SLASH_SEPARATOR.len()..].to_string());
}
(path.to_string(), "".to_string())
}
pub fn path_to_bucket_object(s: &str) -> (String, String) {
path_to_bucket_object_with_base_path("", s)
}
pub fn base_dir_from_prefix(prefix: &str) -> String {
let mut base_dir = dir(prefix).to_owned();
if base_dir == "." || base_dir == "./" || base_dir == "/" {
base_dir = "".to_owned();
}
if !prefix.contains('/') {
base_dir = "".to_owned();
}
if !base_dir.is_empty() && !base_dir.ends_with(SLASH_SEPARATOR) {
base_dir.push_str(SLASH_SEPARATOR);
}
base_dir
}
pub struct LazyBuf {
s: String,
buf: Option<Vec<u8>>,
w: usize,
}
impl LazyBuf {
pub fn new(s: String) -> Self {
LazyBuf { s, buf: None, w: 0 }
}
pub fn index(&self, i: usize) -> u8 {
if let Some(ref buf) = self.buf {
buf[i]
} else {
self.s.as_bytes()[i]
}
}
pub fn append(&mut self, c: u8) {
if self.buf.is_none() {
if self.w < self.s.len() && self.s.as_bytes()[self.w] == c {
self.w += 1;
return;
}
let mut new_buf = vec![0; self.s.len()];
new_buf[..self.w].copy_from_slice(&self.s.as_bytes()[..self.w]);
self.buf = Some(new_buf);
}
if let Some(ref mut buf) = self.buf {
buf[self.w] = c;
self.w += 1;
}
}
pub fn string(&self) -> String {
if let Some(ref buf) = self.buf {
String::from_utf8(buf[..self.w].to_vec()).unwrap()
} else {
self.s[..self.w].to_string()
}
}
}
pub fn clean(path: &str) -> String {
if path.is_empty() {
return ".".to_string();
}
let rooted = path.starts_with('/');
let n = path.len();
let mut out = LazyBuf::new(path.to_string());
let mut r = 0;
let mut dotdot = 0;
if rooted {
out.append(b'/');
r = 1;
dotdot = 1;
}
while r < n {
match path.as_bytes()[r] {
b'/' => {
// Empty path element
r += 1;
}
b'.' if r + 1 == n || path.as_bytes()[r + 1] == b'/' => {
// . element
r += 1;
}
b'.' if path.as_bytes()[r + 1] == b'.' && (r + 2 == n || path.as_bytes()[r + 2] == b'/') => {
// .. element: remove to last /
r += 2;
if out.w > dotdot {
// Can backtrack
out.w -= 1;
while out.w > dotdot && out.index(out.w) != b'/' {
out.w -= 1;
}
} else if !rooted {
// Cannot backtrack but not rooted, so append .. element.
if out.w > 0 {
out.append(b'/');
}
out.append(b'.');
out.append(b'.');
dotdot = out.w;
}
}
_ => {
// Real path element.
// Add slash if needed
if (rooted && out.w != 1) || (!rooted && out.w != 0) {
out.append(b'/');
}
// Copy element
while r < n && path.as_bytes()[r] != b'/' {
out.append(path.as_bytes()[r]);
r += 1;
}
}
}
}
// Turn empty string into "."
if out.w == 0 {
return ".".to_string();
}
out.string()
}
pub fn split(path: &str) -> (&str, &str) {
// Find the last occurrence of the '/' character
if let Some(i) = path.rfind('/') {
// Return the directory (up to and including the last '/') and the file name
return (&path[..i + 1], &path[i + 1..]);
}
// If no '/' is found, return an empty string for the directory and the whole path as the file name
(path, "")
}
pub fn dir(path: &str) -> String {
let (a, _) = split(path);
clean(a)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_base_dir_from_prefix() {
let a = "da/";
println!("---- in {}", a);
let a = base_dir_from_prefix(a);
println!("---- out {}", a);
}
#[test]
fn test_clean() {
assert_eq!(clean(""), ".");
assert_eq!(clean("abc"), "abc");
assert_eq!(clean("abc/def"), "abc/def");
assert_eq!(clean("a/b/c"), "a/b/c");
assert_eq!(clean("."), ".");
assert_eq!(clean(".."), "..");
assert_eq!(clean("../.."), "../..");
assert_eq!(clean("../../abc"), "../../abc");
assert_eq!(clean("/abc"), "/abc");
assert_eq!(clean("/"), "/");
assert_eq!(clean("abc/"), "abc");
assert_eq!(clean("abc/def/"), "abc/def");
assert_eq!(clean("a/b/c/"), "a/b/c");
assert_eq!(clean("./"), ".");
assert_eq!(clean("../"), "..");
assert_eq!(clean("../../"), "../..");
assert_eq!(clean("/abc/"), "/abc");
assert_eq!(clean("abc//def//ghi"), "abc/def/ghi");
assert_eq!(clean("//abc"), "/abc");
assert_eq!(clean("///abc"), "/abc");
assert_eq!(clean("//abc//"), "/abc");
assert_eq!(clean("abc//"), "abc");
assert_eq!(clean("abc/./def"), "abc/def");
assert_eq!(clean("/./abc/def"), "/abc/def");
assert_eq!(clean("abc/."), "abc");
assert_eq!(clean("abc/./../def"), "def");
assert_eq!(clean("abc//./../def"), "def");
assert_eq!(clean("abc/../../././../def"), "../../def");
assert_eq!(clean("abc/def/ghi/../jkl"), "abc/def/jkl");
assert_eq!(clean("abc/def/../ghi/../jkl"), "abc/jkl");
assert_eq!(clean("abc/def/.."), "abc");
assert_eq!(clean("abc/def/../.."), ".");
assert_eq!(clean("/abc/def/../.."), "/");
assert_eq!(clean("abc/def/../../.."), "..");
assert_eq!(clean("/abc/def/../../.."), "/");
assert_eq!(clean("abc/def/../../../ghi/jkl/../../../mno"), "../../mno");
}
}
+666
View File
@@ -0,0 +1,666 @@
use lazy_static::*;
use regex::Regex;
use std::io::{Error, Result};
pub fn parse_bool(str: &str) -> Result<bool> {
match str {
"1" | "t" | "T" | "true" | "TRUE" | "True" | "on" | "ON" | "On" | "enabled" => Ok(true),
"0" | "f" | "F" | "false" | "FALSE" | "False" | "off" | "OFF" | "Off" | "disabled" => Ok(false),
_ => Err(Error::other(format!("ParseBool: parsing {}", str))),
}
}
pub fn match_simple(pattern: &str, name: &str) -> bool {
if pattern.is_empty() {
return name == pattern;
}
if pattern == "*" {
return true;
}
// Do an extended wildcard '*' and '?' match.
deep_match_rune(name.as_bytes(), pattern.as_bytes(), true)
}
pub fn match_pattern(pattern: &str, name: &str) -> bool {
if pattern.is_empty() {
return name == pattern;
}
if pattern == "*" {
return true;
}
// Do an extended wildcard '*' and '?' match.
deep_match_rune(name.as_bytes(), pattern.as_bytes(), false)
}
pub fn has_pattern(patterns: &[&str], match_str: &str) -> bool {
for pattern in patterns {
if match_simple(pattern, match_str) {
return true;
}
}
false
}
pub fn has_string_suffix_in_slice(str: &str, list: &[&str]) -> bool {
let str = str.to_lowercase();
for v in list {
if *v == "*" {
return true;
}
if str.ends_with(&v.to_lowercase()) {
return true;
}
}
false
}
fn deep_match_rune(str_: &[u8], pattern: &[u8], simple: bool) -> bool {
let (mut str_, mut pattern) = (str_, pattern);
while !pattern.is_empty() {
match pattern[0] as char {
'*' => {
return if pattern.len() == 1 {
true
} else {
deep_match_rune(str_, &pattern[1..], simple)
|| (!str_.is_empty() && deep_match_rune(&str_[1..], pattern, simple))
};
}
'?' => {
if str_.is_empty() {
return simple;
}
}
_ => {
if str_.is_empty() || str_[0] != pattern[0] {
return false;
}
}
}
str_ = &str_[1..];
pattern = &pattern[1..];
}
str_.is_empty() && pattern.is_empty()
}
pub fn match_as_pattern_prefix(pattern: &str, text: &str) -> bool {
let mut i = 0;
while i < text.len() && i < pattern.len() {
match pattern.as_bytes()[i] as char {
'*' => return true,
'?' => i += 1,
_ => {
if pattern.as_bytes()[i] != text.as_bytes()[i] {
return false;
}
}
}
i += 1;
}
text.len() <= pattern.len()
}
lazy_static! {
static ref ELLIPSES_RE: Regex = Regex::new(r"(.*)(\{[0-9a-z]*\.\.\.[0-9a-z]*\})(.*)").unwrap();
}
/// Ellipses constants
const OPEN_BRACES: &str = "{";
const CLOSE_BRACES: &str = "}";
const ELLIPSES: &str = "...";
/// ellipses pattern, describes the range and also the
/// associated prefix and suffixes.
#[derive(Debug, Default, PartialEq, Eq)]
pub struct Pattern {
pub prefix: String,
pub suffix: String,
pub seq: Vec<String>,
}
impl Pattern {
/// expands a ellipses pattern.
pub fn expand(&self) -> Vec<String> {
let mut ret = Vec::with_capacity(self.suffix.len());
for v in self.seq.iter() {
match (self.prefix.is_empty(), self.suffix.is_empty()) {
(false, true) => ret.push(format!("{}{}", self.prefix, v)),
(true, false) => ret.push(format!("{}{}", v, self.suffix)),
(true, true) => ret.push(v.to_string()),
(false, false) => ret.push(format!("{}{}{}", self.prefix, v, self.suffix)),
}
}
ret
}
pub fn len(&self) -> usize {
self.seq.len()
}
pub fn is_empty(&self) -> bool {
self.seq.is_empty()
}
}
/// contains a list of patterns provided in the input.
#[derive(Debug, PartialEq, Eq)]
pub struct ArgPattern {
inner: Vec<Pattern>,
}
impl AsRef<Vec<Pattern>> for ArgPattern {
fn as_ref(&self) -> &Vec<Pattern> {
&self.inner
}
}
impl AsMut<Vec<Pattern>> for ArgPattern {
fn as_mut(&mut self) -> &mut Vec<Pattern> {
&mut self.inner
}
}
impl ArgPattern {
pub fn new(inner: Vec<Pattern>) -> Self {
Self { inner }
}
/// expands all the ellipses patterns in the given argument.
pub fn expand(&self) -> Vec<Vec<String>> {
let ret: Vec<Vec<String>> = self.inner.iter().map(|v| v.expand()).collect();
Self::arg_expander(&ret)
}
/// recursively expands labels into its respective forms.
fn arg_expander(lbs: &[Vec<String>]) -> Vec<Vec<String>> {
if lbs.len() == 1 {
return lbs[0].iter().map(|v| vec![v.to_string()]).collect();
}
let mut ret = Vec::new();
let (first, others) = lbs.split_at(1);
for bs in first[0].iter() {
let ots = Self::arg_expander(others);
for mut obs in ots {
obs.push(bs.to_string());
ret.push(obs);
}
}
ret
}
/// returns the total number of sizes in the given patterns.
pub fn total_sizes(&self) -> usize {
self.inner.iter().fold(1, |acc, v| acc * v.seq.len())
}
}
/// finds all ellipses patterns, recursively and parses the ranges numerically.
pub fn find_ellipses_patterns(arg: &str) -> Result<ArgPattern> {
let mut parts = match ELLIPSES_RE.captures(arg) {
Some(caps) => caps,
None => {
return Err(Error::other(format!(
"Invalid ellipsis format in ({}), Ellipsis range must be provided in format {{N...M}} where N and M are positive integers, M must be greater than N, with an allowed minimum range of 4",
arg
)));
}
};
let mut patterns = Vec::new();
while let Some(prefix) = parts.get(1) {
let seq = parse_ellipses_range(parts[2].into())?;
match ELLIPSES_RE.captures(prefix.into()) {
Some(cs) => {
patterns.push(Pattern {
seq,
prefix: String::new(),
suffix: parts[3].into(),
});
parts = cs;
}
None => {
patterns.push(Pattern {
seq,
prefix: prefix.as_str().to_owned(),
suffix: parts[3].into(),
});
break;
}
};
}
// Check if any of the prefix or suffixes now have flower braces
// left over, in such a case we generally think that there is
// perhaps a typo in users input and error out accordingly.
for p in patterns.iter() {
if p.prefix.contains(OPEN_BRACES)
|| p.prefix.contains(CLOSE_BRACES)
|| p.suffix.contains(OPEN_BRACES)
|| p.suffix.contains(CLOSE_BRACES)
{
return Err(Error::other(format!(
"Invalid ellipsis format in ({}), Ellipsis range must be provided in format {{N...M}} where N and M are positive integers, M must be greater than N, with an allowed minimum range of 4",
arg
)));
}
}
Ok(ArgPattern::new(patterns))
}
/// returns true if input arg has ellipses type pattern.
pub fn has_ellipses<T: AsRef<str>>(s: &[T]) -> bool {
let pattern = [ELLIPSES, OPEN_BRACES, CLOSE_BRACES];
s.iter().any(|v| pattern.iter().any(|p| v.as_ref().contains(p)))
}
/// Parses an ellipses range pattern of following style
///
/// example:
/// {1...64}
/// {33...64}
pub fn parse_ellipses_range(pattern: &str) -> Result<Vec<String>> {
if !pattern.contains(OPEN_BRACES) {
return Err(Error::other("Invalid argument"));
}
if !pattern.contains(CLOSE_BRACES) {
return Err(Error::other("Invalid argument"));
}
let ellipses_range: Vec<&str> = pattern
.trim_start_matches(OPEN_BRACES)
.trim_end_matches(CLOSE_BRACES)
.split(ELLIPSES)
.collect();
if ellipses_range.len() != 2 {
return Err(Error::other("Invalid argument"));
}
// TODO: Add support for hexadecimals.
let start = ellipses_range[0].parse::<usize>().map_err(Error::other)?;
let end = ellipses_range[1].parse::<usize>().map_err(Error::other)?;
if start > end {
return Err(Error::other("Invalid argument:range start cannot be bigger than end"));
}
let mut ret: Vec<String> = Vec::with_capacity(end - start + 1);
for i in start..=end {
if ellipses_range[0].starts_with('0') && ellipses_range[0].len() > 1 {
ret.push(format!("{:0width$}", i, width = ellipses_range[1].len()));
} else {
ret.push(format!("{}", i));
}
}
Ok(ret)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_has_ellipses() {
// Tests for all args without ellipses.
let test_cases = [
(1, vec!["64"], false),
// Found flower braces, still attempt to parse and throw an error.
(2, vec!["{1..64}"], true),
(3, vec!["{1..2..}"], true),
// Test for valid input.
(4, vec!["1...64"], true),
(5, vec!["{1...2O}"], true),
(6, vec!["..."], true),
(7, vec!["{-1...1}"], true),
(8, vec!["{0...-1}"], true),
(9, vec!["{1....4}"], true),
(10, vec!["{1...64}"], true),
(11, vec!["{...}"], true),
(12, vec!["{1...64}", "{65...128}"], true),
(13, vec!["http://rustfs{2...3}/export/set{1...64}"], true),
(
14,
vec![
"http://rustfs{2...3}/export/set{1...64}",
"http://rustfs{2...3}/export/set{65...128}",
],
true,
),
(15, vec!["mydisk-{a...z}{1...20}"], true),
(16, vec!["mydisk-{1...4}{1..2.}"], true),
];
for (i, args, expected) in test_cases {
let ret = has_ellipses(&args);
assert_eq!(ret, expected, "Test{}: Expected {}, got {}", i, expected, ret);
}
}
#[test]
fn test_find_ellipses_patterns() {
#[derive(Default)]
struct TestCase<'a> {
num: usize,
pattern: &'a str,
success: bool,
want: Vec<Vec<&'a str>>,
}
let test_cases = [
TestCase {
num: 1,
pattern: "{1..64}",
..Default::default()
},
TestCase {
num: 2,
pattern: "1...64",
..Default::default()
},
TestCase {
num: 2,
pattern: "...",
..Default::default()
},
TestCase {
num: 3,
pattern: "{1...",
..Default::default()
},
TestCase {
num: 4,
pattern: "...64}",
..Default::default()
},
TestCase {
num: 5,
pattern: "{...}",
..Default::default()
},
TestCase {
num: 6,
pattern: "{-1...1}",
..Default::default()
},
TestCase {
num: 7,
pattern: "{0...-1}",
..Default::default()
},
TestCase {
num: 8,
pattern: "{1...2O}",
..Default::default()
},
TestCase {
num: 9,
pattern: "{64...1}",
..Default::default()
},
TestCase {
num: 10,
pattern: "{1....4}",
..Default::default()
},
TestCase {
num: 11,
pattern: "mydisk-{a...z}{1...20}",
..Default::default()
},
TestCase {
num: 12,
pattern: "mydisk-{1...4}{1..2.}",
..Default::default()
},
TestCase {
num: 13,
pattern: "{1..2.}-mydisk-{1...4}",
..Default::default()
},
TestCase {
num: 14,
pattern: "{{1...4}}",
..Default::default()
},
TestCase {
num: 16,
pattern: "{4...02}",
..Default::default()
},
TestCase {
num: 17,
pattern: "{f...z}",
..Default::default()
},
// Test for valid input.
TestCase {
num: 18,
pattern: "{1...64}",
success: true,
want: vec![
vec!["1"],
vec!["2"],
vec!["3"],
vec!["4"],
vec!["5"],
vec!["6"],
vec!["7"],
vec!["8"],
vec!["9"],
vec!["10"],
vec!["11"],
vec!["12"],
vec!["13"],
vec!["14"],
vec!["15"],
vec!["16"],
vec!["17"],
vec!["18"],
vec!["19"],
vec!["20"],
vec!["21"],
vec!["22"],
vec!["23"],
vec!["24"],
vec!["25"],
vec!["26"],
vec!["27"],
vec!["28"],
vec!["29"],
vec!["30"],
vec!["31"],
vec!["32"],
vec!["33"],
vec!["34"],
vec!["35"],
vec!["36"],
vec!["37"],
vec!["38"],
vec!["39"],
vec!["40"],
vec!["41"],
vec!["42"],
vec!["43"],
vec!["44"],
vec!["45"],
vec!["46"],
vec!["47"],
vec!["48"],
vec!["49"],
vec!["50"],
vec!["51"],
vec!["52"],
vec!["53"],
vec!["54"],
vec!["55"],
vec!["56"],
vec!["57"],
vec!["58"],
vec!["59"],
vec!["60"],
vec!["61"],
vec!["62"],
vec!["63"],
vec!["64"],
],
},
TestCase {
num: 19,
pattern: "{1...5} {65...70}",
success: true,
want: vec![
vec!["1 ", "65"],
vec!["2 ", "65"],
vec!["3 ", "65"],
vec!["4 ", "65"],
vec!["5 ", "65"],
vec!["1 ", "66"],
vec!["2 ", "66"],
vec!["3 ", "66"],
vec!["4 ", "66"],
vec!["5 ", "66"],
vec!["1 ", "67"],
vec!["2 ", "67"],
vec!["3 ", "67"],
vec!["4 ", "67"],
vec!["5 ", "67"],
vec!["1 ", "68"],
vec!["2 ", "68"],
vec!["3 ", "68"],
vec!["4 ", "68"],
vec!["5 ", "68"],
vec!["1 ", "69"],
vec!["2 ", "69"],
vec!["3 ", "69"],
vec!["4 ", "69"],
vec!["5 ", "69"],
vec!["1 ", "70"],
vec!["2 ", "70"],
vec!["3 ", "70"],
vec!["4 ", "70"],
vec!["5 ", "70"],
],
},
TestCase {
num: 20,
pattern: "{01...036}",
success: true,
want: vec![
vec!["001"],
vec!["002"],
vec!["003"],
vec!["004"],
vec!["005"],
vec!["006"],
vec!["007"],
vec!["008"],
vec!["009"],
vec!["010"],
vec!["011"],
vec!["012"],
vec!["013"],
vec!["014"],
vec!["015"],
vec!["016"],
vec!["017"],
vec!["018"],
vec!["019"],
vec!["020"],
vec!["021"],
vec!["022"],
vec!["023"],
vec!["024"],
vec!["025"],
vec!["026"],
vec!["027"],
vec!["028"],
vec!["029"],
vec!["030"],
vec!["031"],
vec!["032"],
vec!["033"],
vec!["034"],
vec!["035"],
vec!["036"],
],
},
TestCase {
num: 21,
pattern: "{001...036}",
success: true,
want: vec![
vec!["001"],
vec!["002"],
vec!["003"],
vec!["004"],
vec!["005"],
vec!["006"],
vec!["007"],
vec!["008"],
vec!["009"],
vec!["010"],
vec!["011"],
vec!["012"],
vec!["013"],
vec!["014"],
vec!["015"],
vec!["016"],
vec!["017"],
vec!["018"],
vec!["019"],
vec!["020"],
vec!["021"],
vec!["022"],
vec!["023"],
vec!["024"],
vec!["025"],
vec!["026"],
vec!["027"],
vec!["028"],
vec!["029"],
vec!["030"],
vec!["031"],
vec!["032"],
vec!["033"],
vec!["034"],
vec!["035"],
vec!["036"],
],
},
];
for test_case in test_cases {
let ret = find_ellipses_patterns(test_case.pattern);
match ret {
Ok(v) => {
if !test_case.success {
panic!("Test{}: Expected failure but passed instead", test_case.num);
}
let got = v.expand();
if got.len() != test_case.want.len() {
panic!("Test{}: Expected {}, got {}", test_case.num, test_case.want.len(), got.len());
}
assert_eq!(got, test_case.want, "Test{}: Expected {:?}, got {:?}", test_case.num, test_case.want, got);
}
Err(e) => {
if test_case.success {
panic!("Test{}: Expected success but failed instead {:?}", test_case.num, e);
}
}
}
}
}
}
+4 -4
View File
@@ -608,8 +608,8 @@ mod tests {
#[tokio::test]
async fn test_decompress_with_invalid_format() {
// Test decompression with invalid format
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
let sample_content = b"Hello, compression world!";
let cursor = Cursor::new(sample_content);
@@ -634,8 +634,8 @@ mod tests {
#[tokio::test]
async fn test_decompress_with_zip_format() {
// Test decompression with Zip format (currently not supported)
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
let sample_content = b"Hello, compression world!";
let cursor = Cursor::new(sample_content);
@@ -660,8 +660,8 @@ mod tests {
#[tokio::test]
async fn test_decompress_error_propagation() {
// Test error propagation during decompression process
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
let sample_content = b"Hello, compression world!";
let cursor = Cursor::new(sample_content);
@@ -690,8 +690,8 @@ mod tests {
#[tokio::test]
async fn test_decompress_callback_execution() {
// Test callback function execution during decompression
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
let sample_content = b"Hello, compression world!";
let cursor = Cursor::new(sample_content);