distributed_lock_postgres/
key.rs

1//! PostgreSQL advisory lock key encoding.
2
3use distributed_lock_core::error::{LockError, LockResult};
4use sha1::{Digest, Sha1};
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-1 (taking first 8 bytes).
123    #[allow(clippy::disallowed_methods)]
124    fn hash_string(name: &str) -> i64 {
125        let mut hasher = Sha1::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        // This matches the C# implementation:
131        // for (var i = sizeof(long) - 1; i >= 0; --i) { result = (result << 8) | hashBytes[i]; }
132        let mut result = 0i64;
133        for i in (0..8).rev() {
134            result = (result << 8) | (hash_bytes[i] as i64);
135        }
136        result
137    }
138
139    /// Returns true if this is a single key.
140    pub fn has_single_key(&self) -> bool {
141        matches!(self, Self::Single(_))
142    }
143
144    /// Gets the single key value (panics if Pair).
145    pub fn key(&self) -> i64 {
146        match self {
147            Self::Single(k) => *k,
148            Self::Pair(_, _) => panic!("key() called on Pair variant"),
149        }
150    }
151
152    /// Gets the key pair (splits Single into two i32s).
153    pub fn keys(&self) -> (i32, i32) {
154        match self {
155            Self::Single(k) => {
156                // Split i64 into two i32s
157                let upper = (*k >> 32) as i32;
158                let lower = (*k & 0xFFFFFFFF) as i32;
159                (upper, lower)
160            }
161            Self::Pair(k1, k2) => (*k1, *k2),
162        }
163    }
164
165    /// Convert to SQL function arguments.
166    #[allow(clippy::wrong_self_convention)]
167    pub fn to_sql_args(&self) -> String {
168        match self {
169            Self::Single(k) => format!("{k}"),
170            Self::Pair(k1, k2) => format!("{k1}, {k2}"),
171        }
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    #[test]
180    fn test_ascii_encoding() {
181        let key = PostgresAdvisoryLockKey::from_name("test", false).unwrap();
182        assert!(key.has_single_key());
183    }
184
185    #[test]
186    fn test_hex_encoding() {
187        let hex_str = "0000000000000001";
188        let key = PostgresAdvisoryLockKey::from_name(hex_str, false).unwrap();
189        assert!(key.has_single_key());
190        assert_eq!(key.key(), 1);
191    }
192
193    #[test]
194    fn test_pair_encoding() {
195        let pair_str = "00000001,00000002";
196        let key = PostgresAdvisoryLockKey::from_name(pair_str, false).unwrap();
197        match key {
198            PostgresAdvisoryLockKey::Pair(k1, k2) => {
199                assert_eq!(k1, 1);
200                assert_eq!(k2, 2);
201            }
202            _ => panic!("Expected Pair variant"),
203        }
204    }
205
206    #[test]
207    fn test_hash_encoding() {
208        let long_name = "this is a very long lock name that needs hashing";
209        let key = PostgresAdvisoryLockKey::from_name(long_name, true).unwrap();
210        assert!(key.has_single_key());
211    }
212}