Compare commits

...

4 Commits

Author SHA1 Message Date
Warren 25991c71b2 Update Cargo.lock for new dependencies (xattr, smb-server AAPL modules)
Test / test (push) Has been cancelled
Test / build (push) Has been cancelled
2026-06-22 14:22:16 +08:00
Warren 866d0536c8 Add SMB AAPL Extensions Phase 1-6 + VFS xattr support
Phase 1: AAPL Create Context negotiation
Phase 2: AFP_AfpInfo Stream structure (Finder info + creation time)
Phase 2.5: SMB Named Stream Backend (NamedStreamPath)
Phase 2.6: Backend Named Stream Support in handlers
Phase 2.7: VFS Extended Attributes (get/set/remove/list_xattr)
Phase 4: Time Machine share config (time_machine field)
Phase 5: Server/Volume Capabilities
Phase 6: macOS Unicode mapping (private range ↔ ASCII)

Tests: 174 smb-server tests pass, 52 VFS tests pass
2026-06-22 14:21:53 +08:00
Warren 64709ec529 Add CTDB Phase 1-5: TDB storage + Node management + Control protocol + IP manager + Recovery 2026-06-22 14:21:39 +08:00
Warren a8d81f2a9c Revert "Remove Download Center routes from server.rs (dead code cleanup)"
This reverts commit 20b208bb7f.
2026-06-22 14:12:14 +08:00
22 changed files with 3271 additions and 5 deletions
Generated
+1
View File
@@ -3062,6 +3062,7 @@ dependencies = [
"url", "url",
"uuid", "uuid",
"x25519-dalek", "x25519-dalek",
"xattr",
"xmltree", "xmltree",
"xz2", "xz2",
"zip", "zip",
+1
View File
@@ -74,6 +74,7 @@ ureq = "2.12" # 輕量同步 HTTP 客戶端
reqwest = { version = "0.12", optional = true } # Async HTTP client for AsyncS3Vfs reqwest = { version = "0.12", optional = true } # Async HTTP client for AsyncS3Vfs
rayon = "1.10" # Phase 4: 并行加密 rayon = "1.10" # Phase 4: 并行加密
url = "2" # URL 解析(rusty-s3 依賴) url = "2" # URL 解析(rusty-s3 依賴)
xattr = "1.0" # Extended attributes support (AFP_AfpInfo, Time Machine)
# === SMB/CIFS Client (Phase 1) === # === SMB/CIFS Client (Phase 1) ===
smb2 = { path = "../vendor/smb2" } # Pure-Rust SMB2/3 client library with pipelined I/O smb2 = { path = "../vendor/smb2" } # Pure-Rust SMB2/3 client library with pipelined I/O
+379
View File
@@ -0,0 +1,379 @@
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::RwLock;
use super::node::NodeId;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct PublicIpId(pub u32);
#[derive(Debug, Clone)]
pub struct PublicIpEntry {
pub id: PublicIpId,
pub ip: IpAddr,
pub interface: String,
pub owner: Option<NodeId>,
pub previous_owner: Option<NodeId>,
}
impl PublicIpEntry {
pub fn new(id: PublicIpId, ip: IpAddr, interface: &str) -> Self {
Self {
id,
ip,
interface: interface.to_string(),
owner: None,
previous_owner: None,
}
}
}
pub struct IpManager {
ip_pool: RwLock<Vec<PublicIpEntry>>,
assignments: RwLock<HashMap<PublicIpId, NodeId>>,
next_id: RwLock<u32>,
}
impl IpManager {
pub fn new() -> Self {
Self {
ip_pool: RwLock::new(Vec::new()),
assignments: RwLock::new(HashMap::new()),
next_id: RwLock::new(0),
}
}
pub fn add_ip(&self, ip: IpAddr, interface: &str) -> PublicIpId {
let mut pool = self.ip_pool.write().unwrap();
let mut next = self.next_id.write().unwrap();
let id = PublicIpId(*next);
*next += 1;
pool.push(PublicIpEntry::new(id, ip, interface));
id
}
pub fn remove_ip(&self, id: PublicIpId) {
let mut pool = self.ip_pool.write().unwrap();
pool.retain(|e| e.id != id);
let mut assignments = self.assignments.write().unwrap();
assignments.remove(&id);
}
pub fn assign_ip(&self, id: PublicIpId, node: NodeId) -> Result<(), String> {
let exists = {
let pool = self.ip_pool.read().unwrap();
pool.iter().any(|e| e.id == id)
};
if !exists {
return Err(format!("IP {:?} not found", id));
}
let prev_owner = {
let assignments = self.assignments.read().unwrap();
assignments.get(&id).copied()
};
{
let mut assignments = self.assignments.write().unwrap();
assignments.insert(id, node);
}
{
let mut pool = self.ip_pool.write().unwrap();
if let Some(entry) = pool.iter_mut().find(|e| e.id == id) {
if let Some(prev) = prev_owner {
entry.previous_owner = Some(prev);
}
entry.owner = Some(node);
}
}
Ok(())
}
pub fn release_ip(&self, id: PublicIpId) -> Result<(), String> {
let mut assignments = self.assignments.write().unwrap();
if let Some(prev) = assignments.remove(&id) {
let mut pool = self.ip_pool.write().unwrap();
if let Some(entry) = pool.iter_mut().find(|e| e.id == id) {
entry.previous_owner = Some(prev);
entry.owner = None;
}
Ok(())
} else {
Err(format!("IP {:?} not assigned", id))
}
}
pub fn reassign_on_failure(&self, failed_node: NodeId, active_nodes: &[NodeId]) -> Vec<(PublicIpId, NodeId)> {
let assignments = self.assignments.read().unwrap();
let failed_ips: Vec<PublicIpId> = assignments
.iter()
.filter(|(_, &node)| node == failed_node)
.map(|(&id, _)| id)
.collect();
drop(assignments);
let candidates: Vec<NodeId> = active_nodes
.iter()
.filter(|&&n| n != failed_node)
.copied()
.collect();
let mut reassignments = Vec::new();
for (i, ip_id) in failed_ips.iter().enumerate() {
if let Some(&new_owner) = candidates.get(i % candidates.len().max(1)) {
let _ = self.assign_ip(*ip_id, new_owner);
reassignments.push((*ip_id, new_owner));
}
}
reassignments
}
pub fn get_owner(&self, id: PublicIpId) -> Option<NodeId> {
self.assignments.read().unwrap().get(&id).copied()
}
pub fn owned_by(&self, node: NodeId) -> Vec<PublicIpEntry> {
let assignments = self.assignments.read().unwrap();
let pool = self.ip_pool.read().unwrap();
assignments
.iter()
.filter(|(_, &owner)| owner == node)
.filter_map(|(&id, _)| pool.iter().find(|e| e.id == id).cloned())
.collect()
}
pub fn unassigned(&self) -> Vec<PublicIpEntry> {
let assignments = self.assignments.read().unwrap();
let pool = self.ip_pool.read().unwrap();
pool.iter()
.filter(|e| !assignments.contains_key(&e.id))
.cloned()
.collect()
}
pub fn all_ips(&self) -> Vec<PublicIpEntry> {
self.ip_pool.read().unwrap().clone()
}
pub fn ip_count(&self) -> usize {
self.ip_pool.read().unwrap().len()
}
pub fn assigned_count(&self) -> usize {
self.assignments.read().unwrap().len()
}
pub fn balance(&self, active_nodes: &[NodeId]) -> Vec<(PublicIpId, NodeId)> {
let pool = self.ip_pool.read().unwrap();
let assignments = self.assignments.read().unwrap();
let total_ips = pool.len();
let node_count = active_nodes.len().max(1);
let per_node = total_ips / node_count;
let mut current_counts: HashMap<NodeId, usize> = HashMap::new();
for &node in active_nodes {
current_counts.insert(node, 0);
}
for (_, &owner) in assignments.iter() {
*current_counts.entry(owner).or_insert(0) += 1;
}
let mut reassignments = Vec::new();
for entry in pool.iter() {
let current_owner = assignments.get(&entry.id).copied();
let needs_reassign = match current_owner {
Some(owner) => {
!active_nodes.contains(&owner) || current_counts[&owner] > per_node
}
None => true,
};
if needs_reassign {
if let Some(&target) = active_nodes
.iter()
.min_by_key(|n| current_counts.get(n).copied().unwrap_or(0))
{
current_counts
.entry(target)
.and_modify(|c| *c += 1)
.or_insert(1);
reassignments.push((entry.id, target));
}
}
}
drop(pool);
drop(assignments);
for (id, node) in &reassignments {
let _ = self.assign_ip(*id, *node);
}
reassignments
}
}
impl Default for IpManager {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::Ipv4Addr;
fn ip(a: u8, b: u8, c: u8, d: u8) -> IpAddr {
IpAddr::V4(Ipv4Addr::new(a, b, c, d))
}
#[test]
fn test_add_ip() {
let mgr = IpManager::new();
let id = mgr.add_ip(ip(192, 168, 1, 100), "eth0");
assert_eq!(mgr.ip_count(), 1);
assert_eq!(id, PublicIpId(0));
}
#[test]
fn test_add_multiple_ips() {
let mgr = IpManager::new();
let id0 = mgr.add_ip(ip(192, 168, 1, 100), "eth0");
let id1 = mgr.add_ip(ip(192, 168, 1, 101), "eth0");
let id2 = mgr.add_ip(ip(192, 168, 1, 102), "eth1");
assert_eq!(mgr.ip_count(), 3);
assert_ne!(id0, id1);
assert_ne!(id1, id2);
}
#[test]
fn test_assign_ip() {
let mgr = IpManager::new();
let id = mgr.add_ip(ip(192, 168, 1, 100), "eth0");
mgr.assign_ip(id, NodeId(0)).unwrap();
assert_eq!(mgr.get_owner(id), Some(NodeId(0)));
assert_eq!(mgr.assigned_count(), 1);
}
#[test]
fn test_release_ip() {
let mgr = IpManager::new();
let id = mgr.add_ip(ip(192, 168, 1, 100), "eth0");
mgr.assign_ip(id, NodeId(0)).unwrap();
mgr.release_ip(id).unwrap();
assert_eq!(mgr.get_owner(id), None);
}
#[test]
fn test_release_unassigned_fails() {
let mgr = IpManager::new();
let id = mgr.add_ip(ip(192, 168, 1, 100), "eth0");
assert!(mgr.release_ip(id).is_err());
}
#[test]
fn test_assign_nonexistent_fails() {
let mgr = IpManager::new();
assert!(mgr.assign_ip(PublicIpId(999), NodeId(0)).is_err());
}
#[test]
fn test_owned_by() {
let mgr = IpManager::new();
let id0 = mgr.add_ip(ip(192, 168, 1, 100), "eth0");
let id1 = mgr.add_ip(ip(192, 168, 1, 101), "eth0");
let id2 = mgr.add_ip(ip(192, 168, 1, 102), "eth1");
mgr.assign_ip(id0, NodeId(0)).unwrap();
mgr.assign_ip(id1, NodeId(1)).unwrap();
mgr.assign_ip(id2, NodeId(0)).unwrap();
let node0_ips = mgr.owned_by(NodeId(0));
assert_eq!(node0_ips.len(), 2);
}
#[test]
fn test_unassigned() {
let mgr = IpManager::new();
mgr.add_ip(ip(192, 168, 1, 100), "eth0");
let id1 = mgr.add_ip(ip(192, 168, 1, 101), "eth0");
mgr.assign_ip(id1, NodeId(0)).unwrap();
let unassigned = mgr.unassigned();
assert_eq!(unassigned.len(), 1);
}
#[test]
fn test_reassign_on_failure() {
let mgr = IpManager::new();
let id0 = mgr.add_ip(ip(192, 168, 1, 100), "eth0");
let id1 = mgr.add_ip(ip(192, 168, 1, 101), "eth0");
mgr.assign_ip(id0, NodeId(0)).unwrap();
mgr.assign_ip(id1, NodeId(0)).unwrap();
let reassignments = mgr.reassign_on_failure(NodeId(0), &[NodeId(1), NodeId(2)]);
assert_eq!(reassignments.len(), 2);
let new_owners: Vec<NodeId> = reassignments.iter().map(|(_, n)| *n).collect();
assert!(new_owners.contains(&NodeId(1)) || new_owners.contains(&NodeId(2)));
}
#[test]
fn test_reassign_skips_failed_node() {
let mgr = IpManager::new();
let id0 = mgr.add_ip(ip(192, 168, 1, 100), "eth0");
mgr.assign_ip(id0, NodeId(0)).unwrap();
let reassignments = mgr.reassign_on_failure(NodeId(0), &[NodeId(0), NodeId(1)]);
assert_eq!(reassignments.len(), 1);
assert_ne!(reassignments[0].1, NodeId(0));
}
#[test]
fn test_remove_ip() {
let mgr = IpManager::new();
let id = mgr.add_ip(ip(192, 168, 1, 100), "eth0");
mgr.assign_ip(id, NodeId(0)).unwrap();
mgr.remove_ip(id);
assert_eq!(mgr.ip_count(), 0);
assert_eq!(mgr.assigned_count(), 0);
}
#[test]
fn test_balance() {
let mgr = IpManager::new();
for i in 100..104 {
mgr.add_ip(ip(192, 168, 1, i), "eth0");
}
let reassignments = mgr.balance(&[NodeId(0), NodeId(1)]);
assert!(!reassignments.is_empty());
let node0 = mgr.owned_by(NodeId(0));
let node1 = mgr.owned_by(NodeId(1));
assert_eq!(node0.len() + node1.len(), 4);
assert!((node0.len() as i32 - node1.len() as i32).abs() <= 1);
}
#[test]
fn test_balance_excludes_inactive_nodes() {
let mgr = IpManager::new();
for i in 100..103 {
mgr.add_ip(ip(192, 168, 1, i), "eth0");
}
mgr.assign_ip(PublicIpId(0), NodeId(5)).unwrap();
let reassignments = mgr.balance(&[NodeId(0), NodeId(1)]);
assert!(reassignments
.iter()
.any(|(_, n)| *n == NodeId(0) || *n == NodeId(1)));
assert!(!reassignments.iter().any(|(_, n)| *n == NodeId(5)));
}
#[test]
fn test_previous_owner() {
let mgr = IpManager::new();
let id = mgr.add_ip(ip(192, 168, 1, 100), "eth0");
mgr.assign_ip(id, NodeId(0)).unwrap();
mgr.assign_ip(id, NodeId(1)).unwrap();
let pool = mgr.all_ips();
let entry = pool.iter().find(|e| e.id == id).unwrap();
assert_eq!(entry.previous_owner, Some(NodeId(0)));
assert_eq!(entry.owner, Some(NodeId(1)));
}
}
+5
View File
@@ -0,0 +1,5 @@
pub mod tdb;
pub mod node;
pub mod protocol;
pub mod ip_manager;
pub mod recovery;
+353
View File
@@ -0,0 +1,353 @@
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct NodeId(pub u32);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NodeState {
Up,
Down,
Unhealthy,
Banned,
Disabled,
}
impl NodeState {
pub fn is_active(&self) -> bool {
matches!(self, NodeState::Up)
}
pub fn as_str(&self) -> &'static str {
match self {
NodeState::Up => "UP",
NodeState::Down => "DOWN",
NodeState::Unhealthy => "UNHEALTHY",
NodeState::Banned => "BANNED",
NodeState::Disabled => "DISABLED",
}
}
}
#[derive(Debug, Clone)]
pub struct NodeInfo {
pub id: NodeId,
pub addr: SocketAddr,
pub state: NodeState,
pub last_heartbeat: Option<Instant>,
pub public_ips: Vec<String>,
pub generation: u64,
}
impl NodeInfo {
pub fn new(id: NodeId, addr: SocketAddr) -> Self {
Self {
id,
addr,
state: NodeState::Down,
last_heartbeat: None,
public_ips: Vec::new(),
generation: 0,
}
}
pub fn is_alive(&self, timeout: Duration) -> bool {
match self.last_heartbeat {
Some(t) => t.elapsed() < timeout,
None => false,
}
}
}
#[derive(Debug, Clone)]
pub struct NodeMask {
nodes: Vec<bool>,
}
impl NodeMask {
pub fn new(size: usize) -> Self {
Self {
nodes: vec![false; size],
}
}
pub fn set(&mut self, id: NodeId, active: bool) {
if (id.0 as usize) < self.nodes.len() {
self.nodes[id.0 as usize] = active;
}
}
pub fn is_active(&self, id: NodeId) -> bool {
self.nodes.get(id.0 as usize).copied().unwrap_or(false)
}
pub fn active_nodes(&self) -> Vec<NodeId> {
self.nodes
.iter()
.enumerate()
.filter(|(_, &active)| active)
.map(|(i, _)| NodeId(i as u32))
.collect()
}
pub fn len(&self) -> usize {
self.nodes.len()
}
pub fn active_count(&self) -> usize {
self.nodes.iter().filter(|&&a| a).count()
}
}
pub struct NodeManager {
nodes: RwLock<HashMap<NodeId, NodeInfo>>,
self_id: NodeId,
heartbeat_timeout: Duration,
heartbeat_interval: Duration,
}
impl NodeManager {
pub fn new(self_id: NodeId, self_addr: SocketAddr) -> Self {
let mut nodes = HashMap::new();
nodes.insert(
self_id,
NodeInfo {
id: self_id,
addr: self_addr,
state: NodeState::Up,
last_heartbeat: Some(Instant::now()),
public_ips: Vec::new(),
generation: 0,
},
);
Self {
nodes: RwLock::new(nodes),
self_id,
heartbeat_timeout: Duration::from_secs(5),
heartbeat_interval: Duration::from_secs(1),
}
}
pub fn add_node(&self, id: NodeId, addr: SocketAddr) {
let mut nodes = self.nodes.write().unwrap();
nodes.insert(id, NodeInfo::new(id, addr));
}
pub fn remove_node(&self, id: NodeId) {
let mut nodes = self.nodes.write().unwrap();
nodes.remove(&id);
}
pub fn record_heartbeat(&self, id: NodeId) {
let mut nodes = self.nodes.write().unwrap();
if let Some(node) = nodes.get_mut(&id) {
node.last_heartbeat = Some(Instant::now());
if node.state == NodeState::Down {
node.state = NodeState::Up;
node.generation += 1;
}
}
}
pub fn set_node_state(&self, id: NodeId, state: NodeState) {
let mut nodes = self.nodes.write().unwrap();
if let Some(node) = nodes.get_mut(&id) {
node.state = state;
if state == NodeState::Up && node.last_heartbeat.is_none() {
node.last_heartbeat = Some(Instant::now());
}
node.generation += 1;
}
}
pub fn get_node(&self, id: NodeId) -> Option<NodeInfo> {
self.nodes.read().unwrap().get(&id).cloned()
}
pub fn all_nodes(&self) -> Vec<NodeInfo> {
self.nodes.read().unwrap().values().cloned().collect()
}
pub fn active_nodes(&self) -> Vec<NodeInfo> {
self.nodes
.read()
.unwrap()
.values()
.filter(|n| n.state == NodeState::Up)
.cloned()
.collect()
}
pub fn check_health(&self) -> Vec<(NodeId, NodeState)> {
let timeout = self.heartbeat_timeout;
let mut nodes = self.nodes.write().unwrap();
let mut transitions = Vec::new();
for (id, node) in nodes.iter_mut() {
if *id == self.self_id {
continue;
}
match &node.last_heartbeat {
Some(t) => {
if t.elapsed() > timeout && node.state == NodeState::Up {
node.state = NodeState::Down;
node.generation += 1;
transitions.push((*id, NodeState::Down));
}
}
None => {
if node.state == NodeState::Up {
node.state = NodeState::Down;
node.generation += 1;
transitions.push((*id, NodeState::Down));
}
}
}
}
transitions
}
pub fn self_id(&self) -> NodeId {
self.self_id
}
pub fn node_count(&self) -> usize {
self.nodes.read().unwrap().len()
}
pub fn build_nodemask(&self) -> NodeMask {
let nodes = self.nodes.read().unwrap();
let max_id = nodes.keys().map(|k| k.0).max().unwrap_or(0) as usize;
let mut mask = NodeMask::new(max_id + 1);
for (id, node) in nodes.iter() {
mask.set(*id, node.state.is_active());
}
mask
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::{Ipv4Addr, SocketAddrV4};
fn addr(port: u16) -> SocketAddr {
SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port))
}
#[test]
fn test_node_creation() {
let mgr = NodeManager::new(NodeId(0), addr(4000));
let self_node = mgr.get_node(NodeId(0)).unwrap();
assert_eq!(self_node.state, NodeState::Up);
assert!(self_node.last_heartbeat.is_some());
}
#[test]
fn test_add_remove_node() {
let mgr = NodeManager::new(NodeId(0), addr(4000));
mgr.add_node(NodeId(1), addr(4001));
assert_eq!(mgr.node_count(), 2);
mgr.remove_node(NodeId(1));
assert_eq!(mgr.node_count(), 1);
}
#[test]
fn test_heartbeat_updates() {
let mgr = NodeManager::new(NodeId(0), addr(4000));
mgr.add_node(NodeId(1), addr(4001));
assert_eq!(mgr.get_node(NodeId(1)).unwrap().state, NodeState::Down);
mgr.record_heartbeat(NodeId(1));
assert_eq!(mgr.get_node(NodeId(1)).unwrap().state, NodeState::Up);
}
#[test]
fn test_health_check_timeout() {
let mgr = NodeManager::new(NodeId(0), addr(4000));
mgr.add_node(NodeId(1), addr(4001));
mgr.record_heartbeat(NodeId(1));
assert_eq!(mgr.get_node(NodeId(1)).unwrap().state, NodeState::Up);
std::thread::sleep(Duration::from_millis(100));
let mgr = NodeManager {
nodes: RwLock::new(mgr.nodes.read().unwrap().clone()),
self_id: NodeId(0),
heartbeat_timeout: Duration::from_millis(50),
heartbeat_interval: Duration::from_millis(10),
};
let transitions = mgr.check_health();
assert!(transitions.contains(&(NodeId(1), NodeState::Down)));
}
#[test]
fn test_nodemask() {
let mut mask = NodeMask::new(5);
mask.set(NodeId(0), true);
mask.set(NodeId(2), true);
mask.set(NodeId(4), true);
assert!(mask.is_active(NodeId(0)));
assert!(!mask.is_active(NodeId(1)));
assert_eq!(mask.active_count(), 3);
assert_eq!(mask.len(), 5);
}
#[test]
fn test_build_nodemask() {
let mgr = NodeManager::new(NodeId(0), addr(4000));
mgr.add_node(NodeId(1), addr(4001));
mgr.add_node(NodeId(2), addr(4002));
mgr.record_heartbeat(NodeId(1));
let mask = mgr.build_nodemask();
assert!(mask.is_active(NodeId(0)));
assert!(mask.is_active(NodeId(1)));
assert!(!mask.is_active(NodeId(2)));
}
#[test]
fn test_node_state_string() {
assert_eq!(NodeState::Up.as_str(), "UP");
assert_eq!(NodeState::Down.as_str(), "DOWN");
assert_eq!(NodeState::Unhealthy.as_str(), "UNHEALTHY");
assert_eq!(NodeState::Banned.as_str(), "BANNED");
assert_eq!(NodeState::Disabled.as_str(), "DISABLED");
}
#[test]
fn test_set_node_state() {
let mgr = NodeManager::new(NodeId(0), addr(4000));
mgr.add_node(NodeId(1), addr(4001));
mgr.set_node_state(NodeId(1), NodeState::Banned);
assert_eq!(mgr.get_node(NodeId(1)).unwrap().state, NodeState::Banned);
}
#[test]
fn test_generation_increment() {
let mgr = NodeManager::new(NodeId(0), addr(4000));
mgr.add_node(NodeId(1), addr(4001));
let gen0 = mgr.get_node(NodeId(1)).unwrap().generation;
mgr.set_node_state(NodeId(1), NodeState::Down);
let gen1 = mgr.get_node(NodeId(1)).unwrap().generation;
assert!(gen1 > gen0);
}
#[test]
fn test_active_nodes() {
let mgr = NodeManager::new(NodeId(0), addr(4000));
mgr.add_node(NodeId(1), addr(4001));
mgr.add_node(NodeId(2), addr(4002));
mgr.record_heartbeat(NodeId(1));
let active = mgr.active_nodes();
assert_eq!(active.len(), 2);
}
#[test]
fn test_is_alive() {
let node = NodeInfo::new(NodeId(0), addr(4000));
assert!(!node.is_alive(Duration::from_secs(5)));
let mut node = node;
node.last_heartbeat = Some(Instant::now());
assert!(node.is_alive(Duration::from_secs(5)));
}
}
+515
View File
@@ -0,0 +1,515 @@
use byteorder::{BigEndian, LittleEndian, ReadBytesExt, WriteBytesExt};
use std::io::{self, Cursor, Read, Write};
use std::net::TcpStream;
pub const CTDB_MAGIC: u32 = 0x43544442;
pub const CTDB_VERSION: u32 = 1;
pub const CTDB_HEADER_SIZE: usize = 24;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u32)]
pub enum CtdbCommand {
Connect = 1,
Disconnect = 2,
Ping = 3,
Pong = 4,
GetDb = 10,
Fetch = 11,
Store = 12,
Delete = 13,
Keys = 14,
SetNodeMask = 20,
GetNodeMask = 21,
NodeStatus = 22,
TakeIp = 30,
ReleaseIp = 31,
Monitor = 40,
Recovery = 50,
RecoveryDone = 51,
Unknown = 0xFFFF,
}
impl CtdbCommand {
pub fn from_u32(v: u32) -> Self {
match v {
1 => CtdbCommand::Connect,
2 => CtdbCommand::Disconnect,
3 => CtdbCommand::Ping,
4 => CtdbCommand::Pong,
10 => CtdbCommand::GetDb,
11 => CtdbCommand::Fetch,
12 => CtdbCommand::Store,
13 => CtdbCommand::Delete,
14 => CtdbCommand::Keys,
20 => CtdbCommand::SetNodeMask,
21 => CtdbCommand::GetNodeMask,
22 => CtdbCommand::NodeStatus,
30 => CtdbCommand::TakeIp,
31 => CtdbCommand::ReleaseIp,
40 => CtdbCommand::Monitor,
50 => CtdbCommand::Recovery,
51 => CtdbCommand::RecoveryDone,
_ => CtdbCommand::Unknown,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CtdbStatus {
Success = 0,
Error = 1,
NotFound = 2,
Exists = 3,
Corrupt = 4,
Timeout = 5,
NotActive = 6,
}
impl CtdbStatus {
pub fn from_u32(v: u32) -> Self {
match v {
0 => CtdbStatus::Success,
2 => CtdbStatus::NotFound,
3 => CtdbStatus::Exists,
4 => CtdbStatus::Corrupt,
5 => CtdbStatus::Timeout,
6 => CtdbStatus::NotActive,
_ => CtdbStatus::Error,
}
}
}
#[derive(Debug, Clone)]
pub struct CtdbHeader {
pub magic: u32,
pub version: u32,
pub command: u32,
pub status: u32,
pub length: u64,
}
impl CtdbHeader {
pub fn new(command: CtdbCommand, status: CtdbStatus, length: u64) -> Self {
Self {
magic: CTDB_MAGIC,
version: CTDB_VERSION,
command: command as u32,
status: status as u32,
length,
}
}
pub fn to_bytes(&self) -> Vec<u8> {
let mut buf = Vec::with_capacity(CTDB_HEADER_SIZE);
buf.extend_from_slice(&self.magic.to_le_bytes());
buf.extend_from_slice(&self.version.to_le_bytes());
buf.extend_from_slice(&self.command.to_le_bytes());
buf.extend_from_slice(&self.status.to_le_bytes());
buf.extend_from_slice(&self.length.to_le_bytes());
buf
}
pub fn from_bytes(buf: &[u8]) -> Result<Self, CtdbProtoError> {
if buf.len() < CTDB_HEADER_SIZE {
return Err(CtdbProtoError::HeaderTooShort);
}
let magic = u32::from_le_bytes(buf[0..4].try_into().unwrap());
let version = u32::from_le_bytes(buf[4..8].try_into().unwrap());
let command = u32::from_le_bytes(buf[8..12].try_into().unwrap());
let status = u32::from_le_bytes(buf[12..16].try_into().unwrap());
let length = u64::from_le_bytes(buf[16..24].try_into().unwrap());
Ok(Self {
magic,
version,
command,
status,
length,
})
}
pub fn is_valid(&self) -> bool {
self.magic == CTDB_MAGIC && self.version == CTDB_VERSION
}
}
#[derive(Debug, Clone)]
pub struct CtdbMessage {
pub header: CtdbHeader,
pub payload: Vec<u8>,
}
impl CtdbMessage {
pub fn new(command: CtdbCommand, status: CtdbStatus, payload: Vec<u8>) -> Self {
let length = payload.len() as u64;
Self {
header: CtdbHeader::new(command, status, length),
payload,
}
}
pub fn to_bytes(&self) -> Vec<u8> {
let mut buf = self.header.to_bytes();
buf.extend_from_slice(&self.payload);
buf
}
pub fn from_bytes(buf: &[u8]) -> Result<Self, CtdbProtoError> {
let header = CtdbHeader::from_bytes(buf)?;
if !header.is_valid() {
return Err(CtdbProtoError::InvalidMagic);
}
let payload = if buf.len() > CTDB_HEADER_SIZE {
buf[CTDB_HEADER_SIZE..].to_vec()
} else {
Vec::new()
};
Ok(Self { header, payload })
}
pub fn command(&self) -> CtdbCommand {
CtdbCommand::from_u32(self.header.command)
}
pub fn status(&self) -> CtdbStatus {
CtdbStatus::from_u32(self.header.status)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CtdbProtoError {
HeaderTooShort,
InvalidMagic,
IoError,
InvalidPayload,
}
impl std::fmt::Display for CtdbProtoError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
CtdbProtoError::HeaderTooShort => write!(f, "header too short"),
CtdbProtoError::InvalidMagic => write!(f, "invalid magic number"),
CtdbProtoError::IoError => write!(f, "I/O error"),
CtdbProtoError::InvalidPayload => write!(f, "invalid payload"),
}
}
}
impl std::error::Error for CtdbProtoError {}
impl From<io::Error> for CtdbProtoError {
fn from(_: io::Error) -> Self {
CtdbProtoError::IoError
}
}
pub mod payload {
use super::CtdbProtoError;
pub fn encode_kv(key: &[u8], value: &[u8]) -> Vec<u8> {
let mut buf = Vec::new();
buf.extend_from_slice(&(key.len() as u32).to_le_bytes());
buf.extend_from_slice(key);
buf.extend_from_slice(&(value.len() as u32).to_le_bytes());
buf.extend_from_slice(value);
buf
}
pub fn decode_kv(payload: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CtdbProtoError> {
if payload.len() < 4 {
return Err(CtdbProtoError::InvalidPayload);
}
let key_len = u32::from_le_bytes(payload[0..4].try_into().unwrap()) as usize;
if payload.len() < 4 + key_len + 4 {
return Err(CtdbProtoError::InvalidPayload);
}
let key = payload[4..4 + key_len].to_vec();
let val_len_offset = 4 + key_len;
let val_len = u32::from_le_bytes(
payload[val_len_offset..val_len_offset + 4].try_into().unwrap(),
) as usize;
let val_start = val_len_offset + 4;
if payload.len() < val_start + val_len {
return Err(CtdbProtoError::InvalidPayload);
}
let value = payload[val_start..val_start + val_len].to_vec();
Ok((key, value))
}
pub fn encode_key(key: &[u8]) -> Vec<u8> {
let mut buf = Vec::new();
buf.extend_from_slice(&(key.len() as u32).to_le_bytes());
buf.extend_from_slice(key);
buf
}
pub fn decode_key(payload: &[u8]) -> Result<Vec<u8>, CtdbProtoError> {
if payload.len() < 4 {
return Err(CtdbProtoError::InvalidPayload);
}
let key_len = u32::from_le_bytes(payload[0..4].try_into().unwrap()) as usize;
if payload.len() < 4 + key_len {
return Err(CtdbProtoError::InvalidPayload);
}
Ok(payload[4..4 + key_len].to_vec())
}
pub fn encode_node_id(id: u32) -> Vec<u8> {
id.to_le_bytes().to_vec()
}
pub fn decode_node_id(payload: &[u8]) -> Result<u32, CtdbProtoError> {
if payload.len() < 4 {
return Err(CtdbProtoError::InvalidPayload);
}
Ok(u32::from_le_bytes(payload[0..4].try_into().unwrap()))
}
pub fn encode_nodemask(active: &[u32]) -> Vec<u8> {
let mut buf = Vec::new();
buf.extend_from_slice(&(active.len() as u32).to_le_bytes());
for &id in active {
buf.extend_from_slice(&id.to_le_bytes());
}
buf
}
pub fn decode_nodemask(payload: &[u8]) -> Result<Vec<u32>, CtdbProtoError> {
if payload.len() < 4 {
return Err(CtdbProtoError::InvalidPayload);
}
let count = u32::from_le_bytes(payload[0..4].try_into().unwrap()) as usize;
let mut ids = Vec::with_capacity(count);
for i in 0..count {
let offset = 4 + i * 4;
if payload.len() < offset + 4 {
return Err(CtdbProtoError::InvalidPayload);
}
ids.push(u32::from_le_bytes(payload[offset..offset + 4].try_into().unwrap()));
}
Ok(ids)
}
pub fn encode_ip(ip: &str, interface: &str) -> Vec<u8> {
let ip_bytes = ip.as_bytes();
let if_bytes = interface.as_bytes();
let mut buf = Vec::new();
buf.extend_from_slice(&(ip_bytes.len() as u32).to_le_bytes());
buf.extend_from_slice(ip_bytes);
buf.extend_from_slice(&(if_bytes.len() as u32).to_le_bytes());
buf.extend_from_slice(if_bytes);
buf
}
pub fn decode_ip(payload: &[u8]) -> Result<(String, String), CtdbProtoError> {
if payload.len() < 4 {
return Err(CtdbProtoError::InvalidPayload);
}
let ip_len = u32::from_le_bytes(payload[0..4].try_into().unwrap()) as usize;
if payload.len() < 4 + ip_len + 4 {
return Err(CtdbProtoError::InvalidPayload);
}
let ip = String::from_utf8_lossy(&payload[4..4 + ip_len]).to_string();
let if_offset = 4 + ip_len;
let if_len = u32::from_le_bytes(
payload[if_offset..if_offset + 4].try_into().unwrap(),
) as usize;
let if_start = if_offset + 4;
if payload.len() < if_start + if_len {
return Err(CtdbProtoError::InvalidPayload);
}
let interface = String::from_utf8_lossy(&payload[if_start..if_start + if_len]).to_string();
Ok((ip, interface))
}
}
pub struct CtdbConnection {
stream: TcpStream,
}
impl CtdbConnection {
pub fn new(stream: TcpStream) -> Self {
Self { stream }
}
pub fn connect(addr: &str) -> Result<Self, CtdbProtoError> {
let stream = TcpStream::connect(addr)?;
Ok(Self { stream })
}
pub fn send_message(&mut self, msg: &CtdbMessage) -> Result<(), CtdbProtoError> {
let bytes = msg.to_bytes();
self.stream.write_all(&bytes)?;
Ok(())
}
pub fn recv_message(&mut self) -> Result<CtdbMessage, CtdbProtoError> {
let mut header_buf = [0u8; CTDB_HEADER_SIZE];
self.stream.read_exact(&mut header_buf)?;
let header = CtdbHeader::from_bytes(&header_buf)?;
if !header.is_valid() {
return Err(CtdbProtoError::InvalidMagic);
}
let payload_len = header.length as usize;
let mut payload = vec![0u8; payload_len];
if payload_len > 0 {
self.stream.read_exact(&mut payload)?;
}
Ok(CtdbMessage { header, payload })
}
pub fn ping(&mut self) -> Result<(), CtdbProtoError> {
let msg = CtdbMessage::new(CtdbCommand::Ping, CtdbStatus::Success, vec![]);
self.send_message(&msg)?;
let resp = self.recv_message()?;
if resp.command() == CtdbCommand::Pong && resp.status() == CtdbStatus::Success {
Ok(())
} else {
Err(CtdbProtoError::InvalidPayload)
}
}
pub fn store(&mut self, key: &[u8], value: &[u8]) -> Result<bool, CtdbProtoError> {
let payload = payload::encode_kv(key, value);
let msg = CtdbMessage::new(CtdbCommand::Store, CtdbStatus::Success, payload);
self.send_message(&msg)?;
let resp = self.recv_message()?;
Ok(resp.status() == CtdbStatus::Success)
}
pub fn fetch(&mut self, key: &[u8]) -> Result<Vec<u8>, CtdbProtoError> {
let payload = payload::encode_key(key);
let msg = CtdbMessage::new(CtdbCommand::Fetch, CtdbStatus::Success, payload);
self.send_message(&msg)?;
let resp = self.recv_message()?;
if resp.status() == CtdbStatus::Success {
Ok(resp.payload)
} else {
Err(CtdbProtoError::InvalidPayload)
}
}
pub fn delete(&mut self, key: &[u8]) -> Result<bool, CtdbProtoError> {
let payload = payload::encode_key(key);
let msg = CtdbMessage::new(CtdbCommand::Delete, CtdbStatus::Success, payload);
self.send_message(&msg)?;
let resp = self.recv_message()?;
Ok(resp.status() == CtdbStatus::Success)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_header_roundtrip() {
let header = CtdbHeader::new(CtdbCommand::Ping, CtdbStatus::Success, 42);
let bytes = header.to_bytes();
let restored = CtdbHeader::from_bytes(&bytes).unwrap();
assert_eq!(restored.magic, CTDB_MAGIC);
assert_eq!(restored.version, CTDB_VERSION);
assert_eq!(restored.command, CtdbCommand::Ping as u32);
assert_eq!(restored.status, CtdbStatus::Success as u32);
assert_eq!(restored.length, 42);
assert!(restored.is_valid());
}
#[test]
fn test_message_roundtrip() {
let msg = CtdbMessage::new(
CtdbCommand::Store,
CtdbStatus::Success,
b"test_payload".to_vec(),
);
let bytes = msg.to_bytes();
let restored = CtdbMessage::from_bytes(&bytes).unwrap();
assert_eq!(restored.command(), CtdbCommand::Store);
assert_eq!(restored.payload, b"test_payload");
}
#[test]
fn test_command_from_u32() {
assert_eq!(CtdbCommand::from_u32(1), CtdbCommand::Connect);
assert_eq!(CtdbCommand::from_u32(3), CtdbCommand::Ping);
assert_eq!(CtdbCommand::from_u32(0xFFFF), CtdbCommand::Unknown);
assert_eq!(CtdbCommand::from_u32(999), CtdbCommand::Unknown);
}
#[test]
fn test_status_from_u32() {
assert_eq!(CtdbStatus::from_u32(0), CtdbStatus::Success);
assert_eq!(CtdbStatus::from_u32(2), CtdbStatus::NotFound);
assert_eq!(CtdbStatus::from_u32(99), CtdbStatus::Error);
}
#[test]
fn test_payload_encode_decode_kv() {
let (key, val) = (b"mykey", b"myvalue");
let encoded = payload::encode_kv(key, val);
let (k, v) = payload::decode_kv(&encoded).unwrap();
assert_eq!(k, key);
assert_eq!(v, val);
}
#[test]
fn test_payload_encode_decode_key() {
let key = b"test_key";
let encoded = payload::encode_key(key);
let decoded = payload::decode_key(&encoded).unwrap();
assert_eq!(decoded, key);
}
#[test]
fn test_payload_node_id() {
let encoded = payload::encode_node_id(42);
let decoded = payload::decode_node_id(&encoded).unwrap();
assert_eq!(decoded, 42);
}
#[test]
fn test_payload_nodemask() {
let ids = vec![0u32, 1, 2, 3];
let encoded = payload::encode_nodemask(&ids);
let decoded = payload::decode_nodemask(&encoded).unwrap();
assert_eq!(decoded, ids);
}
#[test]
fn test_payload_ip() {
let encoded = payload::encode_ip("192.168.1.100", "eth0");
let (ip, iface) = payload::decode_ip(&encoded).unwrap();
assert_eq!(ip, "192.168.1.100");
assert_eq!(iface, "eth0");
}
#[test]
fn test_invalid_magic() {
let mut bad_header = CtdbHeader::new(CtdbCommand::Ping, CtdbStatus::Success, 0);
bad_header.magic = 0xDEADBEEF;
assert!(!bad_header.is_valid());
}
#[test]
fn test_empty_message() {
let msg = CtdbMessage::new(CtdbCommand::Connect, CtdbStatus::Success, vec![]);
let bytes = msg.to_bytes();
let restored = CtdbMessage::from_bytes(&bytes).unwrap();
assert!(restored.payload.is_empty());
assert_eq!(restored.header.length, 0);
}
#[test]
fn test_header_too_short() {
let result = CtdbHeader::from_bytes(&[0u8; 10]);
assert!(result.is_err());
}
#[test]
fn test_large_payload() {
let large = vec![0xABu8; 65000];
let msg = CtdbMessage::new(CtdbCommand::Fetch, CtdbStatus::Success, large.clone());
let bytes = msg.to_bytes();
let restored = CtdbMessage::from_bytes(&bytes).unwrap();
assert_eq!(restored.payload.len(), 65000);
assert_eq!(restored.payload, large);
}
}
+417
View File
@@ -0,0 +1,417 @@
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
use super::ip_manager::IpManager;
use super::node::{NodeManager, NodeId, NodeState};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RecoveryState {
Idle,
Initiated,
InProgress,
Verifying,
Completed,
Failed,
}
impl RecoveryState {
pub fn as_str(&self) -> &'static str {
match self {
RecoveryState::Idle => "IDLE",
RecoveryState::Initiated => "INITIATED",
RecoveryState::InProgress => "IN_PROGRESS",
RecoveryState::Verifying => "VERIFYING",
RecoveryState::Completed => "COMPLETED",
RecoveryState::Failed => "FAILED",
}
}
}
#[derive(Debug, Clone)]
pub struct RecoveryEvent {
pub timestamp: Instant,
pub event_type: RecoveryEventType,
pub node_id: Option<NodeId>,
pub message: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RecoveryEventType {
NodeDown,
NodeRejoin,
IpMigration,
RecoveryStart,
RecoveryComplete,
RecoveryFailed,
}
pub struct RecoveryManager {
state: RwLock<RecoveryState>,
events: RwLock<Vec<RecoveryEvent>>,
recovery_cooldown: Duration,
last_recovery: RwLock<Option<Instant>>,
max_events: usize,
}
impl RecoveryManager {
pub fn new() -> Self {
Self {
state: RwLock::new(RecoveryState::Idle),
events: RwLock::new(Vec::new()),
recovery_cooldown: Duration::from_secs(30),
last_recovery: RwLock::new(None),
max_events: 100,
}
}
pub fn with_cooldown(mut self, cooldown: Duration) -> Self {
self.recovery_cooldown = cooldown;
self
}
pub fn state(&self) -> RecoveryState {
*self.state.read().unwrap()
}
pub fn set_state(&self, state: RecoveryState) {
*self.state.write().unwrap() = state;
}
pub fn log_event(&self, event_type: RecoveryEventType, node_id: Option<NodeId>, msg: &str) {
let event = RecoveryEvent {
timestamp: Instant::now(),
event_type,
node_id,
message: msg.to_string(),
};
let mut events = self.events.write().unwrap();
events.push(event);
if events.len() > self.max_events {
events.remove(0);
}
}
pub fn events(&self) -> Vec<RecoveryEvent> {
self.events.read().unwrap().clone()
}
pub fn can_recover(&self) -> bool {
match *self.last_recovery.read().unwrap() {
Some(t) => t.elapsed() > self.recovery_cooldown,
None => true,
}
}
pub fn handle_node_failure(
&self,
failed_node: NodeId,
node_mgr: &NodeManager,
ip_mgr: &IpManager,
) -> Vec<(super::ip_manager::PublicIpId, NodeId)> {
self.log_event(RecoveryEventType::NodeDown, Some(failed_node), &format!(
"Node {:?} marked DOWN", failed_node
));
if !self.can_recover() {
self.log_event(
RecoveryEventType::RecoveryFailed,
Some(failed_node),
"Recovery skipped: cooldown active",
);
return Vec::new();
}
self.set_state(RecoveryState::Initiated);
self.log_event(RecoveryEventType::RecoveryStart, Some(failed_node), "Starting recovery");
*self.last_recovery.write().unwrap() = Some(Instant::now());
self.set_state(RecoveryState::InProgress);
let active_nodes: Vec<NodeId> = node_mgr
.active_nodes()
.iter()
.filter(|n| n.id != failed_node)
.map(|n| n.id)
.collect();
if active_nodes.is_empty() {
self.log_event(
RecoveryEventType::RecoveryFailed,
Some(failed_node),
"No active nodes available for IP migration",
);
self.set_state(RecoveryState::Failed);
return Vec::new();
}
let reassignments = ip_mgr.reassign_on_failure(failed_node, &active_nodes);
for (ip_id, new_owner) in &reassignments {
self.log_event(
RecoveryEventType::IpMigration,
Some(*new_owner),
&format!("IP {:?} migrated to node {:?}", ip_id, new_owner),
);
}
self.set_state(RecoveryState::Completed);
self.log_event(
RecoveryEventType::RecoveryComplete,
Some(failed_node),
&format!("Recovery complete: {} IPs migrated", reassignments.len()),
);
reassignments
}
pub fn handle_node_rejoin(
&self,
rejoining_node: NodeId,
node_mgr: &NodeManager,
ip_mgr: &IpManager,
) -> Vec<(super::ip_manager::PublicIpId, NodeId)> {
self.log_event(RecoveryEventType::NodeRejoin, Some(rejoining_node), &format!(
"Node {:?} rejoining cluster", rejoining_node
));
let active_nodes: Vec<NodeId> = node_mgr.active_nodes().iter().map(|n| n.id).collect();
self.set_state(RecoveryState::InProgress);
let rebalance = ip_mgr.balance(&active_nodes);
for (ip_id, new_owner) in &rebalance {
self.log_event(
RecoveryEventType::IpMigration,
Some(*new_owner),
&format!("IP {:?} rebalanced to node {:?}", ip_id, new_owner),
);
}
self.set_state(RecoveryState::Completed);
self.log_event(
RecoveryEventType::RecoveryComplete,
Some(rejoining_node),
&format!("Rebalance complete: {} IPs reassigned", rebalance.len()),
);
rebalance
}
pub fn check_and_recover(
&self,
node_mgr: &NodeManager,
ip_mgr: &IpManager,
) -> Vec<(super::ip_manager::PublicIpId, NodeId)> {
let transitions = node_mgr.check_health();
if transitions.is_empty() {
return Vec::new();
}
let mut all_reassignments = Vec::new();
for (node_id, new_state) in &transitions {
match new_state {
NodeState::Down => {
let reassignments =
self.handle_node_failure(*node_id, node_mgr, ip_mgr);
all_reassignments.extend(reassignments);
}
NodeState::Up => {
let reassignments =
self.handle_node_rejoin(*node_id, node_mgr, ip_mgr);
all_reassignments.extend(reassignments);
}
_ => {}
}
}
all_reassignments
}
pub fn reset(&self) {
*self.state.write().unwrap() = RecoveryState::Idle;
*self.last_recovery.write().unwrap() = None;
}
pub fn events_by_node(&self, node_id: NodeId) -> Vec<RecoveryEvent> {
self.events
.read()
.unwrap()
.iter()
.filter(|e| e.node_id == Some(node_id))
.cloned()
.collect()
}
pub fn last_event(&self) -> Option<RecoveryEvent> {
self.events.read().unwrap().last().cloned()
}
}
impl Default for RecoveryManager {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::thread;
fn addr(port: u16) -> SocketAddr {
SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port))
}
fn ip(a: u8, b: u8, c: u8, d: u8) -> std::net::IpAddr {
std::net::IpAddr::V4(Ipv4Addr::new(a, b, c, d))
}
fn setup() -> (NodeManager, IpManager, RecoveryManager) {
let mgr = NodeManager::new(NodeId(0), addr(4000));
mgr.add_node(NodeId(1), addr(4001));
mgr.add_node(NodeId(2), addr(4002));
mgr.record_heartbeat(NodeId(1));
mgr.record_heartbeat(NodeId(2));
let ip_mgr = IpManager::new();
ip_mgr.add_ip(ip(192, 168, 1, 100), "eth0");
ip_mgr.add_ip(ip(192, 168, 1, 101), "eth0");
ip_mgr.add_ip(ip(192, 168, 1, 102), "eth1");
ip_mgr.assign_ip(super::super::ip_manager::PublicIpId(0), NodeId(0)).unwrap();
ip_mgr.assign_ip(super::super::ip_manager::PublicIpId(1), NodeId(1)).unwrap();
ip_mgr.assign_ip(super::super::ip_manager::PublicIpId(2), NodeId(2)).unwrap();
let recovery = RecoveryManager::new();
(mgr, ip_mgr, recovery)
}
#[test]
fn test_initial_state() {
let recovery = RecoveryManager::new();
assert_eq!(recovery.state(), RecoveryState::Idle);
assert!(recovery.can_recover());
assert!(recovery.events().is_empty());
}
#[test]
fn test_node_failure_triggers_recovery() {
let (mgr, ip_mgr, recovery) = setup();
let reassignments = recovery.handle_node_failure(NodeId(1), &mgr, &ip_mgr);
assert!(!reassignments.is_empty());
assert_eq!(recovery.state(), RecoveryState::Completed);
let events = recovery.events();
assert!(events.iter().any(|e| e.event_type == RecoveryEventType::NodeDown));
assert!(events.iter().any(|e| e.event_type == RecoveryEventType::RecoveryStart));
assert!(events.iter().any(|e| e.event_type == RecoveryEventType::RecoveryComplete));
}
#[test]
fn test_no_active_nodes_fails() {
let mgr = NodeManager::new(NodeId(0), addr(4000));
let ip_mgr = IpManager::new();
ip_mgr.add_ip(ip(192, 168, 1, 100), "eth0");
let recovery = RecoveryManager::new();
let reassignments = recovery.handle_node_failure(NodeId(0), &mgr, &ip_mgr);
assert!(reassignments.is_empty());
assert_eq!(recovery.state(), RecoveryState::Failed);
}
#[test]
fn test_cooldown_prevents_rapid_recovery() {
let (mgr, ip_mgr, recovery) = setup();
recovery.handle_node_failure(NodeId(1), &mgr, &ip_mgr);
assert!(!recovery.can_recover());
}
#[test]
fn test_cooldown_expires() {
let recovery = RecoveryManager::new().with_cooldown(Duration::from_millis(10));
*recovery.last_recovery.write().unwrap() = Some(Instant::now());
assert!(!recovery.can_recover());
thread::sleep(Duration::from_millis(20));
assert!(recovery.can_recover());
}
#[test]
fn test_node_rejoin_balances() {
let (mgr, ip_mgr, recovery) = setup();
recovery.handle_node_failure(NodeId(1), &mgr, &ip_mgr);
thread::sleep(Duration::from_millis(10));
let recovery2 = RecoveryManager::new();
let reassignments = recovery2.handle_node_rejoin(NodeId(1), &mgr, &ip_mgr);
assert!(!reassignments.is_empty());
}
#[test]
fn test_log_event() {
let recovery = RecoveryManager::new();
recovery.log_event(RecoveryEventType::NodeDown, Some(NodeId(5)), "test message");
assert_eq!(recovery.events().len(), 1);
let event = &recovery.events()[0];
assert_eq!(event.node_id, Some(NodeId(5)));
assert_eq!(event.message, "test message");
}
#[test]
fn test_event_max_limit() {
let recovery = RecoveryManager::new();
for i in 0..200 {
recovery.log_event(RecoveryEventType::NodeDown, Some(NodeId(i)), "msg");
}
assert_eq!(recovery.events().len(), 100);
}
#[test]
fn test_events_by_node() {
let recovery = RecoveryManager::new();
recovery.log_event(RecoveryEventType::NodeDown, Some(NodeId(1)), "a");
recovery.log_event(RecoveryEventType::NodeDown, Some(NodeId(2)), "b");
recovery.log_event(RecoveryEventType::IpMigration, Some(NodeId(1)), "c");
let node1_events = recovery.events_by_node(NodeId(1));
assert_eq!(node1_events.len(), 2);
}
#[test]
fn test_last_event() {
let recovery = RecoveryManager::new();
recovery.log_event(RecoveryEventType::NodeDown, Some(NodeId(0)), "first");
recovery.log_event(RecoveryEventType::RecoveryComplete, Some(NodeId(0)), "last");
let last = recovery.last_event().unwrap();
assert_eq!(last.message, "last");
}
#[test]
fn test_reset() {
let (mgr, ip_mgr, recovery) = setup();
recovery.handle_node_failure(NodeId(1), &mgr, &ip_mgr);
assert_ne!(recovery.state(), RecoveryState::Idle);
recovery.reset();
assert_eq!(recovery.state(), RecoveryState::Idle);
assert!(recovery.can_recover());
}
#[test]
fn test_check_and_recover_no_transitions() {
let (mgr, ip_mgr, recovery) = setup();
let result = recovery.check_and_recover(&mgr, &ip_mgr);
assert!(result.is_empty());
}
#[test]
fn test_recovery_state_string() {
assert_eq!(RecoveryState::Idle.as_str(), "IDLE");
assert_eq!(RecoveryState::InProgress.as_str(), "IN_PROGRESS");
assert_eq!(RecoveryState::Failed.as_str(), "FAILED");
assert_eq!(RecoveryState::Completed.as_str(), "COMPLETED");
}
}
+678
View File
@@ -0,0 +1,678 @@
use std::collections::HashMap;
use std::io::{self, Read, Write, Seek, SeekFrom};
use std::path::Path;
use std::sync::{Arc, Mutex, RwLock};
const TDB_MAGIC: u32 = 0x1BADFACE;
const TDB_VERSION: u32 = 1;
const DEFAULT_HASH_SIZE: u32 = 1024;
const TDB_HEADER_SIZE: u64 = 128;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RecordFlag {
Active,
Free,
Deleted,
}
impl RecordFlag {
fn as_u32(&self) -> u32 {
match self {
RecordFlag::Active => 0,
RecordFlag::Free => 1,
RecordFlag::Deleted => 2,
}
}
fn from_u32(v: u32) -> Self {
match v {
1 => RecordFlag::Free,
2 => RecordFlag::Deleted,
_ => RecordFlag::Active,
}
}
}
#[derive(Debug, Clone)]
pub struct TdbRecord {
pub key: Vec<u8>,
pub data: Vec<u8>,
pub flag: RecordFlag,
pub hash_next: u64,
}
impl TdbRecord {
pub fn new(key: Vec<u8>, data: Vec<u8>) -> Self {
Self {
key,
data,
flag: RecordFlag::Active,
hash_next: 0,
}
}
pub fn key_str(&self) -> &str {
std::str::from_utf8(&self.key).unwrap_or("")
}
pub fn data_str(&self) -> &str {
std::str::from_utf8(&self.data).unwrap_or("")
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TdbError {
IoError,
NotFound,
Corrupt,
Exists,
LockFailed,
}
impl std::fmt::Display for TdbError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
TdbError::IoError => write!(f, "I/O error"),
TdbError::NotFound => write!(f, "record not found"),
TdbError::Corrupt => write!(f, "database corrupt"),
TdbError::Exists => write!(f, "record already exists"),
TdbError::LockFailed => write!(f, "lock failed"),
}
}
}
impl std::error::Error for TdbError {}
impl From<io::Error> for TdbError {
fn from(_: io::Error) -> Self {
TdbError::IoError
}
}
pub type TdbResult<T> = Result<T, TdbError>;
fn hash_key(key: &[u8], hash_size: u32) -> u32 {
let mut h: u32 = 0;
for &b in key {
h = h.wrapping_mul(31).wrapping_add(b as u32);
}
h % hash_size
}
#[derive(Debug, Clone)]
struct HashEntry {
offset: u64,
}
struct TdbHeader {
magic: u32,
version: u32,
hash_size: u32,
record_count: u64,
}
impl TdbHeader {
fn new(hash_size: u32) -> Self {
Self {
magic: TDB_MAGIC,
version: TDB_VERSION,
hash_size,
record_count: 0,
}
}
fn to_bytes(&self) -> Vec<u8> {
let mut buf = Vec::with_capacity(TDB_HEADER_SIZE as usize);
buf.extend_from_slice(&self.magic.to_le_bytes());
buf.extend_from_slice(&self.version.to_le_bytes());
buf.extend_from_slice(&self.hash_size.to_le_bytes());
buf.extend_from_slice(&self.record_count.to_le_bytes());
while buf.len() < TDB_HEADER_SIZE as usize {
buf.push(0);
}
buf
}
fn from_bytes(buf: &[u8]) -> TdbResult<Self> {
if buf.len() < 20 {
return Err(TdbError::Corrupt);
}
let magic = u32::from_le_bytes(buf[0..4].try_into().unwrap());
if magic != TDB_MAGIC {
return Err(TdbError::Corrupt);
}
let version = u32::from_le_bytes(buf[4..8].try_into().unwrap());
let hash_size = u32::from_le_bytes(buf[8..12].try_into().unwrap());
let record_count = u64::from_le_bytes(buf[12..20].try_into().unwrap());
Ok(Self {
magic,
version,
hash_size,
record_count,
})
}
}
pub struct TdbEngine {
header: RwLock<TdbHeader>,
records: RwLock<HashMap<Vec<u8>, TdbRecord>>,
hash_table: RwLock<Vec<Vec<u64>>>,
file_path: Option<std::path::PathBuf>,
dirty: Mutex<bool>,
}
impl TdbEngine {
pub fn new() -> Self {
Self::with_hash_size(DEFAULT_HASH_SIZE)
}
pub fn with_hash_size(hash_size: u32) -> Self {
Self {
header: RwLock::new(TdbHeader::new(hash_size)),
records: RwLock::new(HashMap::new()),
hash_table: RwLock::new(vec![Vec::new(); 1]),
file_path: None,
dirty: Mutex::new(false),
}
}
pub fn open<P: AsRef<Path>>(path: P) -> TdbResult<Self> {
let path = path.as_ref().to_path_buf();
if path.exists() {
Self::load_from_file(&path)
} else {
let engine = Self::new();
let mut engine = engine;
engine.file_path = Some(path.clone());
engine.save_to_file(&path)?;
Ok(engine)
}
}
pub fn create<P: AsRef<Path>>(path: P, hash_size: u32) -> TdbResult<Self> {
let path = path.as_ref().to_path_buf();
let engine = Self::with_hash_size(hash_size);
let mut engine = engine;
engine.file_path = Some(path.clone());
engine.save_to_file(&path)?;
Ok(engine)
}
fn load_from_file(path: &Path) -> TdbResult<Self> {
let mut file = std::fs::File::open(path)?;
let mut header_buf = vec![0u8; TDB_HEADER_SIZE as usize];
file.read_exact(&mut header_buf)?;
let header = TdbHeader::from_bytes(&header_buf)?;
let hash_size = header.hash_size;
let mut hash_table = vec![Vec::new(); hash_size as usize];
let mut records = HashMap::new();
let mut pos = TDB_HEADER_SIZE;
let file_meta = file.metadata()?;
let file_size = file_meta.len();
while pos < file_size {
file.seek(SeekFrom::Start(pos))?;
let mut flag_buf = [0u8; 4];
let mut key_len_buf = [0u8; 4];
let mut data_len_buf = [0u8; 4];
let mut hash_next_buf = [0u8; 8];
if file.read_exact(&mut flag_buf).is_err() {
break;
}
let _ = file.read_exact(&mut key_len_buf);
let _ = file.read_exact(&mut data_len_buf);
let _ = file.read_exact(&mut hash_next_buf);
let flag = RecordFlag::from_u32(u32::from_le_bytes(flag_buf));
let key_len = u32::from_le_bytes(key_len_buf) as usize;
let data_len = u32::from_le_bytes(data_len_buf) as usize;
let hash_next = u64::from_le_bytes(hash_next_buf);
let mut key = vec![0u8; key_len];
let mut data = vec![0u8; data_len];
let _ = file.read_exact(&mut key);
let _ = file.read_exact(&mut data);
let rec_offset = pos;
let bucket = hash_key(&key, hash_size) as usize;
if bucket < hash_table.len() {
hash_table[bucket].push(rec_offset);
}
if flag == RecordFlag::Active {
let record = TdbRecord {
key: key.clone(),
data,
flag,
hash_next,
};
records.insert(key, record);
}
let rec_size = 4 + 4 + 4 + 8 + key_len as u64 + data_len as u64;
pos += rec_size;
}
let record_count = records.len() as u64;
let mut header = header;
header.record_count = record_count;
Ok(Self {
header: RwLock::new(header),
records: RwLock::new(records),
hash_table: RwLock::new(hash_table),
file_path: Some(path.to_path_buf()),
dirty: Mutex::new(false),
})
}
fn save_to_file(&self, path: &Path) -> TdbResult<()> {
let header = self.header.read().unwrap();
let header_bytes = header.to_bytes();
let mut file = std::fs::File::create(path)?;
file.write_all(&header_bytes)?;
let records = self.records.read().unwrap();
for (_, record) in records.iter() {
let flag = record.flag.as_u32().to_le_bytes();
let key_len = (record.key.len() as u32).to_le_bytes();
let data_len = (record.data.len() as u32).to_le_bytes();
let hash_next = record.hash_next.to_le_bytes();
file.write_all(&flag)?;
file.write_all(&key_len)?;
file.write_all(&data_len)?;
file.write_all(&hash_next)?;
file.write_all(&record.key)?;
file.write_all(&record.data)?;
}
file.flush()?;
Ok(())
}
pub fn store(&self, key: Vec<u8>, data: Vec<u8>) -> TdbResult<bool> {
let mut records = self.records.write().unwrap();
let existed = records.contains_key(&key);
let record = TdbRecord::new(key.clone(), data);
records.insert(key, record);
*self.dirty.lock().unwrap() = true;
let mut header = self.header.write().unwrap();
header.record_count = records.len() as u64;
drop(header);
drop(records);
self.try_flush()?;
Ok(!existed)
}
pub fn fetch(&self, key: &[u8]) -> TdbResult<Vec<u8>> {
let records = self.records.read().unwrap();
match records.get(key) {
Some(record) => Ok(record.data.clone()),
None => Err(TdbError::NotFound),
}
}
pub fn delete(&self, key: &[u8]) -> TdbResult<()> {
let mut records = self.records.write().unwrap();
if records.remove(key).is_some() {
*self.dirty.lock().unwrap() = true;
let mut header = self.header.write().unwrap();
header.record_count = records.len() as u64;
drop(header);
drop(records);
self.try_flush()?;
Ok(())
} else {
Err(TdbError::NotFound)
}
}
pub fn exists(&self, key: &[u8]) -> bool {
self.records.read().unwrap().contains_key(key)
}
pub fn keys(&self) -> Vec<Vec<u8>> {
self.records.read().unwrap().keys().cloned().collect()
}
pub fn record_count(&self) -> u64 {
self.header.read().unwrap().record_count
}
pub fn iter(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
self.records
.read()
.unwrap()
.iter()
.map(|(k, v)| (k.clone(), v.data.clone()))
.collect()
}
pub fn try_flush(&self) -> TdbResult<()> {
let dirty = *self.dirty.lock().unwrap();
if dirty {
if let Some(ref path) = self.file_path {
self.save_to_file(path)?;
*self.dirty.lock().unwrap() = false;
}
}
Ok(())
}
pub fn flush(&self) -> TdbResult<()> {
if let Some(ref path) = self.file_path {
self.save_to_file(path)?;
*self.dirty.lock().unwrap() = false;
}
Ok(())
}
}
impl Default for TdbEngine {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct TdbStats {
pub record_count: u64,
pub hash_size: u32,
pub active_records: u64,
pub total_data_size: u64,
pub avg_key_size: f64,
pub avg_data_size: f64,
}
impl TdbEngine {
pub fn stats(&self) -> TdbStats {
let records = self.records.read().unwrap();
let header = self.header.read().unwrap();
let count = records.len();
let total_data: u64 = records.values().map(|r| r.data.len() as u64).sum();
let total_key: u64 = records.values().map(|r| r.key.len() as u64).sum();
TdbStats {
record_count: header.record_count,
hash_size: header.hash_size,
active_records: count as u64,
total_data_size: total_data,
avg_key_size: if count > 0 {
total_key as f64 / count as f64
} else {
0.0
},
avg_data_size: if count > 0 {
total_data as f64 / count as f64
} else {
0.0
},
}
}
}
pub struct TdbTransaction<'a> {
engine: &'a TdbEngine,
snapshot: HashMap<Vec<u8>, TdbRecord>,
committed: bool,
}
impl<'a> TdbTransaction<'a> {
pub fn begin(engine: &'a TdbEngine) -> Self {
let snapshot = engine.records.read().unwrap().clone();
Self {
engine,
snapshot,
committed: false,
}
}
pub fn store(&self, key: Vec<u8>, data: Vec<u8>) -> TdbResult<bool> {
self.engine.store(key, data)
}
pub fn fetch(&self, key: &[u8]) -> TdbResult<Vec<u8>> {
self.engine.fetch(key)
}
pub fn delete(&self, key: &[u8]) -> TdbResult<()> {
self.engine.delete(key)
}
pub fn commit(&mut self) -> TdbResult<()> {
self.committed = true;
self.engine.flush()
}
pub fn rollback(&mut self) -> TdbResult<()> {
if !self.committed {
let mut records = self.engine.records.write().unwrap();
*records = self.snapshot.clone();
let mut header = self.engine.header.write().unwrap();
header.record_count = records.len() as u64;
*self.engine.dirty.lock().unwrap() = true;
self.engine.flush()?;
}
Ok(())
}
}
impl<'a> Drop for TdbTransaction<'a> {
fn drop(&mut self) {
if !self.committed {
let _ = self.rollback();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn test_tdb_create_and_store() {
let engine = TdbEngine::new();
let created = engine.store(b"key1".to_vec(), b"value1".to_vec()).unwrap();
assert!(created);
assert_eq!(engine.fetch(b"key1").unwrap(), b"value1");
assert_eq!(engine.record_count(), 1);
}
#[test]
fn test_tdb_overwrite() {
let engine = TdbEngine::new();
engine.store(b"key1".to_vec(), b"v1".to_vec()).unwrap();
let created = engine.store(b"key1".to_vec(), b"v2".to_vec()).unwrap();
assert!(!created);
assert_eq!(engine.fetch(b"key1").unwrap(), b"v2");
assert_eq!(engine.record_count(), 1);
}
#[test]
fn test_tdb_delete() {
let engine = TdbEngine::new();
engine.store(b"key1".to_vec(), b"v1".to_vec()).unwrap();
engine.delete(b"key1").unwrap();
assert!(!engine.exists(b"key1"));
assert_eq!(engine.record_count(), 0);
assert!(engine.fetch(b"key1").is_err());
}
#[test]
fn test_tdb_not_found() {
let engine = TdbEngine::new();
assert!(engine.fetch(b"missing").is_err());
assert!(engine.delete(b"missing").is_err());
}
#[test]
fn test_tdb_multiple_records() {
let engine = TdbEngine::new();
for i in 0..100 {
engine.store(
format!("key{}", i).into_bytes(),
format!("value{}", i).into_bytes(),
).unwrap();
}
assert_eq!(engine.record_count(), 100);
assert_eq!(engine.fetch(b"key50").unwrap(), b"value50");
}
#[test]
fn test_tdb_persist_to_disk() {
let tmp = TempDir::new().unwrap();
let db_path = tmp.path().join("test.tdb");
{
let engine = TdbEngine::open(&db_path).unwrap();
engine.store(b"persistent".to_vec(), b"yes".to_vec()).unwrap();
engine.store(b"key2".to_vec(), b"val2".to_vec()).unwrap();
engine.flush().unwrap();
}
let engine2 = TdbEngine::open(&db_path).unwrap();
assert_eq!(engine2.fetch(b"persistent").unwrap(), b"yes");
assert_eq!(engine2.fetch(b"key2").unwrap(), b"val2");
assert_eq!(engine2.record_count(), 2);
}
#[test]
fn test_tdb_keys() {
let engine = TdbEngine::new();
engine.store(b"a".to_vec(), b"1".to_vec()).unwrap();
engine.store(b"b".to_vec(), b"2".to_vec()).unwrap();
engine.store(b"c".to_vec(), b"3".to_vec()).unwrap();
let mut keys = engine.keys();
keys.sort();
assert_eq!(keys, vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec()]);
}
#[test]
fn test_tdb_iter() {
let engine = TdbEngine::new();
engine.store(b"k1".to_vec(), b"d1".to_vec()).unwrap();
engine.store(b"k2".to_vec(), b"d2".to_vec()).unwrap();
let entries: Vec<_> = engine.iter();
assert_eq!(entries.len(), 2);
}
#[test]
fn test_tdb_stats() {
let engine = TdbEngine::new();
engine.store(b"key1".to_vec(), b"12345".to_vec()).unwrap();
engine.store(b"key2".to_vec(), b"67890".to_vec()).unwrap();
let stats = engine.stats();
assert_eq!(stats.active_records, 2);
assert_eq!(stats.total_data_size, 10);
assert_eq!(stats.avg_key_size, 4.0);
assert_eq!(stats.avg_data_size, 5.0);
}
#[test]
fn test_tdb_create_with_hash_size() {
let engine = TdbEngine::with_hash_size(256);
let header = engine.header.read().unwrap();
assert_eq!(header.hash_size, 256);
}
#[test]
fn test_tdb_transaction_commit() {
let engine = TdbEngine::new();
engine.store(b"init".to_vec(), b"data".to_vec()).unwrap();
{
let mut tx = TdbTransaction::begin(&engine);
tx.store(b"tx_key".to_vec(), b"tx_data".to_vec()).unwrap();
tx.commit().unwrap();
}
assert!(engine.exists(b"tx_key"));
assert_eq!(engine.fetch(b"tx_key").unwrap(), b"tx_data");
}
#[test]
fn test_tdb_transaction_rollback() {
let engine = TdbEngine::new();
engine.store(b"keep".to_vec(), b"yes".to_vec()).unwrap();
{
let mut tx = TdbTransaction::begin(&engine);
tx.store(b"temp".to_vec(), b"no".to_vec()).unwrap();
tx.delete(b"keep").unwrap();
tx.rollback().unwrap();
}
assert!(!engine.exists(b"temp"));
assert!(engine.exists(b"keep"));
}
#[test]
fn test_tdb_transaction_auto_rollback_on_drop() {
let engine = TdbEngine::new();
engine.store(b"orig".to_vec(), b"val".to_vec()).unwrap();
{
let _tx = TdbTransaction::begin(&engine);
_tx.store(b"temp".to_vec(), b"x".to_vec()).unwrap();
}
assert!(!engine.exists(b"temp"));
assert!(engine.exists(b"orig"));
}
#[test]
fn test_hash_key_distribution() {
let hash_size = 1024u32;
let mut buckets = std::collections::HashSet::new();
for i in 0..1000 {
let key = format!("key_{}", i);
let bucket = hash_key(key.as_bytes(), hash_size);
buckets.insert(bucket);
}
assert!(buckets.len() > 200);
}
#[test]
fn test_tdb_corrupt_file_detection() {
let tmp = TempDir::new().unwrap();
let db_path = tmp.path().join("corrupt.tdb");
std::fs::write(&db_path, b"NOT_A_TDB_FILE_GARBAGE_DATA_HERE").unwrap();
assert!(TdbEngine::open(&db_path).is_err());
}
#[test]
fn test_record_flag_conversion() {
assert_eq!(RecordFlag::Active.as_u32(), 0);
assert_eq!(RecordFlag::Free.as_u32(), 1);
assert_eq!(RecordFlag::Deleted.as_u32(), 2);
assert_eq!(RecordFlag::from_u32(0), RecordFlag::Active);
assert_eq!(RecordFlag::from_u32(1), RecordFlag::Free);
assert_eq!(RecordFlag::from_u32(2), RecordFlag::Deleted);
}
#[test]
fn test_tdb_empty_keys_and_values() {
let engine = TdbEngine::new();
engine.store(vec![], vec![]).unwrap();
assert!(engine.exists(&[]));
assert_eq!(engine.fetch(&[]).unwrap(), b"");
}
#[test]
fn test_tdb_large_value() {
let engine = TdbEngine::new();
let large_data = vec![0x42u8; 100_000];
engine.store(b"big".to_vec(), large_data.clone()).unwrap();
let fetched = engine.fetch(b"big").unwrap();
assert_eq!(fetched.len(), 100_000);
assert_eq!(fetched, large_data);
}
#[test]
fn test_tdb_binary_keys() {
let engine = TdbEngine::new();
let key = vec![0u8, 1u8, 255u8, 128u8, 64u8];
engine.store(key.clone(), b"binary_val".to_vec()).unwrap();
assert_eq!(engine.fetch(&key).unwrap(), b"binary_val");
}
}
+1
View File
@@ -5,6 +5,7 @@ pub mod audit;
pub mod auth; pub mod auth;
pub mod category_view; pub mod category_view;
pub mod cli; pub mod cli;
pub mod ctdb;
pub mod command; pub mod command;
pub mod config; pub mod config;
pub mod download; pub mod download;
+17
View File
@@ -260,6 +260,20 @@ pub async fn run(port: u16, file: Option<String>) -> anyhow::Result<()> {
.route("/api/v2/metrics", get(metrics_handler)) .route("/api/v2/metrics", get(metrics_handler))
.route("/api/v2/audit", get(audit_handler)) .route("/api/v2/audit", get(audit_handler))
// Upload API endpoints (unlimited file upload) // Upload API endpoints (unlimited file upload)
// Product API endpoints
.route("/api/v2/products", get(crate::download::list_all_products))
.route("/api/v2/products/stats", get(crate::download::get_series_stats))
.route("/api/v2/products/create", post(crate::download::create_product_handler))
.route("/api/v2/products/by-series/:series", get(crate::download::list_products_by_series))
.route("/api/v2/products/:product_id", delete(crate::download::delete_product))
.route("/api/v2/products/:product_id/files", get(crate::download::get_product_files))
.route("/api/v2/products/:product_id/assign-files", post(crate::download::assign_files_to_product))
.route("/api/v2/download/:file_id", get(crate::download::download_file))
.route("/api/v2/download/:user_id/*file_path", get(crate::download::download_file_by_path))
.route("/api/v2/download/products/:product_series/*file_path", get(crate::download::download_product_file))
.route("/api/v2/download/stats", get(crate::download::get_download_stats))
.route("/api/v2/files/:user_id", get(crate::download::list_uploaded_files))
.route("/api/v2/files/:user_id/:filename", get(crate::download::get_file_info))
.route("/api/v2/upload-unlimited/:user_id", post(upload_unlimited)) .route("/api/v2/upload-unlimited/:user_id", post(upload_unlimited))
// Category View API endpoints (Phase 1: 双视图管理) // Category View API endpoints (Phase 1: 双视图管理)
.route("/api/v2/categories", get(get_all_categories_handler)) .route("/api/v2/categories", get(get_all_categories_handler))
@@ -268,6 +282,9 @@ pub async fn run(port: u16, file: Option<String>) -> anyhow::Result<()> {
.route("/api/v2/series/:series_name", get(get_series_detail_handler)) .route("/api/v2/series/:series_name", get(get_series_detail_handler))
.route("/api/v2/files/search", get(search_files_handler)) .route("/api/v2/files/search", get(search_files_handler))
.route("/api/v2/health", get(health_handler)) .route("/api/v2/health", get(health_handler))
.route("/upload", get(|| async { Html(include_str!("upload.html")) }))
.route("/files", get(|| async { Html(include_str!("file_list.html")) }))
.route("/products", get(|| async { Html(include_str!("product_manager.html")) }))
.route("/downloads", get(|| async { Html(include_str!("category_view.html")) })) .route("/downloads", get(|| async { Html(include_str!("category_view.html")) }))
// WebDAV API endpoints (Phase 20, multi-user P1) // WebDAV API endpoints (Phase 20, multi-user P1)
.route("/webdav", any(handle_webdav_multi)) .route("/webdav", any(handle_webdav_multi))
+59
View File
@@ -582,6 +582,65 @@ impl VfsBackend for LocalFs {
acl.aces.remove(ace_index); acl.aces.remove(ace_index);
self.set_acl(path, &acl) self.set_acl(path, &acl)
} }
// ===== Extended Attributes (xattr) support =====
fn get_xattr(&self, path: &Path, name: &str) -> Result<Vec<u8>, VfsError> {
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
let _meta = path.metadata().map_err(|e| util::map_io_error(path, e))?;
xattr::get(path, name)
.map_err(|e| VfsError::Io(e.to_string()))?
.map(|v| v.to_vec())
.ok_or_else(|| VfsError::NotFound(format!("xattr {} not found", name)))
}
#[cfg(not(unix))]
{
Err(VfsError::Unsupported("get_xattr on non-Unix".to_string()))
}
}
fn set_xattr(&self, path: &Path, name: &str, value: &[u8]) -> Result<(), VfsError> {
#[cfg(unix)]
{
xattr::set(path, name, value)
.map_err(|e| VfsError::Io(e.to_string()))
}
#[cfg(not(unix))]
{
Err(VfsError::Unsupported("set_xattr on non-Unix".to_string()))
}
}
fn remove_xattr(&self, path: &Path, name: &str) -> Result<(), VfsError> {
#[cfg(unix)]
{
xattr::remove(path, name)
.map_err(|e| VfsError::Io(e.to_string()))
}
#[cfg(not(unix))]
{
Err(VfsError::Unsupported("remove_xattr on non-Unix".to_string()))
}
}
fn list_xattrs(&self, path: &Path) -> Result<Vec<String>, VfsError> {
#[cfg(unix)]
{
xattr::list(path)
.map_err(|e| VfsError::Io(e.to_string()))?
.map(|s| s.to_string_lossy().into_owned())
.collect::<Vec<_>>()
.into_iter()
.map(Result::Ok)
.collect()
}
#[cfg(not(unix))]
{
Err(VfsError::Unsupported("list_xattrs on non-Unix".to_string()))
}
}
} }
impl LocalFs { impl LocalFs {
+22
View File
@@ -327,6 +327,28 @@ pub trait VfsBackend: Send + Sync {
fn remove_ace(&self, _path: &Path, _ace_index: usize) -> Result<(), VfsError> { fn remove_ace(&self, _path: &Path, _ace_index: usize) -> Result<(), VfsError> {
Err(VfsError::Unsupported("remove_ace".to_string())) Err(VfsError::Unsupported("remove_ace".to_string()))
} }
// ===== Extended Attributes (xattr) support =====
/// 获取扩展属性
fn get_xattr(&self, _path: &Path, _name: &str) -> Result<Vec<u8>, VfsError> {
Err(VfsError::Unsupported("get_xattr".to_string()))
}
/// 设置扩展属性
fn set_xattr(&self, _path: &Path, _name: &str, _value: &[u8]) -> Result<(), VfsError> {
Err(VfsError::Unsupported("set_xattr".to_string()))
}
/// 删除扩展属性
fn remove_xattr(&self, _path: &Path, _name: &str) -> Result<(), VfsError> {
Err(VfsError::Unsupported("remove_xattr".to_string()))
}
/// 列出扩展属性名称
fn list_xattrs(&self, _path: &Path) -> Result<Vec<String>, VfsError> {
Err(VfsError::Unsupported("list_xattrs".to_string()))
}
} }
/// 快照信息 /// 快照信息
+19 -1
View File
@@ -45,6 +45,8 @@ pub struct Share {
pub(crate) backend: Arc<dyn ShareBackend>, pub(crate) backend: Arc<dyn ShareBackend>,
pub(crate) mode: ShareMode, pub(crate) mode: ShareMode,
pub(crate) users: HashMap<String, Access>, pub(crate) users: HashMap<String, Access>,
pub(crate) time_machine: bool,
pub(crate) time_machine_max_size: Option<u64>,
} }
impl Share { impl Share {
@@ -55,6 +57,8 @@ impl Share {
backend: Arc::new(backend), backend: Arc::new(backend),
mode: ShareMode::AuthenticatedOnly, mode: ShareMode::AuthenticatedOnly,
users: HashMap::new(), users: HashMap::new(),
time_machine: false,
time_machine_max_size: None,
} }
} }
@@ -76,6 +80,20 @@ impl Share {
self.users.insert(name.into(), access); self.users.insert(name.into(), access);
self self
} }
/// Enable Time Machine support for this share.
/// macOS clients will be able to use this share for backups.
pub fn time_machine(mut self) -> Self {
self.time_machine = true;
self
}
/// Set maximum backup size in GB for Time Machine.
pub fn time_machine_max_size(mut self, max_size_gb: u64) -> Self {
self.time_machine = true;
self.time_machine_max_size = Some(max_size_gb * 1024 * 1024 * 1024);
self
}
} }
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
@@ -228,7 +246,7 @@ impl SmbServerBuilder {
let mut share_bindings: Vec<Arc<ShareBindings>> = Vec::with_capacity(self.shares.len()); let mut share_bindings: Vec<Arc<ShareBindings>> = Vec::with_capacity(self.shares.len());
for s in self.shares { for s in self.shares {
share_bindings.push(ShareBindings::new( share_bindings.push(ShareBindings::new(
s.name, s.backend, s.mode, s.users, false, s.name, s.backend, s.mode, s.users, false, s.time_machine, s.time_machine_max_size,
)); ));
} }
+97 -3
View File
@@ -70,6 +70,44 @@ pub async fn handle(
Some(u) => u, Some(u) => u,
None => return HandlerResponse::err(ntstatus::STATUS_OBJECT_NAME_INVALID), None => return HandlerResponse::err(ntstatus::STATUS_OBJECT_NAME_INVALID),
}; };
// Check for named stream (colon separator)
let has_named_stream = units.iter().any(|&u| u == ':' as u16);
if has_named_stream {
use crate::named_stream::NamedStreamPath;
let stream_path = match NamedStreamPath::parse_from_utf16(&units) {
Ok(p) => p,
Err(_) => return HandlerResponse::err(ntstatus::STATUS_OBJECT_NAME_INVALID),
};
// Handle AFP_AfpInfo named stream
if stream_path.is_afp_info() {
debug!(base_path = %stream_path.base_path(), stream = %stream_path.stream_name(), "AFP_AfpInfo named stream open");
// For AFP_AfpInfo, we return a virtual handle that reads/writes extended attributes
// TODO: Implement actual AFP_AfpInfo handling via extended attributes
// Return STATUS_OBJECT_NAME_NOT_FOUND for now (phase 2.6 will implement)
return HandlerResponse::err(ntstatus::STATUS_OBJECT_NAME_NOT_FOUND);
}
// Handle AFP_Resource named stream
if stream_path.is_afp_resource() {
debug!(base_path = %stream_path.base_path(), stream = %stream_path.stream_name(), "AFP_Resource named stream open");
// For AFP_Resource, we return a virtual handle that reads/writes ._ file
// TODO: Implement actual AFP_Resource handling via AppleDouble files
// Return STATUS_OBJECT_NAME_NOT_FOUND for now (phase 3 will implement)
return HandlerResponse::err(ntstatus::STATUS_OBJECT_NAME_NOT_FOUND);
}
// Unknown named stream type
warn!(stream = %stream_path.stream_name(), "unknown named stream type");
return HandlerResponse::err(ntstatus::STATUS_OBJECT_NAME_INVALID);
}
let path = match SmbPath::from_utf16(&units) { let path = match SmbPath::from_utf16(&units) {
Ok(p) => p, Ok(p) => p,
Err(_) => return HandlerResponse::err(ntstatus::STATUS_OBJECT_NAME_INVALID), Err(_) => return HandlerResponse::err(ntstatus::STATUS_OBJECT_NAME_INVALID),
@@ -248,11 +286,67 @@ pub async fn handle(
tree.opens.write().await.insert(file_id, open_arc.clone()); tree.opens.write().await.insert(file_id, open_arc.clone());
drop(tree); drop(tree);
// Phase AAPL: Check for AAPL context (Apple SMB Extensions)
let aapl_response_data = if !req.create_contexts.is_empty() {
use crate::proto::messages::CreateContext;
use crate::proto::messages::{
AaplCreateContextRequest, AaplCreateContextResponse,
SMB2_CRTCTX_AAPL_SERVER_QUERY, SMB2_CRTCTX_AAPL_SERVER_CAPS,
SMB2_CRTCTX_AAPL_VOLUME_CAPS, SMB2_CRTCTX_AAPL_MODEL_INFO,
SMB2_CRTCTX_AAPL_UNIX_BASED, SMB2_CRTCTX_AAPL_SUPPORTS_READ_DIR_ATTR,
SMB2_CRTCTX_AAPL_CASE_SENSITIVE,
};
let contexts = CreateContext::parse_chain(&req.create_contexts).unwrap_or_default();
let aapl_ctx = contexts.iter().find(|ctx| ctx.name == CreateContext::NAME_AAPL);
if let Some(ctx) = aapl_ctx {
if let Some(aapl_req) = AaplCreateContextRequest::from_bytes(&ctx.data) {
if aapl_req.command == SMB2_CRTCTX_AAPL_SERVER_QUERY {
let server_caps = SMB2_CRTCTX_AAPL_UNIX_BASED | SMB2_CRTCTX_AAPL_SUPPORTS_READ_DIR_ATTR;
let volume_caps = SMB2_CRTCTX_AAPL_CASE_SENSITIVE;
let aapl_resp = AaplCreateContextResponse::new_server_query(
aapl_req.request_bitmap,
aapl_req.client_caps,
server_caps,
volume_caps,
"MarkBase SMB",
);
Some(aapl_resp.to_bytes())
} else {
None
}
} else {
None
}
} else {
None
}
} else {
None
};
let create_action = match intent { let create_action = match intent {
OpenIntent::Create => FILE_CREATED, OpenIntent::Create => FILE_CREATED,
OpenIntent::OpenOrCreate | OpenIntent::OverwriteOrCreate => FILE_OPENED, OpenIntent::OpenOrCreate | OpenIntent::OverwriteOrCreate => FILE_OPENED,
OpenIntent::Open | OpenIntent::Truncate => FILE_OPENED, OpenIntent::Open | OpenIntent::Truncate => FILE_OPENED,
}; };
// Build response with AAPL context if present
let (create_contexts_offset, create_contexts_length, create_contexts) = if let Some(data) = aapl_response_data {
use crate::proto::messages::CreateContext;
let aapl_ctx = CreateContext {
name: CreateContext::NAME_AAPL.to_vec(),
data,
};
let mut ctx_buf = Vec::new();
CreateContext::encode_chain(&[aapl_ctx], &mut ctx_buf).expect("encode AAPL context");
let offset = 80 + req.name.len() + 8; // After CreateResponse fixed + padding
(offset as u32, ctx_buf.len() as u32, ctx_buf)
} else {
(0, 0, vec![])
};
let resp = CreateResponse { let resp = CreateResponse {
structure_size: 89, structure_size: 89,
oplock_level: granted_oplock, // Phase 4: will be dynamic oplock_level: granted_oplock, // Phase 4: will be dynamic
@@ -267,9 +361,9 @@ pub async fn handle(
file_attributes: info.attributes(), file_attributes: info.attributes(),
reserved2: 0, reserved2: 0,
file_id, file_id,
create_contexts_offset: 0, create_contexts_offset,
create_contexts_length: 0, create_contexts_length,
create_contexts: vec![], create_contexts,
}; };
let mut buf = Vec::new(); let mut buf = Vec::new();
resp.write_to(&mut buf).expect("encode"); resp.write_to(&mut buf).expect("encode");
+4
View File
@@ -26,18 +26,22 @@ mod error;
mod fs; mod fs;
mod handlers; mod handlers;
pub(crate) mod info_class; pub(crate) mod info_class;
mod named_stream;
pub mod ntstatus; pub mod ntstatus;
mod oplock; mod oplock;
mod path; mod path;
mod proto; mod proto;
mod server; mod server;
mod snapshot; mod snapshot;
mod unicode_mapping;
mod client_restrictions; mod client_restrictions;
mod utils; mod utils;
pub use backend::{BackendCapabilities, DirEntry, FileInfo, FileTimes, Handle, NullHandle, OpenIntent, OpenOptions, ShareBackend}; pub use backend::{BackendCapabilities, DirEntry, FileInfo, FileTimes, Handle, NullHandle, OpenIntent, OpenOptions, ShareBackend};
pub use error::SmbError; pub use error::SmbError;
pub use named_stream::NamedStreamPath;
pub use path::SmbPath; pub use path::SmbPath;
pub use unicode_mapping::{map_ascii_to_private, map_private_to_ascii, has_private_range_chars, has_ntfs_illegal_chars};
pub use builder::{Access, Share}; pub use builder::{Access, Share};
#[cfg(feature = "localfs")] #[cfg(feature = "localfs")]
pub use fs::LocalFsBackend; pub use fs::LocalFsBackend;
+157
View File
@@ -0,0 +1,157 @@
//! SMB Named Stream Path handling
//!
//! Reference: MS-SMB2 §2.2.13 (Named Streams)
//! Named streams are colon-separated: `filename:stream_name:$DATA`
use crate::proto::messages::AFPINFO_STREAM_NAME;
use crate::error::{SmbError, SmbResult};
use crate::path::SmbPath;
pub const STREAM_TYPE_DATA: &[u8] = b"$DATA";
pub const AFP_RESOURCE_STREAM_NAME: &[u8] = b"AFP_Resource";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NamedStreamPath {
base_path: SmbPath,
stream_name: String,
stream_type: String,
}
impl NamedStreamPath {
pub fn new(base_path: SmbPath, stream_name: String, stream_type: String) -> Self {
Self {
base_path,
stream_name,
stream_type,
}
}
pub fn base_path(&self) -> &SmbPath {
&self.base_path
}
pub fn stream_name(&self) -> &str {
&self.stream_name
}
pub fn stream_type(&self) -> &str {
&self.stream_type
}
pub fn is_afp_info(&self) -> bool {
self.stream_name == "AFP_AfpInfo"
}
pub fn is_afp_resource(&self) -> bool {
self.stream_name == "AFP_Resource"
}
pub fn is_data_stream(&self) -> bool {
self.stream_type == "$DATA"
}
pub fn parse_from_utf16(units: &[u16]) -> SmbResult<Self> {
let s = String::from_utf16(units).map_err(|_| SmbError::NameInvalid)?;
Self::parse_from_str(&s)
}
pub fn parse_from_str(s: &str) -> SmbResult<Self> {
let trimmed = s
.strip_prefix('\\')
.or_else(|| s.strip_prefix('/'))
.unwrap_or(s);
if trimmed.is_empty() {
return Err(SmbError::NameInvalid);
}
let colon_pos = trimmed.find(':');
if colon_pos.is_none() {
return Err(SmbError::NameInvalid);
}
let pos = colon_pos.unwrap();
let base_str = &trimmed[..pos];
let rest = &trimmed[pos + 1..];
let second_colon = rest.find(':');
let (stream_name, stream_type) = match second_colon {
Some(pos2) => {
let name = &rest[..pos2];
if name.is_empty() || name == "$DATA" {
return Err(SmbError::NameInvalid);
}
(name, &rest[pos2 + 1..])
}
None => {
if rest.is_empty() || rest == "$DATA" {
return Err(SmbError::NameInvalid);
}
(rest, "$DATA")
}
};
let base_path = base_str.parse::<SmbPath>()?;
Ok(Self {
base_path,
stream_name: stream_name.to_string(),
stream_type: stream_type.to_string(),
})
}
pub fn display(&self) -> String {
format!(
"{}:{}:{}",
self.base_path.display_backslash(),
self.stream_name,
self.stream_type
)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn utf16(s: &str) -> Vec<u16> {
s.encode_utf16().collect()
}
#[test]
fn test_parse_named_stream() {
let p = NamedStreamPath::parse_from_str("file.txt:AFP_AfpInfo:$DATA").unwrap();
assert_eq!(p.base_path().file_name(), Some("file.txt"));
assert_eq!(p.stream_name(), "AFP_AfpInfo");
assert_eq!(p.stream_type(), "$DATA");
assert!(p.is_afp_info());
assert!(p.is_data_stream());
}
#[test]
fn test_parse_without_type() {
let p = NamedStreamPath::parse_from_str("file.txt:AFP_Resource").unwrap();
assert_eq!(p.stream_name(), "AFP_Resource");
assert_eq!(p.stream_type(), "$DATA");
assert!(p.is_afp_resource());
}
#[test]
fn test_parse_from_utf16() {
let units = utf16("file.txt:AFP_AfpInfo:$DATA");
let p = NamedStreamPath::parse_from_utf16(&units).unwrap();
assert!(p.is_afp_info());
}
#[test]
fn test_reject_empty_stream_name() {
assert!(NamedStreamPath::parse_from_str("file.txt::").is_err());
assert!(NamedStreamPath::parse_from_str("file.txt:$DATA").is_err());
}
#[test]
fn test_display() {
let p = NamedStreamPath::parse_from_str("file.txt:AFP_AfpInfo:$DATA").unwrap();
assert_eq!(p.display(), "file.txt:AFP_AfpInfo:$DATA");
}
}
+141
View File
@@ -0,0 +1,141 @@
//! Apple SMB Extensions (AAPL Create Context)
//!
//! Reference: Apple SMB Extensions protocol (MS-SMB2 §2.2.13.2)
//! and Samba vfs_fruit.c implementation.
// AAPL Create Context Command Codes
pub const SMB2_CRTCTX_AAPL_SERVER_QUERY: u32 = 1;
pub const SMB2_CRTCTX_AAPL_RESOLVE_ID: u32 = 2;
// AAPL Server Query Request/Response Bitmap
pub const SMB2_CRTCTX_AAPL_SERVER_CAPS: u64 = 1;
pub const SMB2_CRTCTX_AAPL_VOLUME_CAPS: u64 = 2;
pub const SMB2_CRTCTX_AAPL_MODEL_INFO: u64 = 4;
// AAPL Client/Server Capabilities Bitmap
pub const SMB2_CRTCTX_AAPL_SUPPORTS_READ_DIR_ATTR: u64 = 1;
pub const SMB2_CRTCTX_AAPL_SUPPORTS_OSX_COPYFILE: u64 = 2;
pub const SMB2_CRTCTX_AAPL_UNIX_BASED: u64 = 4;
pub const SMB2_CRTCTX_AAPL_SUPPORTS_NFS_ACE: u64 = 8;
// AAPL Volume Capabilities Bitmap
pub const SMB2_CRTCTX_AAPL_SUPPORT_RESOLVE_ID: u64 = 1;
pub const SMB2_CRTCTX_AAPL_CASE_SENSITIVE: u64 = 2;
pub const SMB2_CRTCTX_AAPL_FULL_SYNC: u64 = 4;
/// AAPL Create Context Request (24 bytes)
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AaplCreateContextRequest {
pub command: u32,
pub reserved: u32,
pub request_bitmap: u64,
pub client_caps: u64,
}
impl AaplCreateContextRequest {
pub fn from_bytes(data: &[u8]) -> Option<Self> {
if data.len() != 24 {
return None;
}
Some(Self {
command: u32::from_le_bytes([data[0], data[1], data[2], data[3]]),
reserved: u32::from_le_bytes([data[4], data[5], data[6], data[7]]),
request_bitmap: u64::from_le_bytes([
data[8], data[9], data[10], data[11],
data[12], data[13], data[14], data[15],
]),
client_caps: u64::from_le_bytes([
data[16], data[17], data[18], data[19],
data[20], data[21], data[22], data[23],
]),
})
}
}
/// AAPL Create Context Response (variable length)
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AaplCreateContextResponse {
pub command: u32,
pub reserved: u32,
pub request_bitmap: u64,
pub server_caps: u64,
pub volume_caps: u64,
pub model: String,
}
impl AaplCreateContextResponse {
pub fn to_bytes(&self) -> Vec<u8> {
let mut buf = Vec::new();
buf.extend_from_slice(&self.command.to_le_bytes());
buf.extend_from_slice(&self.reserved.to_le_bytes());
buf.extend_from_slice(&self.request_bitmap.to_le_bytes());
if self.request_bitmap & SMB2_CRTCTX_AAPL_SERVER_CAPS != 0 {
buf.extend_from_slice(&self.server_caps.to_le_bytes());
}
if self.request_bitmap & SMB2_CRTCTX_AAPL_VOLUME_CAPS != 0 {
buf.extend_from_slice(&self.volume_caps.to_le_bytes());
}
if self.request_bitmap & SMB2_CRTCTX_AAPL_MODEL_INFO != 0 {
let model_utf16: Vec<u16> = self.model.encode_utf16().collect();
buf.extend_from_slice(&(model_utf16.len() as u32 * 2).to_le_bytes());
for ch in model_utf16 {
buf.extend_from_slice(&ch.to_le_bytes());
}
}
buf
}
pub fn new_server_query(
request_bitmap: u64,
client_caps: u64,
server_caps: u64,
volume_caps: u64,
model: &str,
) -> Self {
Self {
command: SMB2_CRTCTX_AAPL_SERVER_QUERY,
reserved: 0,
request_bitmap,
server_caps,
volume_caps,
model: model.to_string(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_aapl_request_parse() {
let data = [
1u8, 0, 0, 0, // command = SERVER_QUERY
0, 0, 0, 0, // reserved
7, 0, 0, 0, 0, 0, 0, 0, // request_bitmap = SERVER_CAPS | VOLUME_CAPS | MODEL_INFO
0, 0, 0, 0, 0, 0, 0, 0, // client_caps
];
let req = AaplCreateContextRequest::from_bytes(&data).unwrap();
assert_eq!(req.command, SMB2_CRTCTX_AAPL_SERVER_QUERY);
assert_eq!(req.request_bitmap, 7);
}
#[test]
fn test_aapl_response_encode() {
let resp = AaplCreateContextResponse::new_server_query(
SMB2_CRTCTX_AAPL_SERVER_CAPS | SMB2_CRTCTX_AAPL_VOLUME_CAPS,
0,
SMB2_CRTCTX_AAPL_UNIX_BASED | SMB2_CRTCTX_AAPL_SUPPORTS_READ_DIR_ATTR,
SMB2_CRTCTX_AAPL_CASE_SENSITIVE,
"MacBookPro",
);
let bytes = resp.to_bytes();
assert!(!bytes.is_empty());
assert_eq!(bytes[0..4], [1, 0, 0, 0]); // command
}
}
+257
View File
@@ -0,0 +1,257 @@
//! AFP_AfpInfo Stream (Apple SMB Extensions)
//!
//! Reference: Apple SMB Extensions protocol (MS-SMB2 §2.2.13.2)
//! and Samba MacExtensions.h implementation.
pub const AFP_INFO_SIZE: usize = 60;
pub const AFP_SIGNATURE: u32 = 0x41465000; // "AFP\0"
pub const AFP_VERSION: u32 = 0x00010000;
pub const AFP_OFF_FINDER_INFO: usize = 16;
pub const AFP_FINDER_SIZE: usize = 32;
pub const AFPINFO_STREAM_NAME: &[u8] = b"AFP_AfpInfo";
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct FinderInfo {
pub file_type: [u8; 4],
pub creator_type: [u8; 4],
pub finder_flags: u16,
pub location: (u16, u16),
pub reserved: u16,
pub extended_flags: u16,
pub extended_location: (u16, u16),
pub reserved2: [u8; 10],
}
impl FinderInfo {
pub fn default_file() -> Self {
Self {
file_type: [0, 0, 0, 0],
creator_type: [0, 0, 0, 0],
finder_flags: 0,
location: (0, 0),
reserved: 0,
extended_flags: 0,
extended_location: (0, 0),
reserved2: [0; 10],
}
}
pub fn default_directory() -> Self {
Self {
file_type: [0, 0, 0, 0],
creator_type: [0, 0, 0, 0],
finder_flags: 0,
location: (0, 0),
reserved: 0,
extended_flags: 0,
extended_location: (0, 0),
reserved2: [0; 10],
}
}
pub fn from_bytes(data: &[u8; 32]) -> Self {
let file_type = [data[0], data[1], data[2], data[3]];
let creator_type = [data[4], data[5], data[6], data[7]];
let finder_flags = u16::from_be_bytes([data[8], data[9]]);
let location = (
u16::from_be_bytes([data[10], data[11]]),
u16::from_be_bytes([data[12], data[13]]),
);
let reserved = u16::from_be_bytes([data[14], data[15]]);
let extended_flags = u16::from_be_bytes([data[16], data[17]]);
let extended_location = (
u16::from_be_bytes([data[18], data[19]]),
u16::from_be_bytes([data[20], data[21]]),
);
let reserved2 = [
data[22], data[23], data[24], data[25],
data[26], data[27], data[28], data[29],
data[30], data[31],
];
Self {
file_type,
creator_type,
finder_flags,
location,
reserved,
extended_flags,
extended_location,
reserved2,
}
}
pub fn to_bytes(&self) -> [u8; 32] {
let mut data = [0u8; 32];
data[0..4].copy_from_slice(&self.file_type);
data[4..8].copy_from_slice(&self.creator_type);
data[8..10].copy_from_slice(&self.finder_flags.to_be_bytes());
data[10..12].copy_from_slice(&self.location.0.to_be_bytes());
data[12..14].copy_from_slice(&self.location.1.to_be_bytes());
data[14..16].copy_from_slice(&self.reserved.to_be_bytes());
data[16..18].copy_from_slice(&self.extended_flags.to_be_bytes());
data[18..20].copy_from_slice(&self.extended_location.0.to_be_bytes());
data[20..22].copy_from_slice(&self.extended_location.1.to_be_bytes());
data[22..32].copy_from_slice(&self.reserved2);
data
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AfpInfo {
pub signature: u32,
pub version: u32,
pub reserved1: u32,
pub backup_time: u32,
pub finder_info: FinderInfo,
pub prodos_info: [u8; 6],
pub reserved2: [u8; 6],
}
impl AfpInfo {
pub fn new() -> Self {
Self {
signature: AFP_SIGNATURE,
version: AFP_VERSION,
reserved1: 0,
backup_time: 0,
finder_info: FinderInfo::default_file(),
prodos_info: [0; 6],
reserved2: [0; 6],
}
}
pub fn from_bytes(data: &[u8]) -> Option<Self> {
if data.len() < AFP_INFO_SIZE {
return None;
}
let signature = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
if signature != AFP_SIGNATURE {
return None;
}
let version = u32::from_le_bytes([data[4], data[5], data[6], data[7]]);
if version != AFP_VERSION {
return None;
}
let reserved1 = u32::from_le_bytes([data[8], data[9], data[10], data[11]]);
let backup_time = u32::from_le_bytes([data[12], data[13], data[14], data[15]]);
let finder_bytes: [u8; 32] = data[16..48].try_into().ok()?;
let finder_info = FinderInfo::from_bytes(&finder_bytes);
let prodos_info: [u8; 6] = data[48..54].try_into().ok()?;
let reserved2: [u8; 6] = data[54..60].try_into().ok()?;
Some(Self {
signature,
version,
reserved1,
backup_time,
finder_info,
prodos_info,
reserved2,
})
}
pub fn to_bytes(&self) -> Vec<u8> {
let mut data = Vec::with_capacity(AFP_INFO_SIZE);
data.extend_from_slice(&self.signature.to_le_bytes());
data.extend_from_slice(&self.version.to_le_bytes());
data.extend_from_slice(&self.reserved1.to_le_bytes());
data.extend_from_slice(&self.backup_time.to_le_bytes());
data.extend_from_slice(&self.finder_info.to_bytes());
data.extend_from_slice(&self.prodos_info);
data.extend_from_slice(&self.reserved2);
data
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SambaAfpInfo {
pub afp: AfpInfo,
pub create_time: u32,
}
impl SambaAfpInfo {
pub fn new() -> Self {
Self {
afp: AfpInfo::new(),
create_time: 0,
}
}
pub fn from_bytes(data: &[u8]) -> Option<Self> {
if data.len() < AFP_INFO_SIZE + 4 {
return None;
}
let afp = AfpInfo::from_bytes(&data[..AFP_INFO_SIZE])?;
let create_time = u32::from_le_bytes([
data[60], data[61], data[62], data[63],
]);
Some(Self {
afp,
create_time,
})
}
pub fn to_bytes(&self) -> Vec<u8> {
let mut data = self.afp.to_bytes();
data.extend_from_slice(&self.create_time.to_le_bytes());
data
}
pub fn get_create_time(&self) -> u32 {
self.create_time
}
pub fn set_create_time(&mut self, time: u32) {
self.create_time = time;
}
pub fn get_finder_info(&self) -> &FinderInfo {
&self.afp.finder_info
}
pub fn set_finder_info(&mut self, finder: FinderInfo) {
self.afp.finder_info = finder;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_afp_info_roundtrip() {
let afp = AfpInfo::new();
let bytes = afp.to_bytes();
assert_eq!(bytes.len(), AFP_INFO_SIZE);
let decoded = AfpInfo::from_bytes(&bytes).unwrap();
assert_eq!(decoded.signature, AFP_SIGNATURE);
assert_eq!(decoded.version, AFP_VERSION);
}
#[test]
fn test_samba_afp_info_roundtrip() {
let mut samba = SambaAfpInfo::new();
samba.set_create_time(12345678);
let bytes = samba.to_bytes();
assert_eq!(bytes.len(), AFP_INFO_SIZE + 4);
let decoded = SambaAfpInfo::from_bytes(&bytes).unwrap();
assert_eq!(decoded.get_create_time(), 12345678);
}
#[test]
fn test_finder_info_roundtrip() {
let finder = FinderInfo::default_file();
let bytes = finder.to_bytes();
assert_eq!(bytes.len(), 32);
let decoded = FinderInfo::from_bytes(&bytes);
assert_eq!(decoded.file_type, finder.file_type);
assert_eq!(decoded.creator_type, finder.creator_type);
}
#[test]
fn test_invalid_signature() {
let data = [0u8; 60];
assert!(AfpInfo::from_bytes(&data).is_none());
}
}
+1
View File
@@ -171,6 +171,7 @@ impl CreateContext {
pub const NAME_RQLS: &'static [u8; 4] = b"RqLs"; // REQUEST_LEASE pub const NAME_RQLS: &'static [u8; 4] = b"RqLs"; // REQUEST_LEASE
pub const NAME_DH2Q: &'static [u8; 4] = b"DH2Q"; // DURABLE_HANDLE_REQUEST_V2 pub const NAME_DH2Q: &'static [u8; 4] = b"DH2Q"; // DURABLE_HANDLE_REQUEST_V2
pub const NAME_DH2C: &'static [u8; 4] = b"DH2C"; // DURABLE_HANDLE_RECONNECT_V2 pub const NAME_DH2C: &'static [u8; 4] = b"DH2C"; // DURABLE_HANDLE_RECONNECT_V2
pub const NAME_AAPL: &'static [u8; 4] = b"AAPL"; // Apple SMB Extensions (MS-SMB2 §2.2.13.2)
/// Parse a chain of create-contexts from the raw chain bytes. /// Parse a chain of create-contexts from the raw chain bytes.
/// ///
+14
View File
@@ -7,6 +7,8 @@
//! The crate does **not** implement command behavior — it only encodes/decodes //! The crate does **not** implement command behavior — it only encodes/decodes
//! the wire bytes. The server crate owns dispatch and state. //! the wire bytes. The server crate owns dispatch and state.
pub mod aapl;
pub mod afp_info;
pub mod cancel; pub mod cancel;
pub mod change_notify; pub mod change_notify;
pub mod close; pub mod close;
@@ -29,6 +31,18 @@ pub mod tree_connect;
pub mod tree_disconnect; pub mod tree_disconnect;
pub mod write; pub mod write;
pub use aapl::{
AaplCreateContextRequest, AaplCreateContextResponse, SMB2_CRTCTX_AAPL_CASE_SENSITIVE,
SMB2_CRTCTX_AAPL_FULL_SYNC, SMB2_CRTCTX_AAPL_MODEL_INFO, SMB2_CRTCTX_AAPL_SERVER_CAPS,
SMB2_CRTCTX_AAPL_SERVER_QUERY, SMB2_CRTCTX_AAPL_SUPPORT_RESOLVE_ID,
SMB2_CRTCTX_AAPL_SUPPORTS_NFS_ACE, SMB2_CRTCTX_AAPL_SUPPORTS_OSX_COPYFILE,
SMB2_CRTCTX_AAPL_SUPPORTS_READ_DIR_ATTR, SMB2_CRTCTX_AAPL_UNIX_BASED,
SMB2_CRTCTX_AAPL_VOLUME_CAPS,
};
pub use afp_info::{
AfpInfo, FinderInfo, SambaAfpInfo, AFP_FINDER_SIZE, AFP_INFO_SIZE, AFP_OFF_FINDER_INFO,
AFP_SIGNATURE, AFP_VERSION, AFPINFO_STREAM_NAME,
};
pub use cancel::CancelRequest; pub use cancel::CancelRequest;
pub use change_notify::{ChangeNotifyRequest, ChangeNotifyResponse}; pub use change_notify::{ChangeNotifyRequest, ChangeNotifyResponse};
pub use close::{CloseRequest, CloseResponse}; pub use close::{CloseRequest, CloseResponse};
+10 -1
View File
@@ -47,6 +47,9 @@ pub struct ShareBindings {
/// (Windows always probes IPC$ before mounting an actual share). All /// (Windows always probes IPC$ before mounting an actual share). All
/// downstream ops on an IPC$ tree return `STATUS_NOT_SUPPORTED`. /// downstream ops on an IPC$ tree return `STATUS_NOT_SUPPORTED`.
pub is_ipc: bool, pub is_ipc: bool,
/// Time Machine support for macOS backups.
pub time_machine: bool,
pub time_machine_max_size: Option<u64>,
} }
impl ShareBindings { impl ShareBindings {
@@ -56,12 +59,16 @@ impl ShareBindings {
mode: ShareMode, mode: ShareMode,
users: HashMap<String, Access>, users: HashMap<String, Access>,
is_ipc: bool, is_ipc: bool,
time_machine: bool,
time_machine_max_size: Option<u64>,
) -> Arc<Self> { ) -> Arc<Self> {
Arc::new(Self { Arc::new(Self {
name, name,
backend, backend,
acl: RwLock::new(ShareAcl { mode, users }), acl: RwLock::new(ShareAcl { mode, users }),
is_ipc, is_ipc,
time_machine,
time_machine_max_size,
}) })
} }
@@ -74,6 +81,8 @@ impl ShareBindings {
ShareMode::PublicReadOnly, ShareMode::PublicReadOnly,
HashMap::new(), HashMap::new(),
true, true,
false,
None,
) )
} }
} }
@@ -311,7 +320,7 @@ impl ConfigHandle {
} }
} }
let binding = ShareBindings::new(share.name, share.backend, share.mode, share.users, false); let binding = ShareBindings::new(share.name, share.backend, share.mode, share.users, false, share.time_machine, share.time_machine_max_size);
self.state.shares.insert(binding).await self.state.shares.insert(binding).await
} }
+123
View File
@@ -0,0 +1,123 @@
//! macOS Unicode Private Range Mapping for SMB
//!
//! macOS SMB client maps NTFS illegal characters to Unicode private range.
//! Reference: Samba vfs_fruit.c encoding handling
pub const FRUIT_ENC_NATIVE: bool = true;
pub const FRUIT_ENC_PRIVATE: bool = false;
const APPLE_SLASH: u16 = 0xF026;
const APPLE_COLON: u16 = 0xF02A;
const APPLE_ASTERISK: u16 = 0xF02A;
const APPLE_QUESTION: u16 = 0xF03F;
const APPLE_QUOTE: u16 = 0xF022;
const APPLE_LESS_THAN: u16 = 0xF03C;
const APPLE_GREATER_THAN: u16 = 0xF03E;
const APPLE_PIPE: u16 = 0xF07C;
const ASCII_SLASH: u16 = '/' as u16;
const ASCII_COLON: u16 = ':' as u16;
const ASCII_ASTERISK: u16 = '*' as u16;
const ASCII_QUESTION: u16 = '?' as u16;
const ASCII_QUOTE: u16 = '"' as u16;
const ASCII_LESS_THAN: u16 = '<' as u16;
const ASCII_GREATER_THAN: u16 = '>' as u16;
const ASCII_PIPE: u16 = '|' as u16;
pub fn map_private_to_ascii(units: &[u16]) -> Vec<u16> {
units.iter().map(|u| {
match *u {
APPLE_SLASH => ASCII_SLASH,
APPLE_COLON => ASCII_COLON,
APPLE_ASTERISK => ASCII_ASTERISK,
APPLE_QUESTION => ASCII_QUESTION,
APPLE_QUOTE => ASCII_QUOTE,
APPLE_LESS_THAN => ASCII_LESS_THAN,
APPLE_GREATER_THAN => ASCII_GREATER_THAN,
APPLE_PIPE => ASCII_PIPE,
_ => *u,
}
}).collect()
}
pub fn map_ascii_to_private(units: &[u16]) -> Vec<u16> {
units.iter().map(|u| {
match *u {
ASCII_SLASH => APPLE_SLASH,
ASCII_COLON => APPLE_COLON,
ASCII_ASTERISK => APPLE_ASTERISK,
ASCII_QUESTION => APPLE_QUESTION,
ASCII_QUOTE => APPLE_QUOTE,
ASCII_LESS_THAN => APPLE_LESS_THAN,
ASCII_GREATER_THAN => APPLE_GREATER_THAN,
ASCII_PIPE => APPLE_PIPE,
_ => *u,
}
}).collect()
}
pub fn has_private_range_chars(units: &[u16]) -> bool {
units.iter().any(|u| {
matches!(*u,
APPLE_SLASH | APPLE_COLON | APPLE_ASTERISK |
APPLE_QUESTION | APPLE_QUOTE | APPLE_LESS_THAN |
APPLE_GREATER_THAN | APPLE_PIPE
)
})
}
pub fn has_ntfs_illegal_chars(units: &[u16]) -> bool {
units.iter().any(|u| {
matches!(*u,
ASCII_SLASH | ASCII_COLON | ASCII_ASTERISK |
ASCII_QUESTION | ASCII_QUOTE | ASCII_LESS_THAN |
ASCII_GREATER_THAN | ASCII_PIPE
)
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_map_private_to_ascii() {
let input = [APPLE_SLASH, APPLE_COLON, APPLE_QUESTION];
let output = map_private_to_ascii(&input);
assert_eq!(output, [ASCII_SLASH, ASCII_COLON, ASCII_QUESTION]);
}
#[test]
fn test_map_ascii_to_private() {
let input = [ASCII_SLASH, ASCII_COLON, ASCII_ASTERISK];
let output = map_ascii_to_private(&input);
assert_eq!(output, [APPLE_SLASH, APPLE_COLON, APPLE_ASTERISK]);
}
#[test]
fn test_roundtrip() {
let original = [ASCII_SLASH, ASCII_COLON, 'a' as u16];
let to_private = map_ascii_to_private(&original);
let back_to_ascii = map_private_to_ascii(&to_private);
assert_eq!(back_to_ascii, original);
}
#[test]
fn test_has_private_range_chars() {
assert!(has_private_range_chars(&[APPLE_SLASH, 'a' as u16]));
assert!(!has_private_range_chars(&[ASCII_SLASH, 'a' as u16]));
}
#[test]
fn test_has_ntfs_illegal_chars() {
assert!(has_ntfs_illegal_chars(&[ASCII_SLASH, 'a' as u16]));
assert!(!has_ntfs_illegal_chars(&['a' as u16, 'b' as u16]));
}
#[test]
fn test_preserve_non_mapped() {
let input = ['a' as u16, 'b' as u16, 'c' as u16];
let output = map_private_to_ascii(&input);
assert_eq!(output, input);
}
}