use std::path::PathBuf;
use std::sync::Mutex;
use serde::{Deserialize, Serialize};
use crate::error::{ApiError, ApiResult};
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct DataSource {
pub id: String,
pub name: String,
#[serde(rename = "type")]
pub ds_type: String,
pub config: serde_json::Value,
#[serde(default = "default_encryption_enabled")]
pub encryption_enabled: bool,
#[serde(default)]
pub password: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub prev_password: Option<String>,
#[serde(default = "default_volume_enabled")]
pub volume_enabled: bool,
#[serde(default = "default_volume_size")]
pub volume_size: u64,
#[serde(default = "default_volume_strategy")]
pub volume_strategy: String,
#[serde(default = "default_volume_name_format")]
pub volume_name_format: String,
#[serde(default = "default_cache_enabled")]
pub cache_enabled: bool,
pub created_at: u64,
}
pub const DEFAULT_VOLUME_SIZE: u64 = 300 * 1024 * 1024;
pub const MIN_VOLUME_SIZE: u64 = 64 * 1024;
pub fn gen_password() -> String {
use base64::Engine;
let mut raw = [0u8; 18];
getrandom::fill(&mut raw).expect("系统随机数不可用");
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw)
}
fn default_encryption_enabled() -> bool {
true
}
fn default_volume_enabled() -> bool {
true
}
fn default_volume_size() -> u64 {
DEFAULT_VOLUME_SIZE
}
fn default_volume_strategy() -> String {
"random".into()
}
fn default_volume_name_format() -> String {
"{s}_{i}.bin".into()
}
fn default_cache_enabled() -> bool {
true
}
#[derive(Debug, Default, Serialize, Deserialize)]
struct RegistryFile {
version: u32,
datasources: Vec<DataSource>,
}
pub struct Registry {
path: PathBuf,
inner: Mutex<Vec<DataSource>>,
}
impl Registry {
pub fn load(path: PathBuf) -> anyhow::Result<Self> {
let list = match std::fs::read(&path) {
Ok(bytes) => serde_json::from_slice::<RegistryFile>(&bytes)?.datasources,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Vec::new(),
Err(e) => return Err(e.into()),
};
Ok(Self {
path,
inner: Mutex::new(list),
})
}
pub fn list(&self) -> Vec<DataSource> {
self.inner.lock().unwrap().clone()
}
pub fn get(&self, id: &str) -> Option<DataSource> {
self.inner
.lock()
.unwrap()
.iter()
.find(|d| d.id == id)
.cloned()
}
pub fn create(&self, ds: DataSource) -> ApiResult<DataSource> {
let mut guard = self.inner.lock().unwrap();
guard.push(ds.clone());
self.save(&guard)?;
Ok(ds)
}
pub fn update(&self, id: &str, ds: DataSource) -> ApiResult<DataSource> {
let mut guard = self.inner.lock().unwrap();
let slot = guard
.iter_mut()
.find(|d| d.id == id)
.ok_or_else(|| ApiError::NotFound(format!("数据源不存在: {id}")))?;
let mut replacement = DataSource {
id: id.to_string(),
..ds
};
preserve_live_baidu_tokens(slot, &mut replacement);
*slot = replacement;
let saved = slot.clone();
self.save(&guard)?;
Ok(saved)
}
pub fn update_baidu_tokens(
&self,
id: &str,
access_token: &str,
refresh_token: &str,
access_token_expires_at: u64,
) -> ApiResult<()> {
let mut guard = self.inner.lock().unwrap();
let datasource = guard
.iter_mut()
.find(|datasource| datasource.id == id)
.ok_or_else(|| ApiError::NotFound(format!("数据源不存在: {id}")))?;
if datasource.ds_type != "baidupan" {
return Err(ApiError::BadRequest("数据源不是百度网盘".into()));
}
let config = datasource
.config
.as_object_mut()
.ok_or_else(|| ApiError::BadRequest("百度网盘配置不是对象".into()))?;
config.insert("accessToken".into(), access_token.into());
config.insert("refreshToken".into(), refresh_token.into());
config.insert(
"accessTokenExpiresAt".into(),
access_token_expires_at.into(),
);
self.save(&guard)
}
pub fn remove(&self, id: &str) -> ApiResult<()> {
let mut guard = self.inner.lock().unwrap();
let before = guard.len();
guard.retain(|d| d.id != id);
if guard.len() == before {
return Err(ApiError::NotFound(format!("数据源不存在: {id}")));
}
self.save(&guard)?;
Ok(())
}
fn save(&self, list: &[DataSource]) -> ApiResult<()> {
let file = RegistryFile {
version: 1,
datasources: list.to_vec(),
};
let data = serde_json::to_vec_pretty(&file).map_err(|e| anyhow::anyhow!(e))?;
let tmp = self.path.with_extension("json.tmp");
std::fs::write(&tmp, &data)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = std::fs::set_permissions(&tmp, std::fs::Permissions::from_mode(0o600));
}
std::fs::rename(&tmp, &self.path)?;
Ok(())
}
}
fn preserve_live_baidu_tokens(current: &DataSource, replacement: &mut DataSource) {
if current.ds_type != "baidupan" || replacement.ds_type != "baidupan" {
return;
}
let identity_fields = ["bduss", "clientId", "clientSecret"];
let same_identity = identity_fields
.iter()
.all(|field| current.config.get(field) == replacement.config.get(field));
if !same_identity {
return;
}
let Some(target) = replacement.config.as_object_mut() else {
return;
};
for field in ["accessToken", "refreshToken", "accessTokenExpiresAt"] {
if let Some(value) = current.config.get(field) {
target.insert(field.into(), value.clone());
}
}
}
pub fn now_ms() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use super::*;
fn ds(id: &str) -> DataSource {
DataSource {
id: id.into(),
name: format!("ds-{id}"),
ds_type: "localfs".into(),
config: serde_json::json!({"root": "/tmp/x"}),
encryption_enabled: true,
password: "test-password".into(),
prev_password: None,
volume_enabled: true,
volume_size: DEFAULT_VOLUME_SIZE,
volume_strategy: "random".into(),
volume_name_format: "{s}_{i}.bin".into(),
cache_enabled: true,
created_at: 1,
}
}
#[test]
fn crud_roundtrip_persists() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("datasources.json");
let reg = Registry::load(path.clone()).unwrap();
reg.create(ds("a")).unwrap();
reg.create(ds("b")).unwrap();
reg.update(
"a",
DataSource {
name: "renamed".into(),
..ds("a")
},
)
.unwrap();
reg.remove("b").unwrap();
assert!(reg.remove("b").is_err());
let reg2 = Registry::load(path).unwrap();
let list = reg2.list();
assert_eq!(list.len(), 1);
assert_eq!(list[0].name, "renamed");
assert_eq!(reg2.get("a").unwrap().id, "a");
}
#[test]
fn refreshed_baidu_tokens_are_persisted() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("datasources.json");
let registry = Registry::load(path.clone()).unwrap();
let mut source = ds("baidu");
source.ds_type = "baidupan".into();
source.config = serde_json::json!({
"accessToken": "old-access",
"refreshToken": "old-refresh"
});
registry.create(source).unwrap();
registry
.update_baidu_tokens("baidu", "new-access", "new-refresh", 1234)
.unwrap();
let reloaded = Registry::load(path).unwrap();
let config = reloaded.get("baidu").unwrap().config;
assert_eq!(config["accessToken"], "new-access");
assert_eq!(config["refreshToken"], "new-refresh");
assert_eq!(config["accessTokenExpiresAt"], 1234);
}
#[test]
fn datasource_update_does_not_overwrite_fresh_baidu_tokens() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("datasources.json");
let registry = Registry::load(path.clone()).unwrap();
let mut source = ds("baidu");
source.ds_type = "baidupan".into();
source.config = serde_json::json!({
"bduss": "same-account",
"clientId": "same-client",
"clientSecret": "same-secret",
"accessToken": "old-access",
"refreshToken": "old-refresh",
"accessTokenExpiresAt": 100
});
registry.create(source.clone()).unwrap();
registry
.update_baidu_tokens("baidu", "new-access", "new-refresh", 1234)
.unwrap();
source.name = "renamed".into();
registry.update("baidu", source).unwrap();
let reloaded = Registry::load(path).unwrap();
let updated = reloaded.get("baidu").unwrap();
assert_eq!(updated.name, "renamed");
assert_eq!(updated.config["accessToken"], "new-access");
assert_eq!(updated.config["refreshToken"], "new-refresh");
assert_eq!(updated.config["accessTokenExpiresAt"], 1234);
}
}