distributed_lock_postgres/
key.rs

1//! PostgreSQL advisory lock key encoding.
2
3use distributed_lock_core::error::{LockError, LockResult};
4use sha2::{Digest, Sha256};
5
6/// Key for PostgreSQL advisory locks.
7///
8/// Advisory locks use either a single 64-bit key or a pair of 32-bit keys.
9/// These represent different key spaces and do not overlap.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11pub enum PostgresAdvisoryLockKey {
12    /// Single 64-bit key.
13    Single(i64),
14    /// Pair of 32-bit keys.
15    Pair(i32, i32),
16}
17
18impl PostgresAdvisoryLockKey {
19    /// Maximum length for ASCII encoding (9 characters).
20    const MAX_ASCII_LENGTH: usize = 9;
21    /// Bits per ASCII character (7 bits).
22    const ASCII_CHAR_BITS: u32 = 7;
23    /// Maximum ASCII value (127).
24    const MAX_ASCII_VALUE: u32 = (1 << Self::ASCII_CHAR_BITS) - 1;
25    /// Hash string length (16 hex chars for i64).
26    const HASH_STRING_LENGTH: usize = 16;
27    /// Hash part length (8 hex chars for i32).
28    const HASH_PART_LENGTH: usize = 8;
29    /// Hash string separator.
30    const HASH_STRING_SEPARATOR: char = ',';
31
32    /// Create a key from a string name.
33    ///
34    /// - ASCII strings up to 9 chars are encoded directly (collision-free)
35    /// - 16-char hex strings are parsed as i64
36    /// - "XXXXXXXX,XXXXXXXX" format parsed as (i32, i32)
37    /// - Other strings are hashed to i64 (if `allow_hashing` is true)
38    pub fn from_name(name: &str, allow_hashing: bool) -> LockResult<Self> {
39        if name.is_empty() {
40            return Err(LockError::InvalidName(
41                "lock name cannot be empty".to_string(),
42            ));
43        }
44
45        // Try ASCII encoding first
46        if let Some(key) = Self::try_encode_ascii(name) {
47            return Ok(Self::Single(key));
48        }
49
50        // Try parsing as hex string
51        if let Some(key) = Self::try_parse_hex_string(name) {
52            return Ok(key);
53        }
54
55        // Try parsing as pair format
56        if let Some(key) = Self::try_parse_pair_string(name) {
57            return Ok(key);
58        }
59
60        // Hash if allowed
61        if allow_hashing {
62            let hash = Self::hash_string(name);
63            return Ok(Self::Single(hash));
64        }
65
66        Err(LockError::InvalidName(format!(
67            "Name '{}' could not be encoded as a PostgresAdvisoryLockKey. Please specify allow_hashing or use one of the following formats: (1) a 0-{} character string using only ASCII characters, (2) a {} character hex string, or (3) a 2-part, {} character string of the form XXXXXXXX{}XXXXXXXX",
68            name,
69            Self::MAX_ASCII_LENGTH,
70            Self::HASH_STRING_LENGTH,
71            Self::HASH_PART_LENGTH * 2 + 1,
72            Self::HASH_STRING_SEPARATOR
73        )))
74    }
75
76    /// Try to encode as ASCII string (up to 9 chars).
77    fn try_encode_ascii(name: &str) -> Option<i64> {
78        if name.len() > Self::MAX_ASCII_LENGTH {
79            return None;
80        }
81
82        let mut result = 0i64;
83        for ch in name.chars() {
84            let ch_val = ch as u32;
85            if ch_val > Self::MAX_ASCII_VALUE {
86                return None;
87            }
88            result = (result << Self::ASCII_CHAR_BITS) | (ch_val as i64);
89        }
90
91        // Add padding: shift by 1 (zero bit), then fill remaining with 1s
92        result <<= 1;
93        for _ in name.len()..Self::MAX_ASCII_LENGTH {
94            result = (result << Self::ASCII_CHAR_BITS) | (Self::MAX_ASCII_VALUE as i64);
95        }
96
97        Some(result)
98    }
99
100    /// Try to parse as hex string (16 chars for i64).
101    fn try_parse_hex_string(name: &str) -> Option<Self> {
102        if name.len() != Self::HASH_STRING_LENGTH {
103            return None;
104        }
105
106        i64::from_str_radix(name, 16).ok().map(Self::Single)
107    }
108
109    /// Try to parse as pair format "XXXXXXXX,XXXXXXXX".
110    fn try_parse_pair_string(name: &str) -> Option<Self> {
111        let parts: Vec<&str> = name.split(Self::HASH_STRING_SEPARATOR).collect();
112        if parts.len() != 2 {
113            return None;
114        }
115
116        let key1 = i32::from_str_radix(parts[0], 16).ok()?;
117        let key2 = i32::from_str_radix(parts[1], 16).ok()?;
118
119        Some(Self::Pair(key1, key2))
120    }
121
122    /// Hash a string to i64 using SHA-256 (taking first 8 bytes).
123    #[allow(clippy::disallowed_methods)]
124    fn hash_string(name: &str) -> i64 {
125        let mut hasher = Sha256::new();
126        hasher.update(name.as_bytes());
127        let hash_bytes = hasher.finalize();
128
129        // Take first 8 bytes and convert to i64 (little-endian)
130        let mut result = 0i64;
131        for i in (0..8).rev() {
132            result = (result << 8) | (hash_bytes[i] as i64);
133        }
134        result
135    }
136
137    /// Returns true if this is a single key.
138    pub fn has_single_key(&self) -> bool {
139        matches!(self, Self::Single(_))
140    }
141
142    /// Gets the single key value (panics if Pair).
143    pub fn key(&self) -> i64 {
144        match self {
145            Self::Single(k) => *k,
146            Self::Pair(_, _) => panic!("key() called on Pair variant"),
147        }
148    }
149
150    /// Gets the key pair (splits Single into two i32s).
151    pub fn keys(&self) -> (i32, i32) {
152        match self {
153            Self::Single(k) => {
154                // Split i64 into two i32s
155                let upper = (*k >> 32) as i32;
156                let lower = (*k & 0xFFFFFFFF) as i32;
157                (upper, lower)
158            }
159            Self::Pair(k1, k2) => (*k1, *k2),
160        }
161    }
162
163    /// Convert to SQL function arguments.
164    #[allow(clippy::wrong_self_convention)]
165    pub fn to_sql_args(&self) -> String {
166        match self {
167            Self::Single(k) => format!("{k}"),
168            Self::Pair(k1, k2) => format!("{k1}, {k2}"),
169        }
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    #[test]
178    fn test_ascii_encoding() {
179        let key = PostgresAdvisoryLockKey::from_name("test", false).unwrap();
180        assert!(key.has_single_key());
181    }
182
183    #[test]
184    fn test_hex_encoding() {
185        let hex_str = "0000000000000001";
186        let key = PostgresAdvisoryLockKey::from_name(hex_str, false).unwrap();
187        assert!(key.has_single_key());
188        assert_eq!(key.key(), 1);
189    }
190
191    #[test]
192    fn test_pair_encoding() {
193        let pair_str = "00000001,00000002";
194        let key = PostgresAdvisoryLockKey::from_name(pair_str, false).unwrap();
195        match key {
196            PostgresAdvisoryLockKey::Pair(k1, k2) => {
197                assert_eq!(k1, 1);
198                assert_eq!(k2, 2);
199            }
200            _ => panic!("Expected Pair variant"),
201        }
202    }
203
204    #[test]
205    fn test_hash_encoding() {
206        let long_name = "this is a very long lock name that needs hashing";
207        let key = PostgresAdvisoryLockKey::from_name(long_name, true).unwrap();
208        assert!(key.has_single_key());
209    }
210}
distributed_lock_postgres/key.rs

distributed_lock_postgres/
key.rs
