use crate::crypto::encryption::EncryptionManager;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Reference {
pub name: String,
pub hash: String,
}
pub fn get_lit_dir(repo_path: &Path) -> PathBuf {
repo_path.join(".lit")
}
fn ref_encryption(repo_path: &Path) -> EncryptionManager {
let config = crate::crypto::encryption::EncryptionConfig::load(repo_path).unwrap_or_default();
EncryptionManager::new_auto(config, repo_path)
}
fn write_ref_file(path: &Path, repo_path: &Path, text: &str) -> Result<(), String> {
let data = ref_encryption(repo_path).encrypt(text.as_bytes())?;
fs::write(path, data).map_err(|e| format!("Failed to write reference: {}", e))
}
fn read_ref_file(path: &Path, repo_path: &Path) -> Result<String, String> {
let data = fs::read(path).map_err(|e| format!("Failed to read reference: {}", e))?;
let plain = if EncryptionManager::is_encrypted_payload(&data) {
ref_encryption(repo_path).decrypt(&data)?
} else {
data
};
String::from_utf8(plain)
.map_err(|e| format!("Reference is not valid UTF-8: {}", e))
.map(|s| s.trim().to_string())
}
fn refs_index_path(repo_path: &Path) -> PathBuf {
get_lit_dir(repo_path).join("refs.enc")
}
fn encryption_enabled(repo_path: &Path) -> bool {
crate::crypto::encryption::EncryptionConfig::load(repo_path)
.map(|config| config.enabled)
.unwrap_or(false)
}
fn load_refs_index(repo_path: &Path) -> Result<BTreeMap<String, String>, String> {
let path = refs_index_path(repo_path);
if !path.exists() {
return Ok(BTreeMap::new());
}
let data = fs::read(&path).map_err(|e| format!("Failed to read ref index: {}", e))?;
let plain = if EncryptionManager::is_encrypted_payload(&data) {
ref_encryption(repo_path).decrypt(&data)?
} else {
data
};
serde_json::from_slice(&plain).map_err(|e| format!("Failed to parse ref index: {}", e))
}
fn save_refs_index(repo_path: &Path, refs: &BTreeMap<String, String>) -> Result<(), String> {
let json =
serde_json::to_vec(refs).map_err(|e| format!("Failed to serialize ref index: {}", e))?;
let data = ref_encryption(repo_path).encrypt(&json)?;
fs::write(refs_index_path(repo_path), data)
.map_err(|e| format!("Failed to write ref index: {}", e))
}
pub fn is_lit_repo(path: &Path) -> bool {
get_lit_dir(path).exists()
}
pub fn find_repo_root() -> Result<PathBuf, String> {
let mut current =
std::env::current_dir().map_err(|e| format!("Failed to get current directory: {}", e))?;
loop {
if is_lit_repo(¤t) {
return Ok(current);
}
if !current.pop() {
return Err("Not in a Lit repository".to_string());
}
}
}
pub fn read_ref(repo_path: &Path, ref_name: &str) -> Result<String, String> {
if encryption_enabled(repo_path) {
if let Some(hash) = load_refs_index(repo_path)?.get(ref_name) {
return Ok(hash.clone());
}
}
let ref_path = get_lit_dir(repo_path).join("refs").join(ref_name);
if !ref_path.exists() {
return Err(format!("Reference '{}' not found", ref_name));
}
read_ref_file(&ref_path, repo_path)
}
pub fn read_ref_encrypted(
repo_path: &Path,
ref_name: &str,
encryption: &Arc<Mutex<EncryptionManager>>,
) -> Result<String, String> {
let ref_path = get_lit_dir(repo_path).join("refs").join(ref_name);
if !ref_path.exists() {
return Err(format!("Reference '{}' not found", ref_name));
}
let encrypted_data =
fs::read(&ref_path).map_err(|e| format!("Failed to read reference: {}", e))?;
let enc_guard = encryption
.lock()
.map_err(|_| "Failed to lock encryption manager".to_string())?;
let decrypted = enc_guard.decrypt(&encrypted_data)?;
String::from_utf8(decrypted)
.map_err(|e| format!("Invalid UTF-8 in decrypted reference: {}", e))
.map(|s| s.trim().to_string())
}
pub fn write_ref(repo_path: &Path, ref_name: &str, hash: &str) -> Result<(), String> {
if encryption_enabled(repo_path) {
let mut refs = load_refs_index(repo_path)?;
refs.insert(ref_name.to_string(), hash.to_string());
return save_refs_index(repo_path, &refs);
}
let ref_path = get_lit_dir(repo_path).join("refs").join(ref_name);
if let Some(parent) = ref_path.parent() {
fs::create_dir_all(parent).map_err(|e| format!("Failed to create ref directory: {}", e))?;
}
write_ref_file(&ref_path, repo_path, &format!("{}\n", hash))
}
pub fn write_ref_encrypted(
repo_path: &Path,
ref_name: &str,
hash: &str,
encryption: &Arc<Mutex<EncryptionManager>>,
) -> Result<(), String> {
let ref_path = get_lit_dir(repo_path).join("refs").join(ref_name);
if let Some(parent) = ref_path.parent() {
fs::create_dir_all(parent).map_err(|e| format!("Failed to create ref directory: {}", e))?;
}
let data = format!("{}\n", hash);
let enc_guard = encryption
.lock()
.map_err(|_| "Failed to lock encryption manager".to_string())?;
let encrypted = enc_guard.encrypt(data.as_bytes())?;
drop(enc_guard);
fs::write(&ref_path, encrypted).map_err(|e| format!("Failed to write reference: {}", e))
}
pub fn delete_ref(repo_path: &Path, ref_name: &str) -> Result<(), String> {
let ref_path = get_lit_dir(repo_path).join("refs").join(ref_name);
let mut removed = false;
if encryption_enabled(repo_path) {
let mut refs = load_refs_index(repo_path)?;
if refs.remove(ref_name).is_some() {
save_refs_index(repo_path, &refs)?;
removed = true;
}
}
if ref_path.exists() {
fs::remove_file(&ref_path).map_err(|e| format!("Failed to delete reference: {}", e))?;
removed = true;
}
if removed {
Ok(())
} else {
Err(format!("Reference '{}' not found", ref_name))
}
}
pub fn list_refs(repo_path: &Path, prefix: &str) -> Result<Vec<Reference>, String> {
let refs_dir = get_lit_dir(repo_path).join("refs").join(prefix);
let mut refs = Vec::new();
let mut seen = std::collections::HashSet::new();
if encryption_enabled(repo_path) {
let with_slash = format!("{}/", prefix);
for (name, hash) in load_refs_index(repo_path)? {
if let Some(short) = name.strip_prefix(&with_slash) {
seen.insert(short.to_string());
refs.push(Reference {
name: short.to_string(),
hash,
});
}
}
}
if !refs_dir.exists() {
return Ok(refs);
}
for entry in walkdir::WalkDir::new(&refs_dir) {
let entry = entry.map_err(|e| format!("Failed to read refs: {}", e))?;
if entry.file_type().is_file() {
let path = entry.path();
let name = path
.strip_prefix(&refs_dir)
.map_err(|e| format!("Path error: {}", e))?
.to_string_lossy()
.to_string();
if seen.contains(&name) {
continue;
}
let hash = read_ref_file(path, repo_path)?;
refs.push(Reference { name, hash });
}
}
Ok(refs)
}
pub fn read_head(repo_path: &Path) -> Result<String, String> {
let head_path = get_lit_dir(repo_path).join("HEAD");
if !head_path.exists() {
return Err("HEAD not found".to_string());
}
let content = read_ref_file(&head_path, repo_path)?;
let content = content.trim();
if let Some(ref_name) = content.strip_prefix("ref: ") {
read_ref(
repo_path,
ref_name.strip_prefix("refs/").unwrap_or(ref_name),
)
} else {
Ok(content.to_string())
}
}
pub fn get_current_branch(repo_path: &Path) -> Result<String, String> {
let head_path = get_lit_dir(repo_path).join("HEAD");
let content = read_ref_file(&head_path, repo_path)?;
let content = content.trim();
if let Some(branch) = content.strip_prefix("ref: refs/heads/") {
Ok(branch.to_string())
} else {
Err("HEAD is detached".to_string())
}
}
pub fn update_head(repo_path: &Path, branch: &str) -> Result<(), String> {
let head_path = get_lit_dir(repo_path).join("HEAD");
write_ref_file(
&head_path,
repo_path,
&format!("ref: refs/heads/{}\n", branch),
)
}
pub fn update_head_encrypted(
repo_path: &Path,
branch: &str,
encryption: &Arc<Mutex<EncryptionManager>>,
) -> Result<(), String> {
let head_path = get_lit_dir(repo_path).join("HEAD");
let data = format!("ref: refs/heads/{}\n", branch);
let enc_guard = encryption
.lock()
.map_err(|_| "Failed to lock encryption manager".to_string())?;
let encrypted = enc_guard.encrypt(data.as_bytes())?;
drop(enc_guard);
fs::write(&head_path, encrypted).map_err(|e| format!("Failed to update HEAD: {}", e))
}
pub fn set_head_detached(repo_path: &Path, hash: &str) -> Result<(), String> {
let head_path = get_lit_dir(repo_path).join("HEAD");
write_ref_file(&head_path, repo_path, &format!("{}\n", hash))
}
pub fn set_head_detached_encrypted(
repo_path: &Path,
hash: &str,
encryption: &Arc<Mutex<EncryptionManager>>,
) -> Result<(), String> {
let head_path = get_lit_dir(repo_path).join("HEAD");
let data = format!("{}\n", hash);
let enc_guard = encryption
.lock()
.map_err(|_| "Failed to lock encryption manager".to_string())?;
let encrypted = enc_guard.encrypt(data.as_bytes())?;
drop(enc_guard);
fs::write(&head_path, encrypted).map_err(|e| format!("Failed to set HEAD: {}", e))
}
pub fn read_head_encrypted(
repo_path: &Path,
encryption: &Arc<Mutex<EncryptionManager>>,
) -> Result<String, String> {
let head_path = get_lit_dir(repo_path).join("HEAD");
if !head_path.exists() {
return Err("HEAD not found".to_string());
}
let encrypted_data = fs::read(&head_path).map_err(|e| format!("Failed to read HEAD: {}", e))?;
let enc_guard = encryption
.lock()
.map_err(|_| "Failed to lock encryption manager".to_string())?;
let decrypted = enc_guard.decrypt(&encrypted_data)?;
drop(enc_guard);
let content = String::from_utf8(decrypted)
.map_err(|e| format!("Invalid UTF-8 in decrypted HEAD: {}", e))?;
let content = content.trim();
if let Some(ref_name) = content.strip_prefix("ref: ") {
read_ref_encrypted(
repo_path,
ref_name.strip_prefix("refs/").unwrap_or(ref_name),
encryption,
)
} else {
Ok(content.to_string())
}
}
pub fn get_current_branch_encrypted(
repo_path: &Path,
encryption: &Arc<Mutex<EncryptionManager>>,
) -> Result<String, String> {
let head_path = get_lit_dir(repo_path).join("HEAD");
let encrypted_data = fs::read(&head_path).map_err(|e| format!("Failed to read HEAD: {}", e))?;
let enc_guard = encryption
.lock()
.map_err(|_| "Failed to lock encryption manager".to_string())?;
let decrypted = enc_guard.decrypt(&encrypted_data)?;
drop(enc_guard);
let content = String::from_utf8(decrypted)
.map_err(|e| format!("Invalid UTF-8 in decrypted HEAD: {}", e))?;
let content = content.trim();
if let Some(branch) = content.strip_prefix("ref: refs/heads/") {
Ok(branch.to_string())
} else {
Err("HEAD is detached".to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::crypto::encryption::EncryptionConfig;
use tempfile::TempDir;
#[test]
fn test_encrypted_ref_write_read() {
let temp = TempDir::new().unwrap();
let temp_dir = temp.path().to_path_buf();
fs::create_dir_all(get_lit_dir(&temp_dir).join("refs/heads")).unwrap();
let config = EncryptionConfig {
enabled: true,
key_file: temp_dir
.join("encryption.key")
.to_string_lossy()
.to_string(),
..Default::default()
};
let mut enc_manager = EncryptionManager::new(config);
enc_manager.initialize("test-passphrase-refs").unwrap();
let encryption = Arc::new(Mutex::new(enc_manager));
let test_hash = "abc123def456";
let ref_name = "heads/test-branch";
write_ref_encrypted(&temp_dir, ref_name, test_hash, &encryption).unwrap();
let read_hash = read_ref_encrypted(&temp_dir, ref_name, &encryption).unwrap();
assert_eq!(read_hash, test_hash);
}
#[test]
fn test_encrypted_head_operations() {
let temp = TempDir::new().unwrap();
let temp_dir = temp.path().to_path_buf();
fs::create_dir_all(get_lit_dir(&temp_dir).join("refs/heads")).unwrap();
let config = EncryptionConfig {
enabled: true,
key_file: temp_dir
.join("encryption.key")
.to_string_lossy()
.to_string(),
..Default::default()
};
let mut enc_manager = EncryptionManager::new(config);
enc_manager.initialize("test-passphrase-head").unwrap();
let encryption = Arc::new(Mutex::new(enc_manager));
let branch_name = "main";
let commit_hash = "deadbeef123456";
write_ref_encrypted(&temp_dir, "heads/main", commit_hash, &encryption).unwrap();
update_head_encrypted(&temp_dir, branch_name, &encryption).unwrap();
let current = get_current_branch_encrypted(&temp_dir, &encryption).unwrap();
assert_eq!(current, branch_name);
let head_commit = read_head_encrypted(&temp_dir, &encryption).unwrap();
assert_eq!(head_commit, commit_hash);
}
#[test]
fn test_encrypted_detached_head() {
let temp = TempDir::new().unwrap();
let temp_dir = temp.path().to_path_buf();
fs::create_dir_all(get_lit_dir(&temp_dir)).unwrap();
let config = EncryptionConfig {
enabled: true,
key_file: temp_dir
.join("encryption.key")
.to_string_lossy()
.to_string(),
..Default::default()
};
let mut enc_manager = EncryptionManager::new(config);
enc_manager.initialize("test-passphrase-detached").unwrap();
let encryption = Arc::new(Mutex::new(enc_manager));
let commit_hash = "cafebabe987654";
set_head_detached_encrypted(&temp_dir, commit_hash, &encryption).unwrap();
let head = read_head_encrypted(&temp_dir, &encryption).unwrap();
assert_eq!(head, commit_hash);
assert!(get_current_branch_encrypted(&temp_dir, &encryption).is_err());
}
#[test]
fn test_encrypted_ref_tamper_detection() {
let temp = TempDir::new().unwrap();
let temp_dir = temp.path().to_path_buf();
fs::create_dir_all(get_lit_dir(&temp_dir).join("refs/heads")).unwrap();
let config = EncryptionConfig {
enabled: true,
key_file: temp_dir
.join("encryption.key")
.to_string_lossy()
.to_string(),
..Default::default()
};
let mut enc_manager = EncryptionManager::new(config);
enc_manager.initialize("test-passphrase-tamper").unwrap();
let encryption = Arc::new(Mutex::new(enc_manager));
let test_hash = "original123";
let ref_name = "heads/tamper-test";
write_ref_encrypted(&temp_dir, ref_name, test_hash, &encryption).unwrap();
let ref_path = get_lit_dir(&temp_dir).join("refs").join(ref_name);
let mut data = fs::read(&ref_path).unwrap();
let len = data.len();
data[len - 1] ^= 0x01; fs::write(&ref_path, data).unwrap();
assert!(read_ref_encrypted(&temp_dir, ref_name, &encryption).is_err());
}
}