1  use std::fs;
  2  use std::fs::OpenOptions;
  3  use std::io::Write;
  4  use std::path::Path;
  5  
  6  use automerge::AutoCommit;
  7  
  8  const STORAGE_MAGIC: [u8; 8] = *b"PALUN\0\0\0";
  9  pub const STORAGE_VERSION: u16 = 1;
 10  const HEADER_LEN: usize = 8 + 2 + 2 + 4;
 11  
 12  #[derive(Debug)]
 13  pub enum StorageError {
 14      Io(std::io::Error),
 15      Automerge(automerge::AutomergeError),
 16      InvalidHeader(String),
 17      UnsupportedVersion(u16),
 18  }
 19  
 20  impl std::fmt::Display for StorageError {
 21      fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 22          match self {
 23              StorageError::Io(err) => write!(f, "io error: {err}"),
 24              StorageError::Automerge(err) => write!(f, "automerge error: {err}"),
 25              StorageError::InvalidHeader(message) => write!(f, "invalid header: {message}"),
 26              StorageError::UnsupportedVersion(version) => {
 27                  write!(f, "unsupported storage version {version}")
 28              }
 29          }
 30      }
 31  }
 32  
 33  impl std::error::Error for StorageError {
 34      fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
 35          match self {
 36              StorageError::Io(err) => Some(err),
 37              StorageError::Automerge(err) => Some(err),
 38              _ => None,
 39          }
 40      }
 41  }
 42  
 43  impl From<std::io::Error> for StorageError {
 44      fn from(err: std::io::Error) -> Self {
 45          StorageError::Io(err)
 46      }
 47  }
 48  
 49  impl From<automerge::AutomergeError> for StorageError {
 50      fn from(err: automerge::AutomergeError) -> Self {
 51          StorageError::Automerge(err)
 52      }
 53  }
 54  
 55  /// Save the Automerge document to disk with a versioned header.
 56  pub fn save(doc: &mut AutoCommit, path: impl AsRef<Path>) -> Result<(), StorageError> {
 57      let document = doc.save();
 58      if document.len() > u32::MAX as usize {
 59          return Err(StorageError::InvalidHeader(
 60              "document too large".to_string(),
 61          ));
 62      }
 63      let payload = encode_document(&document);
 64      atomic_write(path.as_ref(), &payload)?;
 65      Ok(())
 66  }
 67  
 68  /// Load an Automerge document from disk.
 69  pub fn load(path: impl AsRef<Path>) -> Result<AutoCommit, StorageError> {
 70      let bytes = fs::read(path)?;
 71      let document = decode_document(&bytes)?;
 72      Ok(AutoCommit::load(document)?)
 73  }
 74  
 75  fn encode_document(document: &[u8]) -> Vec<u8> {
 76      let mut buffer = Vec::with_capacity(HEADER_LEN + document.len());
 77      buffer.extend_from_slice(&STORAGE_MAGIC);
 78      buffer.extend_from_slice(&STORAGE_VERSION.to_le_bytes());
 79      buffer.extend_from_slice(&0u16.to_le_bytes());
 80      buffer.extend_from_slice(&(document.len() as u32).to_le_bytes());
 81      buffer.extend_from_slice(document);
 82      buffer
 83  }
 84  
 85  fn decode_document(bytes: &[u8]) -> Result<&[u8], StorageError> {
 86      if bytes.len() < HEADER_LEN {
 87          return Err(StorageError::InvalidHeader("file too small".to_string()));
 88      }
 89      if bytes[0..8] != STORAGE_MAGIC {
 90          return Err(StorageError::InvalidHeader("magic mismatch".to_string()));
 91      }
 92      let version = u16::from_le_bytes([bytes[8], bytes[9]]);
 93      if version != STORAGE_VERSION {
 94          return Err(StorageError::UnsupportedVersion(version));
 95      }
 96      let length = u32::from_le_bytes([bytes[12], bytes[13], bytes[14], bytes[15]]) as usize;
 97      let payload = &bytes[HEADER_LEN..];
 98      if payload.len() != length {
 99          return Err(StorageError::InvalidHeader("length mismatch".to_string()));
100      }
101      Ok(payload)
102  }
103  
104  fn atomic_write(path: &Path, data: &[u8]) -> Result<(), StorageError> {
105      if let Some(parent) = path.parent() {
106          if !parent.as_os_str().is_empty() {
107              fs::create_dir_all(parent)?;
108          }
109      }
110  
111      let file_name = path
112          .file_name()
113          .and_then(|name| name.to_str())
114          .unwrap_or("workspace");
115      let tmp_name = format!(".{file_name}.{}.tmp", uuid::Uuid::new_v4());
116      let tmp_path = path.with_file_name(tmp_name);
117  
118      {
119          let mut file = OpenOptions::new()
120              .write(true)
121              .create_new(true)
122              .open(&tmp_path)?;
123          file.write_all(data)?;
124          file.sync_all()?;
125      }
126  
127      if let Err(err) = fs::rename(&tmp_path, path) {
128          if err.kind() == std::io::ErrorKind::AlreadyExists {
129              fs::remove_file(path)?;
130              fs::rename(&tmp_path, path)?;
131          } else {
132              return Err(StorageError::Io(err));
133          }
134      }
135  
136      Ok(())
137  }