promocrypt_core/

counter.rs

//! Counter management for promotional code generation.
//!
//! Supports multiple counter storage modes:
//! - File: Separate `.counter` file with OS locking
//! - InBin: Counter stored in the .bin mutable section
//! - Manual: User-provided counter values
//! - External: Consumer-managed counter (e.g., database column)

use std::fs::OpenOptions;
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};

use fs2::FileExt;
use serde::{Deserialize, Serialize};

use crate::error::{PromocryptError, Result};

/// Counter storage mode.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum CounterMode {
    /// Store counter in a separate file.
    File {
        /// Path to the counter file.
        path: String,
    },

    /// Store counter in the .bin file's mutable section.
    InBin,

    /// No storage - counter provided manually by caller.
    #[default]
    Manual,

    /// External counter management (e.g., database column).
    /// The consumer is responsible for counter storage and atomicity.
    External,
}

impl CounterMode {
    /// Create a file-based counter mode.
    pub fn file(path: impl Into<String>) -> Self {
        CounterMode::File { path: path.into() }
    }

    /// Check if this mode stores counter in-bin.
    pub fn is_in_bin(&self) -> bool {
        matches!(self, CounterMode::InBin)
    }

    /// Check if this mode requires external counter management.
    pub fn is_external(&self) -> bool {
        matches!(self, CounterMode::External | CounterMode::Manual)
    }

    /// Get the counter file path, if applicable.
    pub fn file_path(&self) -> Option<&str> {
        match self {
            CounterMode::File { path } => Some(path),
            _ => None,
        }
    }
}

/// Counter manager for different storage modes.
pub struct CounterManager {
    mode: CounterMode,
    /// In-memory counter for InBin mode
    in_bin_value: u64,
    /// Flag indicating if in-bin value has been modified
    in_bin_modified: bool,
}

impl CounterManager {
    /// Create a new counter manager.
    pub fn new(mode: CounterMode) -> Self {
        Self {
            mode,
            in_bin_value: 0,
            in_bin_modified: false,
        }
    }

    /// Create a counter manager with an initial in-bin value.
    pub fn with_in_bin_value(mode: CounterMode, initial_value: u64) -> Self {
        Self {
            mode,
            in_bin_value: initial_value,
            in_bin_modified: false,
        }
    }

    /// Get the counter mode.
    pub fn mode(&self) -> &CounterMode {
        &self.mode
    }

    /// Get the current counter value.
    pub fn get(&self) -> Result<u64> {
        match &self.mode {
            CounterMode::File { path } => read_counter_file(Path::new(path)),
            CounterMode::InBin => Ok(self.in_bin_value),
            CounterMode::Manual | CounterMode::External => Err(PromocryptError::InvalidArgument(
                "Cannot get counter in manual/external mode - counter must be provided".to_string(),
            )),
        }
    }

    /// Increment counter and return the starting value for a batch.
    ///
    /// For a batch of `count` codes, returns the starting counter value
    /// and increments the stored counter by `count`.
    pub fn reserve(&mut self, count: u64) -> Result<u64> {
        match &self.mode {
            CounterMode::File { path } => {
                let path = PathBuf::from(path);
                reserve_counter_file(&path, count)
            }
            CounterMode::InBin => {
                let start = self.in_bin_value;
                self.in_bin_value = self
                    .in_bin_value
                    .checked_add(count)
                    .ok_or(PromocryptError::CounterOverflow)?;
                self.in_bin_modified = true;
                Ok(start)
            }
            CounterMode::Manual | CounterMode::External => Err(PromocryptError::InvalidArgument(
                "Cannot reserve counter in manual/external mode".to_string(),
            )),
        }
    }

    /// Set the counter value (for initialization or reset).
    pub fn set(&mut self, value: u64) -> Result<()> {
        match &self.mode {
            CounterMode::File { path } => write_counter_file(Path::new(path), value),
            CounterMode::InBin => {
                self.in_bin_value = value;
                self.in_bin_modified = true;
                Ok(())
            }
            CounterMode::Manual | CounterMode::External => Err(PromocryptError::InvalidArgument(
                "Cannot set counter in manual/external mode".to_string(),
            )),
        }
    }

    /// Check if the in-bin counter has been modified.
    pub fn is_in_bin_modified(&self) -> bool {
        self.in_bin_modified && matches!(self.mode, CounterMode::InBin)
    }

    /// Get the in-bin counter value.
    pub fn in_bin_value(&self) -> u64 {
        self.in_bin_value
    }

    /// Mark the in-bin counter as saved.
    pub fn mark_saved(&mut self) {
        self.in_bin_modified = false;
    }

    /// Update the counter mode.
    pub fn set_mode(&mut self, mode: CounterMode) {
        self.mode = mode;
    }
}

/// Read counter value from a file.
///
/// Creates the file with value 0 if it doesn't exist.
pub fn read_counter_file(path: &Path) -> Result<u64> {
    // Create parent directory if needed
    if let Some(parent) = path.parent()
        && !parent.exists()
    {
        std::fs::create_dir_all(parent)?;
    }

    // Open or create file
    let mut file = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .truncate(false)
        .open(path)?;

    // Lock for reading
    file.lock_shared()
        .map_err(|_| PromocryptError::CounterLocked)?;

    let mut buffer = [0u8; 8];
    let bytes_read = file.read(&mut buffer)?;

    // Unlock
    file.unlock().ok();

    if bytes_read == 0 {
        // New file, initialize to 0
        return Ok(0);
    }

    if bytes_read != 8 {
        return Err(PromocryptError::InvalidFileFormatDetails(
            "Counter file corrupted".to_string(),
        ));
    }

    Ok(u64::from_le_bytes(buffer))
}

/// Write counter value to a file.
pub fn write_counter_file(path: &Path, value: u64) -> Result<()> {
    // Create parent directory if needed
    if let Some(parent) = path.parent()
        && !parent.exists()
    {
        std::fs::create_dir_all(parent)?;
    }

    let mut file = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .truncate(false)
        .open(path)?;

    // Lock for writing
    file.lock_exclusive()
        .map_err(|_| PromocryptError::CounterLocked)?;

    file.seek(SeekFrom::Start(0))?;
    file.write_all(&value.to_le_bytes())?;
    file.set_len(8)?;
    file.sync_all()?;

    // Unlock
    file.unlock().ok();

    Ok(())
}

/// Atomically reserve a range of counter values.
///
/// Returns the starting counter value and increments the file by `count`.
pub fn reserve_counter_file(path: &Path, count: u64) -> Result<u64> {
    // Create parent directory if needed
    if let Some(parent) = path.parent()
        && !parent.exists()
    {
        std::fs::create_dir_all(parent)?;
    }

    let mut file = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .truncate(false)
        .open(path)?;

    // Lock for exclusive access
    file.lock_exclusive()
        .map_err(|_| PromocryptError::CounterLocked)?;

    // Read current value
    let mut buffer = [0u8; 8];
    let bytes_read = file.read(&mut buffer)?;

    let current = if bytes_read == 0 {
        0
    } else if bytes_read == 8 {
        u64::from_le_bytes(buffer)
    } else {
        file.unlock().ok();
        return Err(PromocryptError::InvalidFileFormatDetails(
            "Counter file corrupted".to_string(),
        ));
    };

    // Calculate new value
    let new_value = current
        .checked_add(count)
        .ok_or(PromocryptError::CounterOverflow)?;

    // Write new value
    file.seek(SeekFrom::Start(0))?;
    file.write_all(&new_value.to_le_bytes())?;
    file.set_len(8)?;
    file.sync_all()?;

    // Unlock
    file.unlock().ok();

    Ok(current)
}

/// Increment counter by a specific amount.
///
/// Returns the new counter value.
pub fn increment_counter_file(path: &Path, amount: u64) -> Result<u64> {
    let start = reserve_counter_file(path, amount)?;
    Ok(start + amount)
}

/// Serialize counter value to bytes (for in-bin storage).
pub fn counter_to_bytes(value: u64) -> [u8; 8] {
    value.to_le_bytes()
}

/// Deserialize counter value from bytes.
pub fn counter_from_bytes(bytes: &[u8]) -> Result<u64> {
    if bytes.len() < 8 {
        return Err(PromocryptError::InvalidFileFormatDetails(
            "Counter data too short".to_string(),
        ));
    }

    let mut arr = [0u8; 8];
    arr.copy_from_slice(&bytes[..8]);
    Ok(u64::from_le_bytes(arr))
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    #[test]
    fn test_counter_mode_default() {
        let mode = CounterMode::default();
        assert!(matches!(mode, CounterMode::Manual));
    }

    #[test]
    fn test_counter_mode_file() {
        let mode = CounterMode::file("/tmp/test.counter");
        assert!(matches!(mode, CounterMode::File { .. }));
        assert_eq!(mode.file_path(), Some("/tmp/test.counter"));
    }

    #[test]
    fn test_counter_manager_manual() {
        let manager = CounterManager::new(CounterMode::Manual);
        assert!(manager.get().is_err());
    }

    #[test]
    fn test_counter_manager_in_bin() {
        let mut manager = CounterManager::with_in_bin_value(CounterMode::InBin, 100);

        assert_eq!(manager.get().unwrap(), 100);

        let start = manager.reserve(10).unwrap();
        assert_eq!(start, 100);
        assert_eq!(manager.get().unwrap(), 110);
        assert!(manager.is_in_bin_modified());
    }

    #[test]
    fn test_counter_file_read_write() {
        let temp_dir = TempDir::new().unwrap();
        let path = temp_dir.path().join("test.counter");

        // Initially should be 0 (new file)
        let value = read_counter_file(&path).unwrap();
        assert_eq!(value, 0);

        // Write a value
        write_counter_file(&path, 12345).unwrap();

        // Read it back
        let value = read_counter_file(&path).unwrap();
        assert_eq!(value, 12345);
    }

    #[test]
    fn test_counter_file_reserve() {
        let temp_dir = TempDir::new().unwrap();
        let path = temp_dir.path().join("test.counter");

        // Reserve 100
        let start1 = reserve_counter_file(&path, 100).unwrap();
        assert_eq!(start1, 0);

        // Reserve 50 more
        let start2 = reserve_counter_file(&path, 50).unwrap();
        assert_eq!(start2, 100);

        // Verify final value
        let final_value = read_counter_file(&path).unwrap();
        assert_eq!(final_value, 150);
    }

    #[test]
    fn test_counter_file_increment() {
        let temp_dir = TempDir::new().unwrap();
        let path = temp_dir.path().join("test.counter");

        let new_value = increment_counter_file(&path, 10).unwrap();
        assert_eq!(new_value, 10);

        let new_value = increment_counter_file(&path, 5).unwrap();
        assert_eq!(new_value, 15);
    }

    #[test]
    fn test_counter_bytes_roundtrip() {
        let value = 0x123456789ABCDEF0u64;
        let bytes = counter_to_bytes(value);
        let recovered = counter_from_bytes(&bytes).unwrap();
        assert_eq!(value, recovered);
    }

    #[test]
    fn test_counter_overflow() {
        let mut manager = CounterManager::with_in_bin_value(CounterMode::InBin, u64::MAX - 5);

        // This should fail due to overflow
        let result = manager.reserve(10);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            PromocryptError::CounterOverflow
        ));
    }

    #[test]
    fn test_counter_manager_file_mode() {
        let temp_dir = TempDir::new().unwrap();
        let path = temp_dir.path().join("test.counter");
        let path_str = path.to_str().unwrap().to_string();

        let mut manager = CounterManager::new(CounterMode::File { path: path_str });

        // Initial value
        let value = manager.get().unwrap();
        assert_eq!(value, 0);

        // Reserve
        let start = manager.reserve(100).unwrap();
        assert_eq!(start, 0);

        // Get again
        let value = manager.get().unwrap();
        assert_eq!(value, 100);
    }

    #[test]
    fn test_counter_mode_is_external() {
        assert!(
            !CounterMode::File {
                path: "test".to_string()
            }
            .is_external()
        );
        assert!(!CounterMode::InBin.is_external());
        assert!(CounterMode::Manual.is_external());
        assert!(CounterMode::External.is_external());
    }

    #[test]
    fn test_counter_mode_serde() {
        let modes = vec![
            CounterMode::File {
                path: "/tmp/test.counter".to_string(),
            },
            CounterMode::InBin,
            CounterMode::Manual,
            CounterMode::External,
        ];

        for mode in modes {
            let json = serde_json::to_string(&mode).unwrap();
            let recovered: CounterMode = serde_json::from_str(&json).unwrap();
            assert_eq!(mode, recovered);
        }
    }
}