mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-06 05:17:42 +00:00
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:
@@ -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()),
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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);
|
||||
@@ -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(¬_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"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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";
|
||||
@@ -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::*;
|
||||
@@ -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, ¶ms.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);
|
||||
}
|
||||
}
|
||||
@@ -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,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,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};
|
||||
|
||||
@@ -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,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,5 +1,5 @@
|
||||
use crate::sinks::Sink;
|
||||
use crate::UnifiedLogEntry;
|
||||
use crate::sinks::Sink;
|
||||
use async_trait::async_trait;
|
||||
|
||||
/// Webhook Sink Implementation
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,4 +1,4 @@
|
||||
use crate::{sinks::Sink, UnifiedLogEntry};
|
||||
use crate::{UnifiedLogEntry, sinks::Sink};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc::Receiver;
|
||||
|
||||
|
||||
@@ -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"
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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!");
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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");
|
||||
// }
|
||||
// }
|
||||
@@ -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 {}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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> {}
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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 × {
|
||||
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!(); // 添加空行分隔不同大小的结果
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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::*;
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user