diff --git a/crates/store/src/network.rs b/crates/store/src/network.rs index 7de0da3..2421c27 100644 --- a/crates/store/src/network.rs +++ b/crates/store/src/network.rs @@ -1,1338 +1,1535 @@ //! Network management module for storing and managing network configurations. //! //! This module provides functionality to create, read, update, and delete network //! configurations. Each network is defined by a name, network range (IPv4 or IPv6), //! a provider implementation, and an optional assigner implementation. //! //! # Example //! //! ``` //! use store::network::*; //! # use tempfile::TempDir; //! # fn example() -> store::Result<()> { //! # let temp_dir = tempfile::tempdir().unwrap(); //! # let db = sled::open(temp_dir.path())?; //! let store = NetworkStore::open(&db)?; //! //! // Create a network with IPv4 range //! let netrange = NetRange::from_cidr("192.168.1.0/24")?; //! let provider = BasicProvider::new() //! .with_config("type", "aws") //! .with_config("region", "us-west-2"); //! let assigner = Some(BasicAssigner::new("dhcp")); //! //! store.create("production", netrange, provider, assigner)?; //! //! // Retrieve the network //! let network = store.get("production")?.unwrap(); //! assert_eq!(network.name, "production"); //! //! // List all networks //! let networks = store.list()?; //! assert!(networks.contains(&"production".to_string())); //! # Ok(()) //! # } //! ``` use crate::{Result, Error, RangeStore}; use sled::Transactional; use std::collections::HashMap; use std::net::{Ipv4Addr, Ipv6Addr, IpAddr}; /// Represents an IPv4 or IPv6 network range #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] pub enum NetRange { /// IPv4 network with address and prefix length V4 { addr: Ipv4Addr, prefix: u8 }, /// IPv6 network with address and prefix length V6 { addr: Ipv6Addr, prefix: u8 }, } impl NetRange { /// Create a new IPv4 network range pub fn ipv4(addr: Ipv4Addr, prefix: u8) -> Result { if prefix > 32 { return Err(Error::StoreError(sled::Error::Unsupported( format!("Invalid IPv4 prefix length: {}", prefix) ))); } Ok(NetRange::V4 { addr, prefix }) } /// Create a new IPv6 network range pub fn ipv6(addr: Ipv6Addr, prefix: u8) -> Result { if prefix > 128 { return Err(Error::StoreError(sled::Error::Unsupported( format!("Invalid IPv6 prefix length: {}", prefix) ))); } Ok(NetRange::V6 { addr, prefix }) } /// Parse a network range from CIDR notation pub fn from_cidr(cidr: &str) -> Result { let parts: Vec<&str> = cidr.split('/').collect(); if parts.len() != 2 { return Err(Error::StoreError(sled::Error::Unsupported( "Invalid CIDR format".to_string() ))); } let prefix: u8 = parts[1].parse().map_err(|_| { Error::StoreError(sled::Error::Unsupported("Invalid prefix length".to_string())) })?; if let Ok(ipv4) = parts[0].parse::() { Self::ipv4(ipv4, prefix) } else if let Ok(ipv6) = parts[0].parse::() { Self::ipv6(ipv6, prefix) } else { Err(Error::StoreError(sled::Error::Unsupported( "Invalid IP address format".to_string() ))) } } /// Convert to CIDR notation string pub fn to_cidr(&self) -> String { match self { NetRange::V4 { addr, prefix } => format!("{}/{}", addr, prefix), NetRange::V6 { addr, prefix } => format!("{}/{}", addr, prefix), } } } /// Trait for network providers pub trait NetworkProvider: std::fmt::Debug { /// Get the provider type identifier fn provider_type(&self) -> &'static str; /// Serialize provider configuration to JSON fn to_json(&self) -> Result; /// Create provider from JSON fn from_json(json: &str) -> Result where Self: Sized; } /// Trait for network assigners pub trait NetworkAssigner: std::fmt::Debug { /// Get the assigner type identifier fn assigner_type(&self) -> &'static str; /// Serialize assigner configuration to JSON fn to_json(&self) -> Result; /// Create assigner from JSON fn from_json(json: &str) -> Result where Self: Sized; /// Initialize the assigner for a network during creation fn initialize_for_network(&self, _network_name: &str, _netrange: &NetRange, _range_store: Option<&RangeStore>) -> Result<()> { // Default implementation does nothing Ok(()) } /// Initialize the assigner for a network within a transaction fn initialize_for_network_in_transaction( &self, _network_name: &str, _netrange: &NetRange, _range_trees: Option<(&sled::transaction::TransactionalTree, &sled::transaction::TransactionalTree)> ) -> Result<(), sled::transaction::ConflictableTransactionError<()>> { // Default implementation does nothing Ok(()) } + /// Cleanup the assigner for a network during deletion + fn cleanup_for_network(&self, _network_name: &str, _range_store: Option<&RangeStore>) -> Result<()> { + // Default implementation does nothing + Ok(()) + } + + /// Cleanup the assigner for a network within a transaction + fn cleanup_for_network_in_transaction( + &self, + _network_name: &str, + _range_trees: Option<(&sled::transaction::TransactionalTree, &sled::transaction::TransactionalTree)> + ) -> Result<(), sled::transaction::ConflictableTransactionError<()>> { + // Default implementation does nothing + Ok(()) + } + /// Get an IP assignment for the given identifier fn get_assignment(&self, _network: &Network, _identifier: &str, _range_store: Option<&RangeStore>) -> Result> { // Default implementation returns None (no assignment tracking) Ok(None) } } /// Basic network provider implementation #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct BasicProvider { pub config: HashMap, } impl BasicProvider { pub fn new() -> Self { Self { config: HashMap::new(), } } pub fn with_config(mut self, key: &str, value: &str) -> Self { self.config.insert(key.to_string(), value.to_string()); self } } impl NetworkProvider for BasicProvider { fn provider_type(&self) -> &'static str { "basic" } fn to_json(&self) -> Result { serde_json::to_string(self).map_err(|e| { Error::StoreError(sled::Error::Unsupported(format!("JSON serialization error: {}", e))) }) } fn from_json(json: &str) -> Result { serde_json::from_str(json).map_err(|e| { Error::StoreError(sled::Error::Unsupported(format!("JSON deserialization error: {}", e))) }) } } /// Basic network assigner implementation #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct BasicAssigner { pub strategy: String, pub config: HashMap, } impl BasicAssigner { pub fn new(strategy: &str) -> Self { Self { strategy: strategy.to_string(), config: HashMap::new(), } } pub fn with_config(mut self, key: &str, value: &str) -> Self { self.config.insert(key.to_string(), value.to_string()); self } } impl NetworkAssigner for BasicAssigner { fn assigner_type(&self) -> &'static str { "basic" } fn to_json(&self) -> Result { serde_json::to_string(self).map_err(|e| { Error::StoreError(sled::Error::Unsupported(format!("JSON serialization error: {}", e))) }) } fn from_json(json: &str) -> Result { serde_json::from_str(json).map_err(|e| { Error::StoreError(sled::Error::Unsupported(format!("JSON deserialization error: {}", e))) }) } } /// Range-based network assigner that uses RangeStore for IP assignment /// /// The `RangeAssigner` provides IP address assignment functionality by integrating /// with the `RangeStore`. It automatically manages IP address allocation within /// a network range using a bitmap-based approach for efficient tracking. /// /// # Features /// /// - Automatic range initialization when creating networks /// - Sequential IP assignment within the network range /// - Assignment tracking with custom identifiers /// - IP unassignment and reuse /// - Integration with the NetworkStore lifecycle /// /// # Example /// /// ``` /// use store::network::{RangeAssigner, NetRange, BasicProvider, NetworkStore}; /// # use tempfile::TempDir; /// # fn example() -> store::Result<()> { /// # let temp_dir = tempfile::tempdir().unwrap(); /// # let db = sled::open(temp_dir.path())?; /// # let ranges = store::range::RangeStore::open(&db)?; /// # let store = NetworkStore::open_with_ranges(&db, ranges.clone())?; /// /// // Create a RangeAssigner for a network (uses network name as range name) /// let assigner = RangeAssigner::new() /// .with_config("strategy", "sequential"); /// /// // Create a network with the range assigner /// let netrange = NetRange::from_cidr("192.168.1.0/24")?; /// let provider = BasicProvider::new(); /// store.create("my-network", netrange, provider, Some(assigner))?; /// /// // Use the assigner to allocate IPs /// let range_assigner = RangeAssigner::new(); /// let ip_bit = range_assigner.assign_ip("my-network", &ranges, "device1")?; /// # Ok(()) /// # } /// ``` #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct RangeAssigner { pub range_name: Option, pub config: HashMap, } impl RangeAssigner { /// Create a new RangeAssigner with default settings /// /// The range name will default to the network name when the network is created. /// Use `with_range_name()` to specify a custom range name if needed. pub fn new() -> Self { Self { range_name: None, config: HashMap::new(), } } /// Create a new RangeAssigner with a specific range name /// /// # Arguments /// /// * `range_name` - Custom identifier for the IP range pub fn with_range_name(range_name: &str) -> Self { Self { range_name: Some(range_name.to_string()), config: HashMap::new(), } } /// Add configuration key-value pairs to the assigner /// /// Configuration options can be used to customize the behavior of the /// assigner. Common configuration keys might include: /// - `strategy`: Assignment strategy (e.g., "sequential", "random") /// - `reserve_gateway`: Whether to reserve the first IP for gateway /// - `pool_size`: Maximum number of IPs to manage /// /// # Arguments /// /// * `key` - Configuration key /// * `value` - Configuration value pub fn with_config(mut self, key: &str, value: &str) -> Self { self.config.insert(key.to_string(), value.to_string()); self } /// Initialize the range for this assigner with the given network size /// /// This method calculates the number of available IP addresses in the network /// range and creates a corresponding range in the RangeStore. It's automatically /// called when creating a network with a RangeAssigner. /// /// # Arguments /// /// * `range_name` - Name to use for the range /// * `range_store` - Reference to the RangeStore /// * `network` - Network range (IPv4 or IPv6) to initialize /// /// # Returns /// /// Returns `Ok(())` if the range was successfully initialized, or an error /// if the network parameters are invalid or the range already exists. pub fn initialize_range(&self, range_name: &str, range_store: &RangeStore, network: &NetRange) -> Result<()> { let size = match network { NetRange::V4 { prefix, .. } => { if *prefix >= 32 { return Err(Error::StoreError(sled::Error::Unsupported( "IPv4 prefix must be less than 32".to_string() ))); } 1u64 << (32 - prefix) } NetRange::V6 { prefix, .. } => { if *prefix >= 128 { return Err(Error::StoreError(sled::Error::Unsupported( "IPv6 prefix must be less than 128".to_string() ))); } // For IPv6, we'll limit to a reasonable size to avoid huge ranges let host_bits = 128 - prefix; if host_bits > 32 { 1u64 << 32 // Cap at 2^32 addresses } else { 1u64 << host_bits } } }; range_store.define(range_name, size) } /// Assign an IP address from the range /// /// Assigns the next available IP address in the range to the specified identifier. /// The assignment uses a sequential allocation strategy, finding the first /// available bit position in the range bitmap. /// /// # Arguments /// /// * `range_name` - Name of the range to assign from /// * `range_store` - Reference to the RangeStore /// * `identifier` - Unique identifier for the device/entity receiving the IP /// /// # Returns /// /// Returns the bit position of the assigned IP address, which can be used /// to calculate the actual IP address within the network range. /// /// # Errors /// /// Returns an error if the range is full or doesn't exist. pub fn assign_ip(&self, range_name: &str, range_store: &RangeStore, identifier: &str) -> Result { range_store.assign(range_name, identifier) } /// Get assigned IP information /// /// Retrieves the identifier associated with a specific bit position in the range. /// This is useful for determining which device or entity is assigned to a /// particular IP address. /// /// # Arguments /// /// * `range_name` - Name of the range to query /// * `range_store` - Reference to the RangeStore /// * `bit_position` - The bit position to query /// /// # Returns /// /// Returns `Some(identifier)` if the bit position is assigned, or `None` /// if it's free or out of range. pub fn get_assignment_by_bit(&self, range_name: &str, range_store: &RangeStore, bit_position: u64) -> Result> { range_store.get(range_name, bit_position) } /// Unassign an IP address /// /// Releases an IP address assignment, making it available for future assignments. /// This is typically called when a device is removed or no longer needs its /// assigned IP address. /// /// # Arguments /// /// * `range_name` - Name of the range to unassign from /// * `range_store` - Reference to the RangeStore /// * `bit_position` - The bit position to unassign /// /// # Returns /// /// Returns `true` if the IP was successfully unassigned, or `false` if it /// was already free or the bit position is out of range. pub fn unassign_ip(&self, range_name: &str, range_store: &RangeStore, bit_position: u64) -> Result { range_store.unassign_bit(range_name, bit_position) } /// Get the range name for this assigner, using the network name as default /// /// # Arguments /// /// * `network_name` - Network name to use as default if no custom range name is set /// /// # Returns /// /// Returns the range name to use for operations pub fn get_range_name(&self, network_name: &str) -> String { self.range_name.clone().unwrap_or_else(|| network_name.to_string()) } /// Convert a bit position to an IP address within the network range /// /// # Arguments /// /// * `bit_position` - The bit position within the range /// * `netrange` - The network range to calculate the IP within /// /// # Returns /// /// Returns the IP address corresponding to the bit position pub fn bit_position_to_ip(&self, bit_position: u64, netrange: &NetRange) -> Result { match netrange { NetRange::V4 { addr, prefix: _ } => { let network_addr = u32::from(*addr); let host_addr = network_addr + bit_position as u32; Ok(IpAddr::V4(Ipv4Addr::from(host_addr))) } NetRange::V6 { addr, prefix: _ } => { let network_addr = u128::from(*addr); let host_addr = network_addr + bit_position as u128; Ok(IpAddr::V6(Ipv6Addr::from(host_addr))) } } } /// Find the bit position for a given identifier in the range /// /// # Arguments /// /// * `range_name` - Name of the range to search /// * `range_store` - Reference to the RangeStore /// * `identifier` - The identifier to search for /// /// # Returns /// /// Returns the bit position if found, None otherwise pub fn find_identifier_bit_position(&self, range_name: &str, range_store: &RangeStore, identifier: &str) -> Result> { // Get range ID first let range_id = match range_store.names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(None); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid range data".to_string())))?; u64::from_be_bytes(id_bytes) } None => return Ok(None), }; // Search for the assignment with this identifier let prefix = range_id.to_be_bytes(); for result in range_store.assign.scan_prefix(&prefix) { let (key, stored_value) = result?; if key.len() == 16 { let stored_value_str = String::from_utf8(stored_value.to_vec())?; if stored_value_str == identifier { let bit_position_bytes: [u8; 8] = key[8..16].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid key format".to_string())))?; return Ok(Some(u64::from_be_bytes(bit_position_bytes))); } } } Ok(None) } } impl NetworkAssigner for RangeAssigner { fn assigner_type(&self) -> &'static str { "range" } fn to_json(&self) -> Result { serde_json::to_string(self).map_err(|e| { Error::StoreError(sled::Error::Unsupported(format!("JSON serialization error: {}", e))) }) } fn from_json(json: &str) -> Result { serde_json::from_str(json).map_err(|e| { Error::StoreError(sled::Error::Unsupported(format!("JSON deserialization error: {}", e))) }) } fn initialize_for_network(&self, network_name: &str, netrange: &NetRange, range_store: Option<&RangeStore>) -> Result<()> { if let Some(store) = range_store { let range_name = self.get_range_name(network_name); self.initialize_range(&range_name, store, netrange)?; } Ok(()) } fn initialize_for_network_in_transaction( &self, network_name: &str, netrange: &NetRange, range_trees: Option<(&sled::transaction::TransactionalTree, &sled::transaction::TransactionalTree)> ) -> Result<(), sled::transaction::ConflictableTransactionError<()>> { if let Some((range_names, range_map)) = range_trees { let range_name = self.get_range_name(network_name); let size = match netrange { NetRange::V4 { prefix, .. } => { if *prefix >= 32 { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } 1u64 << (32 - prefix) } NetRange::V6 { prefix, .. } => { if *prefix >= 128 { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } // For IPv6, we'll limit to a reasonable size to avoid huge ranges let host_bits = 128 - prefix; if host_bits > 32 { 1u64 << 32 // Cap at 2^32 addresses } else { 1u64 << host_bits } } }; // Define the range within the transaction match range_names.get(range_name.as_bytes())? { Some(_) => Ok(()), // Range already exists None => { let id = range_names.generate_id()?; // Store name -> (id, size) let mut value = Vec::new(); value.extend_from_slice(&id.to_be_bytes()); value.extend_from_slice(&size.to_be_bytes()); range_names.insert(range_name.as_bytes(), value)?; // Initialize bitmap for the range (all zeros) let bitmap_size = (size + 7) / 8; // Round up to nearest byte let bitmap = vec![0u8; bitmap_size as usize]; range_map.insert(&id.to_be_bytes(), bitmap)?; Ok(()) } } } else { Ok(()) } } fn get_assignment(&self, network: &Network, identifier: &str, range_store: Option<&RangeStore>) -> Result> { if let Some(store) = range_store { let range_name = self.get_range_name(&network.name); // Find the bit position for this identifier if let Some(bit_position) = self.find_identifier_bit_position(&range_name, store, identifier)? { // Convert bit position to IP address let ip = self.bit_position_to_ip(bit_position, &network.netrange)?; Ok(Some(ip)) } else { Ok(None) } } else { Ok(None) } } + + fn cleanup_for_network(&self, network_name: &str, range_store: Option<&RangeStore>) -> Result<()> { + if let Some(store) = range_store { + let range_name = self.get_range_name(network_name); + store.delete_range(&range_name)?; + } + Ok(()) + } + + fn cleanup_for_network_in_transaction( + &self, + network_name: &str, + range_trees: Option<(&sled::transaction::TransactionalTree, &sled::transaction::TransactionalTree)> + ) -> Result<(), sled::transaction::ConflictableTransactionError<()>> { + if let Some((range_names, range_map)) = range_trees { + let range_name = self.get_range_name(network_name); + + // Get the range ID from the name + if let Some(value) = range_names.get(range_name.as_bytes())? { + let id = u64::from_be_bytes([ + value[0], value[1], value[2], value[3], + value[4], value[5], value[6], value[7], + ]); + + // Remove the range data + range_map.remove(&id.to_be_bytes())?; + } + + // Remove the range name + range_names.remove(range_name.as_bytes())?; + } + Ok(()) + } } /// Network configuration #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct Network { pub name: String, pub netrange: NetRange, pub provider_type: String, pub provider_config: String, pub assigner_type: Option, pub assigner_config: Option, } impl Network { /// Get an IP assignment for the given identifier /// /// This method calls the appropriate assigner implementation to retrieve /// the assigned IP address for a given identifier. Different assigners /// may return the IP in different formats (bit position or full IP). /// /// # Arguments /// /// * `identifier` - The identifier to look up (e.g., device name, container ID) /// * `range_store` - Optional reference to RangeStore (required for RangeAssigner) /// /// # Returns /// /// Returns the assigned IP address if found, None if no assignment exists /// /// # Example /// /// ``` /// # use store::network::{Network, NetRange, RangeAssigner}; /// # use tempfile::TempDir; /// # fn example() -> store::Result<()> { /// # let temp_dir = tempfile::tempdir().unwrap(); /// # let db = sled::open(temp_dir.path())?; /// # let ranges = store::range::RangeStore::open(&db)?; /// # let store = store::network::NetworkStore::open_with_ranges(&db, ranges.clone())?; /// # let netrange = NetRange::from_cidr("192.168.1.0/24")?; /// # let provider = store::network::BasicProvider::new(); /// # let assigner = RangeAssigner::new(); /// # store.create("test", netrange, provider, Some(assigner))?; /// # let network = store.get("test")?.unwrap(); /// /// let ip = network.get_assignment("device1", Some(&ranges))?; /// if let Some(addr) = ip { /// println!("Device1 is assigned IP: {}", addr); /// } /// # Ok(()) /// # } /// ``` + /// Create an assigner instance from the network configuration + pub fn create_assigner(&self) -> Result> { + if let (Some(assigner_type), Some(assigner_config)) = (&self.assigner_type, &self.assigner_config) { + match assigner_type.as_str() { + "range" => { + let range_assigner: RangeAssigner = serde_json::from_str(assigner_config) + .map_err(|e| Error::StoreError(sled::Error::Unsupported( + format!("Failed to deserialize RangeAssigner: {}", e) + )))?; + Ok(Box::new(range_assigner)) + } + "basic" => { + let basic_assigner: BasicAssigner = serde_json::from_str(assigner_config) + .map_err(|e| Error::StoreError(sled::Error::Unsupported( + format!("Failed to deserialize BasicAssigner: {}", e) + )))?; + Ok(Box::new(basic_assigner)) + } + _ => { + Err(Error::StoreError(sled::Error::Unsupported( + format!("Unknown assigner type: {}", assigner_type) + ))) + } + } + } else { + Err(Error::StoreError(sled::Error::Unsupported( + "No assigner configured".to_string() + ))) + } + } + pub fn get_assignment(&self, identifier: &str, range_store: Option<&RangeStore>) -> Result> { if let (Some(assigner_type), Some(assigner_config)) = (&self.assigner_type, &self.assigner_config) { match assigner_type.as_str() { "range" => { let range_assigner: RangeAssigner = serde_json::from_str(assigner_config) .map_err(|e| Error::StoreError(sled::Error::Unsupported( format!("Failed to deserialize RangeAssigner: {}", e) )))?; range_assigner.get_assignment(self, identifier, range_store) } "basic" => { let basic_assigner: BasicAssigner = serde_json::from_str(assigner_config) .map_err(|e| Error::StoreError(sled::Error::Unsupported( format!("Failed to deserialize BasicAssigner: {}", e) )))?; basic_assigner.get_assignment(self, identifier, range_store) } _ => { // Unknown assigner type Ok(None) } } } else { // No assigner configured Ok(None) } } /// Get an IP address from a bit position in the network range /// /// This method converts a bit position to the corresponding IP address /// within this network's range. Useful when you already know the bit /// position from range operations. /// /// # Arguments /// /// * `bit_position` - The bit position within the network range /// /// # Returns /// /// Returns the IP address corresponding to the bit position /// /// # Example /// /// ``` /// # use store::network::{Network, NetRange}; /// # fn example() -> store::Result<()> { /// # let netrange = NetRange::from_cidr("192.168.1.0/24")?; /// let network = Network { /// name: "test".to_string(), /// netrange: netrange, /// provider_type: "basic".to_string(), /// provider_config: "{}".to_string(), /// assigner_type: None, /// assigner_config: None, /// }; /// /// let ip = network.get_assignment_by_bit_position(5)?; /// println!("Bit position 5 corresponds to IP: {}", ip); /// # Ok(()) /// # } /// ``` pub fn get_assignment_by_bit_position(&self, bit_position: u64) -> Result { match &self.netrange { NetRange::V4 { addr, prefix: _ } => { let network_addr = u32::from(*addr); let host_addr = network_addr + bit_position as u32; Ok(IpAddr::V4(Ipv4Addr::from(host_addr))) } NetRange::V6 { addr, prefix: _ } => { let network_addr = u128::from(*addr); let host_addr = network_addr + bit_position as u128; Ok(IpAddr::V6(Ipv6Addr::from(host_addr))) } } } } /// Network store for managing network configurations #[derive(Debug, Clone)] pub struct NetworkStore { pub(crate) namespaces: crate::namespace::NamespaceStore, pub(crate) networks: sled::Tree, pub(crate) ranges: Option, } impl NetworkStore { /// Open a new network store pub fn open(db: &sled::Db) -> Result { Ok(NetworkStore { namespaces: crate::namespace::NamespaceStore::open(db)?, networks: db.open_tree("networks/1/data")?, ranges: None, }) } /// Open a new network store with range store pub fn open_with_ranges(db: &sled::Db, ranges: crate::range::RangeStore) -> Result { Ok(NetworkStore { namespaces: crate::namespace::NamespaceStore::open(db)?, networks: db.open_tree("networks/1/data")?, ranges: Some(ranges), }) } /// Create a new network with provider and optional assigner pub fn create(&self, name: &str, netrange: NetRange, provider: P, assigner: Option) -> Result<()> where P: NetworkProvider, A: NetworkAssigner, { // Ensure "networks" namespace exists if !self.namespaces.namespace_exists("networks")? { self.namespaces.define("networks")?; } let network = Network { name: name.to_string(), netrange: netrange.clone(), provider_type: provider.provider_type().to_string(), provider_config: provider.to_json()?, assigner_type: assigner.as_ref().map(|a| a.assigner_type().to_string()), assigner_config: assigner.as_ref().map(|a| a.to_json()).transpose()?, }; let network_json = serde_json::to_string(&network).map_err(|e| { Error::StoreError(sled::Error::Unsupported(format!("JSON serialization error: {}", e))) })?; // Include range store trees in transaction if needed if let (Some(assigner_ref), Some(range_store)) = (&assigner, &self.ranges) { if assigner_ref.assigner_type() == "range" { // Transaction with both network and range stores (&self.namespaces.names, &self.namespaces.spaces, &self.networks, &range_store.names, &range_store.map).transaction( |(names, spaces, networks, range_names, range_map)| { // Reserve the network name in the "networks" namespace if !self.namespaces.reserve_in_transaction(names, spaces, "networks", name, name)? { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } // Initialize the assigner within the transaction assigner_ref.initialize_for_network_in_transaction(name, &netrange, Some((range_names, range_map)))?; // Store the network configuration networks.insert(name.as_bytes(), network_json.as_bytes())?; Ok(()) } ).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => { Error::NamespaceKeyReserved("networks".to_string(), name.to_string()) } sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; } else { // Transaction without range store (&self.namespaces.names, &self.namespaces.spaces, &self.networks).transaction( |(names, spaces, networks)| { // Reserve the network name in the "networks" namespace if !self.namespaces.reserve_in_transaction(names, spaces, "networks", name, name)? { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } // Store the network configuration networks.insert(name.as_bytes(), network_json.as_bytes())?; Ok(()) } ).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => { Error::NamespaceKeyReserved("networks".to_string(), name.to_string()) } sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; } } else { // Transaction without assigner or range store (&self.namespaces.names, &self.namespaces.spaces, &self.networks).transaction( |(names, spaces, networks)| { // Reserve the network name in the "networks" namespace if !self.namespaces.reserve_in_transaction(names, spaces, "networks", name, name)? { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } // Store the network configuration networks.insert(name.as_bytes(), network_json.as_bytes())?; Ok(()) } ).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => { Error::NamespaceKeyReserved("networks".to_string(), name.to_string()) } sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; } Ok(()) } /// Get a network by name pub fn get(&self, name: &str) -> Result> { if let Some(data) = self.networks.get(name.as_bytes())? { let network_json = String::from_utf8(data.to_vec())?; let network: Network = serde_json::from_str(&network_json).map_err(|e| { Error::StoreError(sled::Error::Unsupported(format!("JSON deserialization error: {}", e))) })?; Ok(Some(network)) } else { Ok(None) } } /// Update an existing network pub fn update(&self, name: &str, netrange: NetRange, provider: P, assigner: Option) -> Result<()> where P: NetworkProvider, A: NetworkAssigner, { if !self.namespaces.key_exists("networks", name)? { return Err(Error::StoreError(sled::Error::Unsupported( format!("Network '{}' does not exist", name) ))); } let network = Network { name: name.to_string(), netrange, provider_type: provider.provider_type().to_string(), provider_config: provider.to_json()?, assigner_type: assigner.as_ref().map(|a| a.assigner_type().to_string()), assigner_config: assigner.as_ref().map(|a| a.to_json()).transpose()?, }; let network_json = serde_json::to_string(&network).map_err(|e| { Error::StoreError(sled::Error::Unsupported(format!("JSON serialization error: {}", e))) })?; self.networks.insert(name.as_bytes(), network_json.as_bytes())?; Ok(()) } /// Delete a network by name pub fn delete(&self, name: &str) -> Result { - let result = (&self.namespaces.names, &self.namespaces.spaces, &self.networks).transaction( - |(names, spaces, networks)| { - // Remove from namespace - let removed = self.namespaces.remove_in_transaction(names, spaces, "networks", name)?; - if removed { - // Remove network data - networks.remove(name.as_bytes())?; + let result = if let Some(ranges) = &self.ranges { + // Transaction with range cleanup + (&self.namespaces.names, &self.namespaces.spaces, &self.networks, &ranges.names, &ranges.map).transaction( + |(names, spaces, networks, range_names, range_map)| { + // Get network data before deletion to access assigner + if let Some(network_data) = networks.get(name.as_bytes())? { + let network: Network = serde_json::from_slice(network_data.as_ref()) + .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; + + // Clean up assigner if it exists + if network.assigner_type.is_some() { + let assigner = network.create_assigner() + .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; + assigner.cleanup_for_network_in_transaction(name, Some((range_names, range_map)))?; + } + } + + // Remove from namespace + let removed = self.namespaces.remove_in_transaction(names, spaces, "networks", name)?; + if removed { + // Remove network data + networks.remove(name.as_bytes())?; + } + Ok(removed) } - Ok(removed) - } - ).map_err(|e| match e { + ) + } else { + // Transaction without range cleanup + (&self.namespaces.names, &self.namespaces.spaces, &self.networks).transaction( + |(names, spaces, networks)| { + // Remove from namespace + let removed = self.namespaces.remove_in_transaction(names, spaces, "networks", name)?; + if removed { + // Remove network data + networks.remove(name.as_bytes())?; + } + Ok(removed) + } + ) + }; + + let result = result.map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => { Error::StoreError(sled::Error::Unsupported("Transaction aborted".to_string())) } sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } /// List all network names pub fn list(&self) -> Result> { self.namespaces.list_keys("networks") } /// Check if a network exists pub fn exists(&self, name: &str) -> Result { self.namespaces.key_exists("networks", name) } } #[cfg(test)] mod tests { use super::*; use tempfile::TempDir; fn create_test_store() -> Result<(NetworkStore, TempDir)> { let temp_dir = tempfile::tempdir().unwrap(); let db = sled::open(temp_dir.path())?; let ranges = crate::range::RangeStore::open(&db)?; let store = NetworkStore::open_with_ranges(&db, ranges)?; Ok((store, temp_dir)) } #[test] fn test_netrange_ipv4() -> Result<()> { let netrange = NetRange::ipv4("192.168.1.0".parse().unwrap(), 24)?; assert_eq!(netrange.to_cidr(), "192.168.1.0/24"); Ok(()) } #[test] fn test_netrange_ipv6() -> Result<()> { let netrange = NetRange::ipv6("2001:db8::".parse().unwrap(), 64)?; assert_eq!(netrange.to_cidr(), "2001:db8::/64"); Ok(()) } #[test] fn test_netrange_from_cidr() -> Result<()> { let ipv4_range = NetRange::from_cidr("10.0.0.0/8")?; assert_eq!(ipv4_range.to_cidr(), "10.0.0.0/8"); let ipv6_range = NetRange::from_cidr("fe80::/10")?; assert_eq!(ipv6_range.to_cidr(), "fe80::/10"); Ok(()) } #[test] fn test_basic_provider() -> Result<()> { let provider = BasicProvider::new() .with_config("endpoint", "https://api.example.com") .with_config("timeout", "30"); let json = provider.to_json()?; let restored = BasicProvider::from_json(&json)?; assert_eq!(provider.config, restored.config); assert_eq!(provider.provider_type(), "basic"); Ok(()) } #[test] fn test_basic_assigner() -> Result<()> { let assigner = BasicAssigner::new("round_robin") .with_config("pool_size", "100"); let json = assigner.to_json()?; let restored = BasicAssigner::from_json(&json)?; assert_eq!(assigner.strategy, restored.strategy); assert_eq!(assigner.config, restored.config); assert_eq!(assigner.assigner_type(), "basic"); Ok(()) } #[test] fn test_create_and_get_network() -> Result<()> { let (store, _temp_dir) = create_test_store()?; let netrange = NetRange::ipv4("192.168.1.0".parse().unwrap(), 24)?; let provider = BasicProvider::new().with_config("type", "test"); let assigner = Some(BasicAssigner::new("sequential")); store.create("test_network", netrange.clone(), provider, assigner)?; let network = store.get("test_network")?.unwrap(); assert_eq!(network.name, "test_network"); assert_eq!(network.netrange, netrange); assert_eq!(network.provider_type, "basic"); assert!(network.assigner_type.is_some()); Ok(()) } #[test] fn test_create_network_without_assigner() -> Result<()> { let (store, _temp_dir) = create_test_store()?; let netrange = NetRange::ipv6("2001:db8::".parse().unwrap(), 64)?; let provider = BasicProvider::new(); let assigner: Option = None; store.create("ipv6_network", netrange.clone(), provider, assigner)?; let network = store.get("ipv6_network")?.unwrap(); assert_eq!(network.name, "ipv6_network"); assert_eq!(network.netrange, netrange); assert!(network.assigner_type.is_none()); assert!(network.assigner_config.is_none()); Ok(()) } #[test] fn test_create_duplicate_network() -> Result<()> { let (store, _temp_dir) = create_test_store()?; let netrange = NetRange::ipv4("10.0.0.0".parse().unwrap(), 8)?; let provider = BasicProvider::new(); let assigner: Option = None; // First creation should succeed store.create("duplicate_test", netrange.clone(), provider.clone(), assigner.clone())?; // Second creation should fail let result = store.create("duplicate_test", netrange, provider, assigner); assert!(result.is_err()); Ok(()) } #[test] fn test_update_network() -> Result<()> { let (store, _temp_dir) = create_test_store()?; // Create initial network let netrange = NetRange::ipv4("172.16.0.0".parse().unwrap(), 12)?; let provider = BasicProvider::new().with_config("version", "1"); let assigner: Option = None; store.create("update_test", netrange, provider, assigner)?; // Update the network let new_netrange = NetRange::ipv4("172.16.0.0".parse().unwrap(), 16)?; let new_provider = BasicProvider::new().with_config("version", "2"); let new_assigner = Some(BasicAssigner::new("random")); store.update("update_test", new_netrange.clone(), new_provider, new_assigner)?; // Verify the update let network = store.get("update_test")?.unwrap(); assert_eq!(network.netrange, new_netrange); assert!(network.assigner_type.is_some()); Ok(()) } #[test] fn test_delete_network() -> Result<()> { let (store, _temp_dir) = create_test_store()?; let netrange = NetRange::ipv4("203.0.113.0".parse().unwrap(), 24)?; let provider = BasicProvider::new(); let assigner: Option = None; store.create("delete_test", netrange, provider, assigner)?; assert!(store.exists("delete_test")?); let deleted = store.delete("delete_test")?; assert!(deleted); assert!(!store.exists("delete_test")?); // Try to delete non-existent network let deleted_again = store.delete("delete_test")?; assert!(!deleted_again); Ok(()) } #[test] fn test_list_networks() -> Result<()> { let (store, _temp_dir) = create_test_store()?; // Create multiple networks let networks = vec![ ("net1", "10.1.0.0/16"), ("net2", "10.2.0.0/16"), ("net3", "2001:db8:1::/48"), ]; for (name, cidr) in &networks { let netrange = NetRange::from_cidr(cidr)?; let provider = BasicProvider::new(); let assigner: Option = None; store.create(name, netrange, provider, assigner)?; } let mut network_names = store.list()?; network_names.sort(); let mut expected: Vec = networks.iter().map(|(name, _)| name.to_string()).collect(); expected.sort(); assert_eq!(network_names, expected); Ok(()) } #[test] fn test_update_nonexistent_network() -> Result<()> { let (store, _temp_dir) = create_test_store()?; let netrange = NetRange::ipv4("198.51.100.0".parse().unwrap(), 24)?; let provider = BasicProvider::new(); let assigner: Option = None; let result = store.update("nonexistent", netrange, provider, assigner); assert!(result.is_err()); Ok(()) } #[test] fn test_invalid_cidr() { let result = NetRange::from_cidr("invalid"); assert!(result.is_err()); let result = NetRange::from_cidr("192.168.1.0"); assert!(result.is_err()); let result = NetRange::from_cidr("192.168.1.0/33"); assert!(result.is_err()); } #[test] fn test_range_assigner() -> Result<()> { let (store, _temp_dir) = create_test_store()?; // Create a network with RangeAssigner let netrange = NetRange::from_cidr("192.168.1.0/24")?; let provider = BasicProvider::new() .with_config("type", "test"); let assigner = RangeAssigner::new() .with_config("strategy", "sequential"); store.create("test-network", netrange, provider, Some(assigner))?; // Verify the network was created let network = store.get("test-network")?.unwrap(); assert_eq!(network.name, "test-network"); assert_eq!(network.assigner_type, Some("range".to_string())); // Test range assignment functionality if let Some(ref range_store) = store.ranges { let range_assigner = RangeAssigner::new(); let range_name = "test-network"; // Uses network name as range name // Assign some IPs let bit1 = range_assigner.assign_ip(range_name, range_store, "device1")?; let bit2 = range_assigner.assign_ip(range_name, range_store, "device2")?; // Verify assignments assert_eq!(bit1, 0); assert_eq!(bit2, 1); // Get assignment info let assignment1 = range_assigner.get_assignment_by_bit(range_name, range_store, bit1)?.unwrap(); let assignment2 = range_assigner.get_assignment_by_bit(range_name, range_store, bit2)?.unwrap(); assert_eq!(assignment1, "device1"); assert_eq!(assignment2, "device2"); // Test unassignment let unassigned = range_assigner.unassign_ip(range_name, range_store, bit1)?; assert!(unassigned); let assignment1_after = range_assigner.get_assignment_by_bit(range_name, range_store, bit1)?; assert!(assignment1_after.is_none()); } Ok(()) } #[test] fn test_range_assigner_serialization() -> Result<()> { let assigner = RangeAssigner::with_range_name("test-range") .with_config("strategy", "random") .with_config("pool_size", "100"); // Test serialization let json = assigner.to_json()?; assert!(json.contains("test-range")); assert!(json.contains("random")); assert!(json.contains("100")); // Test deserialization let deserialized: RangeAssigner = RangeAssigner::from_json(&json)?; assert_eq!(deserialized.range_name, Some("test-range".to_string())); assert_eq!(deserialized.config.get("strategy"), Some(&"random".to_string())); assert_eq!(deserialized.config.get("pool_size"), Some(&"100".to_string())); Ok(()) } #[test] fn test_network_get_assignment() -> Result<()> { let (store, _temp_dir) = create_test_store()?; // Create a network with RangeAssigner let netrange = NetRange::from_cidr("192.168.1.0/24")?; let provider = BasicProvider::new() .with_config("type", "test"); let assigner = RangeAssigner::new() .with_config("strategy", "sequential"); store.create("test-network", netrange, provider, Some(assigner))?; // Get the network let network = store.get("test-network")?.unwrap(); // Test that no assignment exists initially let ip = network.get_assignment("device1", store.ranges.as_ref())?; assert!(ip.is_none()); // Assign an IP using the range assigner directly if let Some(ref range_store) = store.ranges { let range_assigner = RangeAssigner::new(); let range_name = "test-network"; // Uses network name as range name // Assign an IP to device1 let bit_position = range_assigner.assign_ip(range_name, range_store, "device1")?; assert_eq!(bit_position, 0); // Now test get_assignment on the network let ip = network.get_assignment("device1", store.ranges.as_ref())?; assert!(ip.is_some()); let assigned_ip = ip.unwrap(); match assigned_ip { std::net::IpAddr::V4(ipv4) => { // Should be 192.168.1.0 + bit_position let expected = std::net::Ipv4Addr::new(192, 168, 1, bit_position as u8); assert_eq!(ipv4, expected); } _ => panic!("Expected IPv4 address"), } // Test assignment for non-existent device let no_ip = network.get_assignment("nonexistent", store.ranges.as_ref())?; assert!(no_ip.is_none()); // Assign another device let bit_position2 = range_assigner.assign_ip(range_name, range_store, "device2")?; assert_eq!(bit_position2, 1); let ip2 = network.get_assignment("device2", store.ranges.as_ref())?; assert!(ip2.is_some()); let assigned_ip2 = ip2.unwrap(); match assigned_ip2 { std::net::IpAddr::V4(ipv4) => { let expected = std::net::Ipv4Addr::new(192, 168, 1, bit_position2 as u8); assert_eq!(ipv4, expected); } _ => panic!("Expected IPv4 address"), } } Ok(()) } #[test] fn test_network_get_assignment_no_assigner() -> Result<()> { let (store, _temp_dir) = create_test_store()?; // Create a network without an assigner let netrange = NetRange::from_cidr("192.168.1.0/24")?; let provider = BasicProvider::new() .with_config("type", "test"); store.create("no-assigner-network", netrange, provider, None::)?; // Get the network let network = store.get("no-assigner-network")?.unwrap(); // Test that get_assignment returns None when no assigner is configured let ip = network.get_assignment("device1", store.ranges.as_ref())?; assert!(ip.is_none()); Ok(()) } + #[test] + fn test_delete_network_with_range_assigner() -> Result<()> { + let (store, _temp_dir) = create_test_store()?; + + // Create a network with a range assigner + let netrange = NetRange::from_cidr("192.168.100.0/24")?; + let provider = BasicProvider::new() + .with_config("type", "test"); + let assigner = RangeAssigner::new() + .with_config("strategy", "sequential"); + + store.create("test-range-network", netrange, provider, Some(assigner))?; + + // Verify the network was created + assert!(store.exists("test-range-network")?); + let network = store.get("test-range-network")?.unwrap(); + + // Verify the range was initialized by checking if the range exists + if let Some(ranges) = &store.ranges { + let range_name = "test-range-network"; + // Check if range exists by trying to get range info + let range_exists = ranges.names.get(range_name.as_bytes())?.is_some(); + assert!(range_exists, "Range should be initialized after network creation"); + + // Make an assignment using the range assigner directly + let range_assigner = RangeAssigner::new(); + let _bit_position = range_assigner.assign_ip(range_name, ranges, "device1")?; + // Assignment should succeed (bit_position is always >= 0 since it's u64) + + // Now verify the assignment exists + let ip = network.get_assignment("device1", store.ranges.as_ref())?; + assert!(ip.is_some(), "Assignment should be retrievable"); + + // Verify assignment data exists + let range_info = ranges.list_range(range_name)?; + assert!(!range_info.is_empty(), "Range should have assignment data"); + } + + // Delete the network + let deleted = store.delete("test-range-network")?; + assert!(deleted); + + // Verify the network no longer exists + assert!(!store.exists("test-range-network")?); + assert!(store.get("test-range-network")?.is_none()); + + // Verify range data was cleaned up + if let Some(ranges) = &store.ranges { + let range_name = "test-range-network"; + let range_exists = ranges.names.get(range_name.as_bytes())?.is_some(); + assert!(!range_exists, "Range should be deleted"); + + let range_info = ranges.list_range(range_name)?; + assert!(range_info.is_empty(), "Range assignments should be cleaned up"); + } + + Ok(()) + } + + #[test] + fn test_delete_network_without_assigner() -> Result<()> { + let (store, _temp_dir) = create_test_store()?; + + // Create a network without an assigner + let netrange = NetRange::from_cidr("10.10.0.0/16")?; + let provider = BasicProvider::new() + .with_config("type", "test"); + + store.create("test-no-assigner", netrange, provider, None::)?; + + // Verify the network was created + assert!(store.exists("test-no-assigner")?); + let network = store.get("test-no-assigner")?.unwrap(); + assert!(network.assigner_type.is_none()); + + // Delete the network + let deleted = store.delete("test-no-assigner")?; + assert!(deleted); + + // Verify the network no longer exists + assert!(!store.exists("test-no-assigner")?); + assert!(store.get("test-no-assigner")?.is_none()); + + Ok(()) + } + #[test] fn test_network_get_assignment_ipv6() -> Result<()> { let (store, _temp_dir) = create_test_store()?; // Create an IPv6 network with RangeAssigner let netrange = NetRange::from_cidr("2001:db8::/64")?; let provider = BasicProvider::new() .with_config("type", "test"); let assigner = RangeAssigner::new() .with_config("strategy", "sequential"); store.create("ipv6-network", netrange, provider, Some(assigner))?; // Get the network let network = store.get("ipv6-network")?.unwrap(); // Assign an IP using the range assigner directly if let Some(ref range_store) = store.ranges { let range_assigner = RangeAssigner::new(); let range_name = "ipv6-network"; // Assign an IP to device1 let bit_position = range_assigner.assign_ip(range_name, range_store, "ipv6-device1")?; assert_eq!(bit_position, 0); // Test get_assignment on the network let ip = network.get_assignment("ipv6-device1", store.ranges.as_ref())?; assert!(ip.is_some()); let assigned_ip = ip.unwrap(); match assigned_ip { std::net::IpAddr::V6(ipv6) => { // Should be 2001:db8:: + bit_position let expected = std::net::Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, bit_position as u16); assert_eq!(ipv6, expected); } _ => panic!("Expected IPv6 address"), } } Ok(()) } } \ No newline at end of file diff --git a/crates/store/src/range.rs b/crates/store/src/range.rs index c599615..686915d 100644 --- a/crates/store/src/range.rs +++ b/crates/store/src/range.rs @@ -1,929 +1,1126 @@ //! # Range Store Module //! //! This module implements "ranges" which are buckets of unique range assignments //! that map values to specific bit positions within fixed-size ranges. This is useful //! for allocating unique identifiers, session tokens, or any scenario where you need //! to assign sequential positions to values. //! //! ## Storage Design //! //! The RangeStore uses three sled trees: //! //! - **`names` tree**: Maps range names to their metadata //! - Key: `name` (string) //! - Value: `id || size` where `id` is a u64 from `generate_id()` and `size` is a u64 //! //! - **`map` tree**: Stores bitmaps for each range to track allocated positions //! - Key: `id` (u64 from the `names` tree) //! - Value: `bitmap` (binary data representing which bits are allocated) //! //! - **`assign` tree**: Maps range positions to their assigned values //! - Key: `id || bit_position` where `id` is the range ID and `bit_position` is a u64 //! - Value: `value` (the string value assigned to this position) //! //! ## API Operations //! //! - `define(name, size)`: Create a new range with the given name and maximum size //! - `assign(name, value)`: Assign a value to the next available position in the range //! - `get(name, position)`: Retrieve the value at a specific position //! - `list_range(name)`: List all assignments in a range //! - `list_ranges()`: List all defined ranges //! - `unassign(name, value)`: Remove a value from the range and free its position //! //! ## Example Usage //! //! ```rust //! use store::{open, Result}; //! //! fn example() -> Result<()> { //! let store = open()?; //! let ranges = store.ranges(); //! //! // Define a range for user IDs //! ranges.define("user_ids", 1000)?; //! //! // Assign some users //! let alice_id = ranges.assign("user_ids", "alice@example.com")?; // Returns 0 //! let bob_id = ranges.assign("user_ids", "bob@example.com")?; // Returns 1 //! //! // Get a user by position //! let email = ranges.get("user_ids", alice_id)?; // Some("alice@example.com") //! //! // Remove a user (frees the position for reuse) //! ranges.unassign("user_ids", "bob@example.com")?; //! //! // Next assignment will reuse Bob's position //! let charlie_id = ranges.assign("user_ids", "charlie@example.com")?; // Returns 1 //! //! Ok(()) //! } //! ``` use crate::Result; use crate::Error; use sled::Transactional; #[derive(Debug, Clone)] pub struct RangeStore { pub(crate) names: sled::Tree, pub(crate) map: sled::Tree, pub(crate) assign: sled::Tree, } impl RangeStore { pub fn open(db: &sled::Db) -> Result { Ok(RangeStore { names: db.open_tree("ranges/1/names")?, map: db.open_tree("ranges/1/map")?, assign: db.open_tree("ranges/1/assign")?, }) } /// Assign a value to a range within an existing transaction context pub fn assign_in_transaction( &self, names: &sled::transaction::TransactionalTree, map: &sled::transaction::TransactionalTree, assign: &sled::transaction::TransactionalTree, range_name: &str, value: &str, ) -> sled::transaction::ConflictableTransactionResult { // Get range info let (range_id, range_size) = match names.get(range_name)? { Some(data) => { if data.len() != 16 { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; let size_bytes: [u8; 8] = data[8..16].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; (u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)) } None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Get current bitmap let mut bitmap = match map.get(&range_id.to_be_bytes())? { Some(data) => data.to_vec(), None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Find first free bit let mut bit_position = None; for bit in 0..range_size { let byte_index = (bit / 8) as usize; let bit_index = bit % 8; if byte_index < bitmap.len() { let mask = 1u8 << bit_index; if (bitmap[byte_index] & mask) == 0 { // This bit is free bitmap[byte_index] |= mask; bit_position = Some(bit); break; } } } let bit_pos = match bit_position { Some(pos) => pos, None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Update bitmap map.insert(&range_id.to_be_bytes(), bitmap)?; // Store assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_pos.to_be_bytes()); assign.insert(assign_key, value.as_bytes())?; Ok(bit_pos) } /// Get a value from a specific bit position within an existing transaction context /// Get range assignment details within an existing transaction context pub fn get_in_transaction( &self, names: &sled::transaction::TransactionalTree, assign: &sled::transaction::TransactionalTree, range_name: &str, bit_position: u64, ) -> sled::transaction::ConflictableTransactionResult, ()> { // Get range info let (range_id, range_size) = match names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(None); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; let size_bytes: [u8; 8] = data[8..16].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; (u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)) } None => return Ok(None), }; if bit_position >= range_size { return Ok(None); } // Get assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_position.to_be_bytes()); match assign.get(&assign_key)? { Some(value_bytes) => { let value = String::from_utf8(value_bytes.to_vec()) .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; Ok(Some(value)) } None => Ok(None), } } /// Unassign a specific bit position within an existing transaction context pub fn unassign_bit_in_transaction( &self, names: &sled::transaction::TransactionalTree, map: &sled::transaction::TransactionalTree, assign: &sled::transaction::TransactionalTree, range_name: &str, bit_position: u64, ) -> sled::transaction::ConflictableTransactionResult { // Get range info let (range_id, range_size) = match names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(false); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; let size_bytes: [u8; 8] = data[8..16].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; (u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)) } None => return Ok(false), }; if bit_position >= range_size { return Ok(false); } // Remove assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_position.to_be_bytes()); let removed = assign.remove(assign_key)?.is_some(); if removed { // Update bitmap let mut bitmap = match map.get(&range_id.to_be_bytes())? { Some(data) => data.to_vec(), None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; let byte_index = (bit_position / 8) as usize; let bit_index = bit_position % 8; if byte_index < bitmap.len() { let mask = 1u8 << bit_index; bitmap[byte_index] &= !mask; // Clear the bit map.insert(&range_id.to_be_bytes(), bitmap)?; } } Ok(removed) } /// Check if a range exists within an existing transaction context pub fn exists_in_transaction( &self, names: &sled::transaction::TransactionalTree, range_name: &str, ) -> sled::transaction::ConflictableTransactionResult { Ok(names.get(range_name)?.is_some()) } /// Get range info (id and size) within an existing transaction context pub fn info_in_transaction( &self, names: &sled::transaction::TransactionalTree, range_name: &str, ) -> sled::transaction::ConflictableTransactionResult, ()> { match names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(None); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; let size_bytes: [u8; 8] = data[8..16].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; Ok(Some((u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)))) } None => Ok(None), } } // Define a new range with a given name and size pub fn define(&self, name: &str, size: u64) -> Result<()> { (&self.names, &self.map).transaction(|(names, map)| { self.define_in_transaction(names, map, name, size) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::StoreError(sled::Error::Unsupported("Range definition failed".to_string())), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(()) } /// Define a new range within an existing transaction pub fn define_in_transaction( &self, names: &sled::transaction::TransactionalTree, map: &sled::transaction::TransactionalTree, name: &str, size: u64, ) -> Result<(), sled::transaction::ConflictableTransactionError<()>> { match names.get(name)? { Some(_) => Ok(()), // Range already exists None => { let id = names.generate_id()?; // Store name -> (id, size) let mut value = Vec::new(); value.extend_from_slice(&id.to_be_bytes()); value.extend_from_slice(&size.to_be_bytes()); names.insert(name, value)?; // Initialize bitmap for the range (all zeros) let bitmap_size = (size + 7) / 8; // Round up to nearest byte let bitmap = vec![0u8; bitmap_size as usize]; map.insert(&id.to_be_bytes(), bitmap)?; Ok(()) } } } // Assign the next free bit in the range to a value pub fn assign(&self, range_name: &str, value: &str) -> Result { let result = (&self.names, &self.map, &self.assign).transaction(|(names, map, assign)| { // Get range info let (range_id, range_size) = match names.get(range_name)? { Some(data) => { if data.len() != 16 { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; let size_bytes: [u8; 8] = data[8..16].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; (u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)) } None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Get current bitmap let mut bitmap = match map.get(&range_id.to_be_bytes())? { Some(data) => data.to_vec(), None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Find first free bit let mut bit_position = None; for bit in 0..range_size { let byte_index = (bit / 8) as usize; let bit_index = bit % 8; if byte_index < bitmap.len() { let mask = 1u8 << bit_index; if (bitmap[byte_index] & mask) == 0 { // This bit is free bitmap[byte_index] |= mask; bit_position = Some(bit); break; } } } let bit_pos = match bit_position { Some(pos) => pos, None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Update bitmap map.insert(&range_id.to_be_bytes(), bitmap)?; // Store assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_pos.to_be_bytes()); assign.insert(assign_key, value.as_bytes())?; Ok(bit_pos) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::RangeFull(range_name.to_string()), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } // Get a value from a specific bit position in the range pub fn get(&self, range_name: &str, bit_position: u64) -> Result> { let result = (&self.names, &self.assign).transaction(|(names, assign)| { // Get range info let (range_id, range_size) = match names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(None); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; let size_bytes: [u8; 8] = data[8..16].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; (u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)) } None => return Ok(None), }; if bit_position >= range_size { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } // Get assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_position.to_be_bytes()); match assign.get(assign_key)? { Some(value_bytes) => { let value = String::from_utf8(value_bytes.to_vec()) .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; Ok(Some(value)) } None => Ok(None), } }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::BitOutOfRange(range_name.to_string(), bit_position), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } // List all assignments in a range pub fn list_range(&self, range_name: &str) -> Result> { // Get range info first let range_id = match self.names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(Vec::new()); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid range data".to_string())))?; u64::from_be_bytes(id_bytes) } None => return Ok(Vec::new()), }; let mut assignments = Vec::new(); let prefix = range_id.to_be_bytes(); for result in self.assign.scan_prefix(&prefix) { let (key, value) = result?; if key.len() == 16 { // 8 bytes for range_id + 8 bytes for bit_position let bit_position_bytes: [u8; 8] = key[8..16].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid key format".to_string())))?; let bit_position = u64::from_be_bytes(bit_position_bytes); let value_str = String::from_utf8(value.to_vec())?; assignments.push((bit_position, value_str)); } } assignments.sort_by_key(|&(bit_pos, _)| bit_pos); Ok(assignments) } // List all ranges pub fn list_ranges(&self) -> Result> { let mut ranges = Vec::new(); for result in self.names.iter() { let (key, value) = result?; if value.len() == 16 { let name = String::from_utf8(key.to_vec())?; let size_bytes: [u8; 8] = value[8..16].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid size data".to_string())))?; let size = u64::from_be_bytes(size_bytes); ranges.push((name, size)); } } ranges.sort_by(|a, b| a.0.cmp(&b.0)); Ok(ranges) } // Unassign a specific bit position from the range pub fn unassign_bit(&self, range_name: &str, bit_position: u64) -> Result { let result = (&self.names, &self.map, &self.assign).transaction(|(names, map, assign)| { self.unassign_bit_in_transaction(names, map, assign, range_name, bit_position) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::BitOutOfRange(range_name.to_string(), bit_position), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } // Unassign a value from the range (find by value and remove) pub fn unassign(&self, range_name: &str, value: &str) -> Result { // First find the assignment outside of transaction let range_id = match self.names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(false); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid range data".to_string())))?; u64::from_be_bytes(id_bytes) } None => return Ok(false), }; // Find the assignment with this value let prefix = range_id.to_be_bytes(); let mut found_bit = None; for result in self.assign.scan_prefix(&prefix) { let (key, stored_value) = result?; if key.len() == 16 { let stored_value_str = String::from_utf8(stored_value.to_vec())?; if stored_value_str == value { let bit_position_bytes: [u8; 8] = key[8..16].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid key format".to_string())))?; found_bit = Some(u64::from_be_bytes(bit_position_bytes)); break; } } } let bit_position = match found_bit { Some(pos) => pos, None => return Ok(false), }; // Now perform the removal in a transaction let result = (&self.map, &self.assign).transaction(|(map, assign)| { // Remove assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_position.to_be_bytes()); assign.remove(assign_key)?; // Update bitmap let mut bitmap = match map.get(&range_id.to_be_bytes())? { Some(data) => data.to_vec(), None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; let byte_index = (bit_position / 8) as usize; let bit_index = bit_position % 8; if byte_index < bitmap.len() { let mask = 1u8 << bit_index; bitmap[byte_index] &= !mask; // Clear the bit map.insert(&range_id.to_be_bytes(), bitmap)?; } Ok(true) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::ValueNotInRange(range_name.to_string(), value.to_string()), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } + /// Delete a range and all its associated data + pub fn delete_range(&self, range_name: &str) -> Result { + // First collect keys to remove outside of transaction + let prefix = format!("{}:", range_name); + let mut keys_to_remove = Vec::new(); + for result in self.assign.scan_prefix(prefix.as_bytes()) { + let (key, _) = result?; + keys_to_remove.push(key.to_vec()); + } + + let result = (&self.names, &self.map, &self.assign).transaction( + |(names, map, assign)| { + // Get the range ID from the name + if let Some(value) = names.get(range_name.as_bytes())? { + let id = u64::from_be_bytes([ + value[0], value[1], value[2], value[3], + value[4], value[5], value[6], value[7], + ]); + + // Remove the range data + map.remove(&id.to_be_bytes())?; + + // Remove all assignments for this range + for key in &keys_to_remove { + assign.remove(key.as_slice())?; + } + + // Remove the range name + names.remove(range_name.as_bytes())?; + + Ok(true) + } else { + Ok(false) + } + } + ).map_err(|e| match e { + sled::transaction::TransactionError::Abort(()) => { + Error::StoreError(sled::Error::Unsupported("Transaction aborted".to_string())) + } + sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), + })?; + + Ok(result) + } + } #[cfg(test)] mod tests { use super::*; use tempfile::tempdir; fn create_test_store() -> Result { let temp_dir = tempdir().unwrap(); let db = sled::open(temp_dir.path())?; RangeStore::open(&db) } #[test] fn test_transaction_methods() -> Result<()> { let store = create_test_store()?; // Define a range store.define("test_range", 100)?; // Test transaction methods let result = (&store.names, &store.map, &store.assign).transaction(|(names, map, assign)| { // Assign in transaction let bit_pos = store.assign_in_transaction(names, map, assign, "test_range", "tx_value")?; assert!(bit_pos < 100); // Get in transaction let value = store.get_in_transaction(names, assign, "test_range", bit_pos)?; assert_eq!(value, Some("tx_value".to_string())); // Check if range exists in transaction let exists = store.exists_in_transaction(names, "test_range")?; assert_eq!(exists, true); // Get range info in transaction let info = store.info_in_transaction(names, "test_range")?; assert!(info.is_some()); let (_, size) = info.unwrap(); assert_eq!(size, 100); // Verify assignment exists let value_check = store.get_in_transaction(names, assign, "test_range", bit_pos)?; assert_eq!(value_check, Some("tx_value".to_string())); // Unassign bit in transaction let unassigned = store.unassign_bit_in_transaction(names, map, assign, "test_range", bit_pos)?; assert_eq!(unassigned, true); // Verify it's gone let value_after = store.get_in_transaction(names, assign, "test_range", bit_pos)?; assert_eq!(value_after, None); Ok(()) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::StoreError(sled::Error::Unsupported("Transaction aborted".to_string())), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } #[test] fn test_transaction_rollback() -> Result<()> { let store = create_test_store()?; // Define a range and assign initial value store.define("test_range", 5)?; // Small range to easily fill let initial_bit = store.assign("test_range", "initial_value")?; // Fill up the range for i in 1..5 { store.assign("test_range", &format!("value_{}", i))?; } // Attempt a transaction that should fail (range full) let result = (&store.names, &store.map, &store.assign).transaction(|(names, map, assign)| { // This should succeed let _unassigned = store.unassign_bit_in_transaction(names, map, assign, "test_range", initial_bit)?; // This should fail (range is still full after we just freed one bit) // Actually, this should succeed since we freed a bit, so let's try to assign to a full range differently // Let's assign multiple values to ensure failure store.assign_in_transaction(names, map, assign, "test_range", "temp1")?; store.assign_in_transaction(names, map, assign, "test_range", "temp2")?; // This should fail Ok(()) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::RangeFull("test_range".to_string()), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), }); // Transaction should have failed assert!(result.is_err()); // Verify rollback - initial value should still be there let value = store.get("test_range", initial_bit)?; assert_eq!(value, Some("initial_value".to_string())); Ok(()) } #[test] fn test_bit_out_of_range() -> Result<()> { let store = create_test_store()?; // Define a small range store.define("small_range", 10)?; // Try to get bit beyond range let result = store.get("small_range", 15); assert!(result.is_err()); // Try to unassign bit beyond range let result = store.unassign_bit("small_range", 15); assert!(result.is_ok()); assert_eq!(result?, false); Ok(()) } #[test] fn test_multiple_ranges() -> Result<()> { let store = create_test_store()?; // Define multiple ranges store.define("range_a", 50)?; store.define("range_b", 100)?; store.define("range_c", 25)?; // Assign values to each range let bit_a = store.assign("range_a", "value_a")?; let bit_b = store.assign("range_b", "value_b")?; let bit_c = store.assign("range_c", "value_c")?; // Verify isolation between ranges assert_eq!(store.get("range_a", bit_a)?, Some("value_a".to_string())); assert_eq!(store.get("range_b", bit_b)?, Some("value_b".to_string())); assert_eq!(store.get("range_c", bit_c)?, Some("value_c".to_string())); // Cross-range access: accessing bit positions that exist in one range but not another // Since all ranges start at bit 0, we need to test with positions beyond smaller ranges // Test accessing valid bit positions across ranges // All ranges have their first assignment at bit 0, so this tests basic isolation assert_eq!(store.get("range_a", bit_a)?, Some("value_a".to_string())); // Test accessing bit positions that are out of range // range_c has only 25 bits (0-24), so bit 25 and above should give an error let result = store.get("range_c", 25); assert!(result.is_err()); // range_a has 50 bits (0-49), so bit 50 and above should give an error let result = store.get("range_a", 50); assert!(result.is_err()); // List all ranges let ranges = store.list_ranges()?; assert_eq!(ranges.len(), 3); assert!(ranges.iter().any(|(name, size)| name == "range_a" && *size == 50)); assert!(ranges.iter().any(|(name, size)| name == "range_b" && *size == 100)); assert!(ranges.iter().any(|(name, size)| name == "range_c" && *size == 25)); // List assignments in each range let assignments_a = store.list_range("range_a")?; assert_eq!(assignments_a.len(), 1); assert_eq!(assignments_a[0], (bit_a, "value_a".to_string())); let assignments_b = store.list_range("range_b")?; assert_eq!(assignments_b.len(), 1); assert_eq!(assignments_b[0], (bit_b, "value_b".to_string())); let assignments_c = store.list_range("range_c")?; assert_eq!(assignments_c.len(), 1); assert_eq!(assignments_c[0], (bit_c, "value_c".to_string())); Ok(()) } #[test] fn test_sequential_bit_assignment() -> Result<()> { let store = create_test_store()?; // Define a range store.define("sequential", 10)?; // Assign values and verify they get sequential bits let mut bits = Vec::new(); for i in 0..5 { let bit = store.assign("sequential", &format!("seq_{}", i))?; bits.push(bit); } // Should get bits 0, 1, 2, 3, 4 in that order for (i, &bit) in bits.iter().enumerate() { assert_eq!(bit, i as u64); } // Unassign bit 2 store.unassign("sequential", "seq_2")?; // Next assignment should reuse bit 2 let next_bit = store.assign("sequential", "reused")?; assert_eq!(next_bit, 2); Ok(()) } #[test] fn test_define_range() -> Result<()> { let store = create_test_store()?; // Define a range store.define("test_range", 100)?; // Defining the same range again should not error store.define("test_range", 100)?; // Check that ranges are listed let ranges = store.list_ranges()?; assert!(ranges.iter().any(|(name, size)| name == "test_range" && *size == 100)); Ok(()) } #[test] fn test_assign_and_get() -> Result<()> { let store = create_test_store()?; // Define a range first store.define("test_range", 100)?; // Assign a value let bit_pos = store.assign("test_range", "test_value")?; assert!(bit_pos < 100); // Get the value let value = store.get("test_range", bit_pos)?; assert_eq!(value, Some("test_value".to_string())); // Test getting non-existent bit let empty = store.get("test_range", 99)?; assert_eq!(empty, None); Ok(()) } #[test] fn test_multiple_assignments() -> Result<()> { let store = create_test_store()?; // Define range store.define("test_range", 64)?; // Assign multiple values let pos1 = store.assign("test_range", "value1")?; let pos2 = store.assign("test_range", "value2")?; let pos3 = store.assign("test_range", "value3")?; assert_eq!(pos1, 0); assert_eq!(pos2, 1); assert_eq!(pos3, 2); // Verify all values assert_eq!(store.get("test_range", 0)?, Some("value1".to_string())); assert_eq!(store.get("test_range", 1)?, Some("value2".to_string())); assert_eq!(store.get("test_range", 2)?, Some("value3".to_string())); Ok(()) } #[test] fn test_list_range() -> Result<()> { let store = create_test_store()?; // Define range store.define("test_range", 64)?; // Assign some values store.assign("test_range", "alpha")?; store.assign("test_range", "beta")?; store.assign("test_range", "gamma")?; // List assignments let assignments = store.list_range("test_range")?; assert_eq!(assignments.len(), 3); assert_eq!(assignments[0], (0, "alpha".to_string())); assert_eq!(assignments[1], (1, "beta".to_string())); assert_eq!(assignments[2], (2, "gamma".to_string())); Ok(()) } #[test] fn test_list_ranges() -> Result<()> { let store = create_test_store()?; // Define multiple ranges store.define("range_a", 32)?; store.define("range_b", 64)?; store.define("range_c", 128)?; // List all ranges let ranges = store.list_ranges()?; assert_eq!(ranges.len(), 3); assert_eq!(ranges[0], ("range_a".to_string(), 32)); assert_eq!(ranges[1], ("range_b".to_string(), 64)); assert_eq!(ranges[2], ("range_c".to_string(), 128)); Ok(()) } #[test] fn test_unassign() -> Result<()> { let store = create_test_store()?; // Define range store.define("test_range", 64)?; // Assign values store.assign("test_range", "value1")?; store.assign("test_range", "value2")?; store.assign("test_range", "value3")?; // Unassign middle value let removed = store.unassign("test_range", "value2")?; assert!(removed); // Verify it's gone assert_eq!(store.get("test_range", 1)?, None); // Other values should still be there assert_eq!(store.get("test_range", 0)?, Some("value1".to_string())); assert_eq!(store.get("test_range", 2)?, Some("value3".to_string())); // Assign a new value - should reuse the freed slot let new_pos = store.assign("test_range", "new_value")?; assert_eq!(new_pos, 1); // Should reuse position 1 Ok(()) } #[test] fn test_range_full() -> Result<()> { let store = create_test_store()?; // Define small range store.define("small_range", 2)?; // Fill it up store.assign("small_range", "value1")?; store.assign("small_range", "value2")?; // Try to assign one more - should fail let result = store.assign("small_range", "value3"); assert!(result.is_err()); Ok(()) } #[test] fn test_undefined_range() -> Result<()> { let store = create_test_store()?; // Try operations on undefined range let result = store.assign("undefined_range", "value"); assert!(result.is_err()); let value = store.get("undefined_range", 0)?; assert_eq!(value, None); let assignments = store.list_range("undefined_range")?; assert_eq!(assignments.len(), 0); let removed = store.unassign("undefined_range", "value")?; assert!(!removed); Ok(()) } + + #[test] + fn test_delete_range() -> Result<()> { + let store = create_test_store()?; + + // Define a range and make some assignments + store.define("delete_test", 100)?; + store.assign("delete_test", "value1")?; + store.assign("delete_test", "value2")?; + store.assign("delete_test", "value3")?; + + // Verify assignments exist + let assignments = store.list_range("delete_test")?; + assert_eq!(assignments.len(), 3); + + // Verify range exists + let range_exists = store.names.get("delete_test".as_bytes())?.is_some(); + assert!(range_exists); + + // Delete the range + let deleted = store.delete_range("delete_test")?; + assert!(deleted); + + // Verify range no longer exists + let range_exists = store.names.get("delete_test".as_bytes())?.is_some(); + assert!(!range_exists); + + // Verify assignments are gone + let assignments = store.list_range("delete_test")?; + assert_eq!(assignments.len(), 0); + + // Test deleting non-existent range + let deleted = store.delete_range("nonexistent")?; + assert!(!deleted); + + Ok(()) + } + + #[test] + fn test_delete_range_comprehensive() -> Result<()> { + let store = create_test_store()?; + + // Test deleting empty range + store.define("empty_range", 50)?; + let deleted = store.delete_range("empty_range")?; + assert!(deleted); + + // Verify it's gone + let range_exists = store.names.get("empty_range".as_bytes())?.is_some(); + assert!(!range_exists); + + // Test deleting range with single assignment + store.define("single_range", 10)?; + store.assign("single_range", "single_value")?; + + let assignments = store.list_range("single_range")?; + assert_eq!(assignments.len(), 1); + + let deleted = store.delete_range("single_range")?; + assert!(deleted); + + let assignments = store.list_range("single_range")?; + assert_eq!(assignments.len(), 0); + + // Test deleting range with maximum assignments + store.define("full_range", 3)?; + store.assign("full_range", "val1")?; + store.assign("full_range", "val2")?; + store.assign("full_range", "val3")?; + + let assignments = store.list_range("full_range")?; + assert_eq!(assignments.len(), 3); + + let deleted = store.delete_range("full_range")?; + assert!(deleted); + + let assignments = store.list_range("full_range")?; + assert_eq!(assignments.len(), 0); + + // Test that we can recreate a range with the same name after deletion + store.define("recreate_test", 20)?; + store.assign("recreate_test", "original")?; + + let deleted = store.delete_range("recreate_test")?; + assert!(deleted); + + // Recreate with different size + store.define("recreate_test", 30)?; + store.assign("recreate_test", "new_value")?; + + let assignments = store.list_range("recreate_test")?; + assert_eq!(assignments.len(), 1); + assert_eq!(assignments[0].1, "new_value"); + + // Test multiple deletions don't interfere + store.define("multi1", 5)?; + store.define("multi2", 5)?; + store.assign("multi1", "m1_val1")?; + store.assign("multi2", "m2_val1")?; + + let deleted1 = store.delete_range("multi1")?; + assert!(deleted1); + + // multi2 should still exist + let assignments = store.list_range("multi2")?; + assert_eq!(assignments.len(), 1); + + let deleted2 = store.delete_range("multi2")?; + assert!(deleted2); + + let assignments = store.list_range("multi2")?; + assert_eq!(assignments.len(), 0); + + Ok(()) + } + + #[test] + fn test_delete_range_atomicity() -> Result<()> { + let store = create_test_store()?; + + // Create multiple ranges to verify isolation + store.define("atomic_test1", 100)?; + store.define("atomic_test2", 50)?; + store.assign("atomic_test1", "value1")?; + store.assign("atomic_test1", "value2")?; + store.assign("atomic_test2", "other_value")?; + + // Verify initial state + let assignments1 = store.list_range("atomic_test1")?; + assert_eq!(assignments1.len(), 2); + let assignments2 = store.list_range("atomic_test2")?; + assert_eq!(assignments2.len(), 1); + + // Delete one range - this should be atomic + let deleted = store.delete_range("atomic_test1")?; + assert!(deleted); + + // Verify first range is completely gone + let range1_exists = store.names.get("atomic_test1".as_bytes())?.is_some(); + assert!(!range1_exists); + let assignments1 = store.list_range("atomic_test1")?; + assert_eq!(assignments1.len(), 0); + + // Verify second range is unaffected + let range2_exists = store.names.get("atomic_test2".as_bytes())?.is_some(); + assert!(range2_exists); + let assignments2 = store.list_range("atomic_test2")?; + assert_eq!(assignments2.len(), 1); + assert_eq!(assignments2[0].1, "other_value"); + + Ok(()) + } }