1use std::{
2 collections::HashMap,
3 fs::{File, OpenOptions},
4 io::{Read, Seek, SeekFrom, Write},
5 path::Path,
6 sync::Mutex,
7};
8
9pub use kcode_k1_access_types::{Authority, GroupId, TxId, UserId};
10pub use kcode_k1_kmap_format::NodeId;
11
12const HEADER: &[u8; 8] = b"K1LNV2\0\0";
13const USER_TAG: u8 = 1;
14const GROUP_TAG: u8 = 2;
15const MAX_TARGET_BYTES: usize = 320;
16const RECORD_PREFIX: usize = 15;
17const RECORD_SUFFIX: usize = 16;
18
19#[derive(Clone, Debug, Eq, Hash, PartialEq)]
20pub struct TargetName(String);
21
22impl TargetName {
23 pub fn new(value: String) -> Result<Self, String> {
24 let characters = value.chars().count();
25 if !(1..=80).contains(&characters) {
26 return Err("target name must contain 1 through 80 Unicode characters".into());
27 }
28 if value.chars().any(char::is_control) {
29 return Err("target name must not contain control characters".into());
30 }
31 if !value.chars().any(|character| !character.is_whitespace()) {
32 return Err("target name must contain a non-whitespace character".into());
33 }
34 if value.len() > MAX_TARGET_BYTES {
35 return Err("target name exceeds its canonical UTF-8 bound".into());
36 }
37 Ok(Self(value))
38 }
39
40 pub fn as_str(&self) -> &str {
41 &self.0
42 }
43
44 pub fn into_string(self) -> String {
45 self.0
46 }
47}
48
49#[derive(Clone, Debug, Eq, Hash, PartialEq)]
50pub struct TargetId {
51 authority: Authority,
52 target: TargetName,
53}
54
55impl TargetId {
56 pub fn new(authority: Authority, target: TargetName) -> Self {
57 Self { authority, target }
58 }
59
60 pub fn authority(&self) -> Authority {
61 self.authority
62 }
63
64 pub fn target(&self) -> &TargetName {
65 &self.target
66 }
67}
68
69struct Inner {
70 file: File,
71 bindings: HashMap<TargetId, NodeId>,
72 failed: bool,
73}
74
75pub struct LaunchNodes {
76 inner: Mutex<Inner>,
77}
78
79impl LaunchNodes {
80 pub fn open(path: &Path) -> Result<Self, String> {
81 let created = !path.exists();
82 let mut file = OpenOptions::new()
83 .create(true)
84 .truncate(false)
85 .read(true)
86 .write(true)
87 .open(path)
88 .map_err(|error| format!("open launch nodes file: {error}"))?;
89
90 if created {
91 file.write_all(HEADER)
92 .and_then(|()| file.sync_data())
93 .map_err(|error| format!("initialize launch nodes file: {error}"))?;
94 sync_parent(path)?;
95 }
96
97 file.seek(SeekFrom::Start(0))
98 .map_err(|error| format!("seek launch nodes file: {error}"))?;
99 let mut bytes = Vec::new();
100 file.read_to_end(&mut bytes)
101 .map_err(|error| format!("read launch nodes file: {error}"))?;
102 if bytes.len() < HEADER.len() || &bytes[..HEADER.len()] != HEADER {
103 return Err("unsupported launch nodes file format".into());
104 }
105
106 let mut bindings = HashMap::new();
107 let mut offset = HEADER.len();
108 while offset < bytes.len() {
109 let remaining = bytes.len() - offset;
110 if remaining < RECORD_PREFIX {
111 repair_tail(&mut file, offset)?;
112 break;
113 }
114 let tag = bytes[offset];
115 if tag != USER_TAG && tag != GROUP_TAG {
116 return Err("launch nodes record has an invalid authority tag".into());
117 }
118
119 let name_len = u16::from_le_bytes([bytes[offset + 13], bytes[offset + 14]]) as usize;
120 if name_len > MAX_TARGET_BYTES {
121 return Err("launch nodes record has an invalid target length".into());
122 }
123 let record_len = RECORD_PREFIX + name_len + RECORD_SUFFIX;
124 if remaining < record_len {
125 repair_tail(&mut file, offset)?;
126 break;
127 }
128
129 let record = &bytes[offset..offset + record_len];
130 let checksum_offset = record_len - 4;
131 let expected = u32::from_le_bytes(
132 record[checksum_offset..]
133 .try_into()
134 .map_err(|_| "launch nodes checksum is unavailable".to_owned())?,
135 );
136 if checksum(&record[..checksum_offset]) != expected {
137 return Err("launch nodes record checksum is invalid".into());
138 }
139
140 let authority_bytes: [u8; 12] = record[1..13]
141 .try_into()
142 .map_err(|_| "launch nodes authority is invalid".to_owned())?;
143 let authority_txid = TxId::from_bytes(authority_bytes);
144 let authority = match tag {
145 USER_TAG => Authority::User(UserId::from_tx_id(authority_txid)),
146 GROUP_TAG => Authority::Group(GroupId::new(authority_txid)),
147 _ => unreachable!(),
148 };
149 let name_start = RECORD_PREFIX;
150 let name_end = name_start + name_len;
151 let name = String::from_utf8(record[name_start..name_end].to_vec())
152 .map_err(|_| "launch nodes target is not valid UTF-8".to_owned())?;
153 let target = TargetName::new(name)
154 .map_err(|error| format!("launch nodes target is invalid: {error}"))?;
155 let node_bytes: [u8; 12] = record[name_end..name_end + 12]
156 .try_into()
157 .map_err(|_| "launch nodes node ID is invalid".to_owned())?;
158 bindings.insert(TargetId::new(authority, target), NodeId(node_bytes));
159 offset += record_len;
160 }
161
162 file.seek(SeekFrom::End(0))
163 .map_err(|error| format!("seek launch nodes append position: {error}"))?;
164 Ok(Self {
165 inner: Mutex::new(Inner {
166 file,
167 bindings,
168 failed: false,
169 }),
170 })
171 }
172
173 pub fn get(&self, target: &TargetId) -> Result<Option<NodeId>, String> {
174 let inner = self
175 .inner
176 .lock()
177 .map_err(|_| "launch nodes lock poisoned".to_owned())?;
178 if inner.failed {
179 return Err("launch nodes store is unavailable".into());
180 }
181 Ok(inner.bindings.get(target).copied())
182 }
183
184 pub fn set(&self, target: TargetId, node: NodeId) -> Result<(), String> {
185 let mut inner = self
186 .inner
187 .lock()
188 .map_err(|_| "launch nodes lock poisoned".to_owned())?;
189 if inner.failed {
190 return Err("launch nodes store is unavailable".into());
191 }
192 let record = encode_record(&target, node);
193 let result = inner
194 .file
195 .seek(SeekFrom::End(0))
196 .and_then(|_| inner.file.write_all(&record))
197 .and_then(|()| inner.file.sync_data());
198 if let Err(error) = result {
199 inner.failed = true;
200 return Err(format!("append launch nodes record: {error}"));
201 }
202 inner.bindings.insert(target, node);
203 Ok(())
204 }
205}
206
207fn encode_record(target: &TargetId, node: NodeId) -> Vec<u8> {
208 let (tag, authority) = match target.authority {
209 Authority::User(user) => (USER_TAG, *user.as_tx_id().as_bytes()),
210 Authority::Group(group) => (GROUP_TAG, *group.txid().as_bytes()),
211 };
212 let name = target.target.as_str().as_bytes();
213 let mut record = Vec::with_capacity(RECORD_PREFIX + name.len() + RECORD_SUFFIX);
214 record.push(tag);
215 record.extend_from_slice(&authority);
216 record.extend_from_slice(&(name.len() as u16).to_le_bytes());
217 record.extend_from_slice(name);
218 record.extend_from_slice(&node.0);
219 let checksum = checksum(&record);
220 record.extend_from_slice(&checksum.to_le_bytes());
221 record
222}
223
224fn checksum(bytes: &[u8]) -> u32 {
225 bytes.iter().fold(0x811c_9dc5, |value, byte| {
226 (value ^ u32::from(*byte)).wrapping_mul(0x0100_0193)
227 })
228}
229
230fn repair_tail(file: &mut File, offset: usize) -> Result<(), String> {
231 file.set_len(offset as u64)
232 .and_then(|()| file.sync_data())
233 .map_err(|error| format!("repair launch nodes incomplete tail: {error}"))
234}
235
236fn sync_parent(path: &Path) -> Result<(), String> {
237 let parent = path
238 .parent()
239 .filter(|parent| !parent.as_os_str().is_empty())
240 .unwrap_or_else(|| Path::new("."));
241 File::open(parent)
242 .and_then(|directory| directory.sync_all())
243 .map_err(|error| format!("synchronize launch nodes parent directory: {error}"))
244}
245
246#[cfg(test)]
247mod tests {
248 use super::*;
249 use std::{
250 fs,
251 path::PathBuf,
252 time::{SystemTime, UNIX_EPOCH},
253 };
254
255 fn path(label: &str) -> PathBuf {
256 let nonce = SystemTime::now()
257 .duration_since(UNIX_EPOCH)
258 .unwrap()
259 .as_nanos();
260 std::env::temp_dir().join(format!(
261 "k1-launch-nodes-{label}-{}-{nonce}",
262 std::process::id()
263 ))
264 }
265
266 fn tx(byte: u8) -> TxId {
267 TxId::from_bytes([byte; 12])
268 }
269
270 fn user(byte: u8) -> Authority {
271 Authority::User(UserId::from_tx_id(tx(byte)))
272 }
273
274 fn group(byte: u8) -> Authority {
275 Authority::Group(GroupId::new(tx(byte)))
276 }
277
278 fn target(authority: Authority, name: &str) -> TargetId {
279 TargetId::new(authority, TargetName::new(name.to_owned()).unwrap())
280 }
281
282 #[test]
283 fn target_names_preserve_exact_slash_and_unicode_boundaries() {
284 let exact = "界".repeat(80);
285 assert_eq!(
286 TargetName::new("model/chatgpt/5.6-sol/date/xhigh".to_owned())
287 .unwrap()
288 .as_str(),
289 "model/chatgpt/5.6-sol/date/xhigh"
290 );
291 assert_eq!(TargetName::new(exact.clone()).unwrap().into_string(), exact);
292 assert!(TargetName::new("界".repeat(81)).is_err());
293 assert!(TargetName::new(String::new()).is_err());
294 assert!(TargetName::new(" ".to_owned()).is_err());
295 assert!(TargetName::new("line\nbreak".to_owned()).is_err());
296 assert_eq!(
297 TargetName::new(" Mixed/Case ".to_owned())
298 .unwrap()
299 .into_string(),
300 " Mixed/Case "
301 );
302 }
303
304 #[test]
305 fn authorities_have_overlapping_names_and_replay_last_write_wins() {
306 let file = path("replay");
307 let user_target = target(user(1), "default/chat");
308 let group_target = target(group(1), "default/chat");
309 {
310 let store = LaunchNodes::open(&file).unwrap();
311 store.set(user_target.clone(), NodeId([1; 12])).unwrap();
312 store.set(group_target.clone(), NodeId([2; 12])).unwrap();
313 store.set(user_target.clone(), NodeId([3; 12])).unwrap();
314 }
315 let reopened = LaunchNodes::open(&file).unwrap();
316 assert_eq!(reopened.get(&user_target).unwrap(), Some(NodeId([3; 12])));
317 assert_eq!(reopened.get(&group_target).unwrap(), Some(NodeId([2; 12])));
318 fs::remove_file(file).unwrap();
319 }
320
321 #[test]
322 fn incomplete_tail_repairs_but_complete_corruption_fails() {
323 let partial = path("partial");
324 let key = target(user(2), "model/x");
325 {
326 let store = LaunchNodes::open(&partial).unwrap();
327 store.set(key.clone(), NodeId([4; 12])).unwrap();
328 }
329 let clean_len = fs::metadata(&partial).unwrap().len();
330 {
331 let mut file = OpenOptions::new().append(true).open(&partial).unwrap();
332 file.write_all(&[USER_TAG, 1, 2]).unwrap();
333 file.sync_data().unwrap();
334 }
335 let reopened = LaunchNodes::open(&partial).unwrap();
336 assert_eq!(reopened.get(&key).unwrap(), Some(NodeId([4; 12])));
337 assert_eq!(fs::metadata(&partial).unwrap().len(), clean_len);
338 fs::remove_file(partial).unwrap();
339
340 let corrupt = path("corrupt");
341 {
342 let store = LaunchNodes::open(&corrupt).unwrap();
343 store
344 .set(target(group(3), "default"), NodeId([5; 12]))
345 .unwrap();
346 }
347 let length = fs::metadata(&corrupt).unwrap().len();
348 {
349 let mut file = OpenOptions::new()
350 .read(true)
351 .write(true)
352 .open(&corrupt)
353 .unwrap();
354 file.seek(SeekFrom::Start(length - 1)).unwrap();
355 let mut byte = [0];
356 file.read_exact(&mut byte).unwrap();
357 file.seek(SeekFrom::Start(length - 1)).unwrap();
358 file.write_all(&[byte[0] ^ 0xff]).unwrap();
359 file.sync_data().unwrap();
360 }
361 assert!(LaunchNodes::open(&corrupt).is_err());
362 fs::remove_file(corrupt).unwrap();
363 }
364
365 #[test]
366 fn headerless_v1_is_rejected() {
367 let file = path("v1");
368 fs::write(&file, [0_u8; 28]).unwrap();
369 assert_eq!(
370 LaunchNodes::open(&file).err().unwrap(),
371 "unsupported launch nodes file format"
372 );
373 fs::remove_file(file).unwrap();
374 }
375
376 #[test]
377 fn complete_package_stays_below_the_managed_limit() {
378 let files = [
379 include_str!("../Cargo.toml"),
380 include_str!("../Documentation.md"),
381 include_str!("lib.rs"),
382 ];
383 let count = files
384 .iter()
385 .flat_map(|file| file.lines())
386 .filter(|line| !line.trim().is_empty())
387 .count();
388 assert!(count < 500, "complete package has {count} nonblank lines");
389 }
390}