Skip to main content

fire_scope/
output_common.rs

1use crate::error::AppError;
2use ipnet::IpNet;
3use std::{collections::BTreeSet, path::{Path, PathBuf}};
4use tokio::fs::{self, OpenOptions};
5use tokio::io::AsyncWriteExt;
6
7/// 出力用の安全な識別子に正規化する
8/// - 非ASCII英数字はアンダースコアに置換
9/// - 先頭末尾のアンダースコアは削除
10/// - 長すぎる場合は64文字に切り詰め
11/// - 空になった場合は "UNKNOWN"
12pub fn sanitize_identifier(input: &str) -> String {
13    let mut s = String::with_capacity(input.len());
14    for ch in input.chars() {
15        if ch.is_ascii_alphanumeric() {
16            s.push(ch);
17        } else {
18            s.push('_');
19        }
20    }
21
22    // 先頭末尾のアンダースコア除去
23    let s = s.trim_matches('_').to_string();
24    let s = if s.len() > 64 { s[..64].to_string() } else { s };
25    if s.is_empty() { "UNKNOWN".to_string() } else { s }
26}
27
28/// 汎用ヘッダー生成
29pub fn make_header(now_str: &str, country_code: &str, as_number: &str) -> String {
30    format!(
31        "# Generated at: {}\n# Country Code: {}\n# AS Number: {}\n\n",
32        now_str, country_code, as_number
33    )
34}
35
36pub async fn write_list_txt<P: AsRef<Path>>(
37    path: P,
38    ipnets: &BTreeSet<IpNet>,
39    header: &str,
40) -> Result<(), AppError> {
41    let body = ipnets
42        .iter()
43        .map(|net| net.to_string())
44        .collect::<Vec<_>>()
45        .join("\n");
46
47    let content = format!("{}{}\n", header, body);
48
49    // 常に上書き(原子的に安全な書き込み)
50    atomic_write(path.as_ref(), content.as_bytes()).await?;
51
52    Ok(())
53}
54
55pub async fn write_list_nft<P: AsRef<Path>>(
56    path: P,
57    ipnets: &BTreeSet<IpNet>,
58    header: &str,
59) -> Result<(), AppError> {
60    let file_path = path.as_ref();
61    let define_name_raw = file_path
62        .file_stem()
63        .and_then(|os| os.to_str())
64        .unwrap_or("unknown_define");
65    let define_name = sanitize_identifier(define_name_raw);
66
67    let mut content = String::new();
68    content.push_str(header);
69    content.push_str(&format!("define {} = {{\n", define_name));
70
71    if !ipnets.is_empty() {
72        let lines: Vec<String> = ipnets.iter().map(|n| format!("    {}", n)).collect();
73        content.push_str(&lines.join(",\n"));
74        content.push('\n');
75    }
76
77    content.push_str("}\n");
78
79    // 常に上書き(原子的に安全な書き込み)
80    atomic_write(file_path, content.as_bytes()).await?;
81
82    Ok(())
83}
84
85/// 一時ファイルに書いてから原子的に`rename`で置換する安全な書き込み
86async fn atomic_write(path: &Path, content: &[u8]) -> Result<(), AppError> {
87    let dir = path.parent().unwrap_or_else(|| Path::new("."));
88    let mut tmp_path = PathBuf::from(dir);
89    let fname = path
90        .file_name()
91        .and_then(|s| s.to_str())
92        .unwrap_or("output");
93    let suffix: u64 = rand::random();
94    tmp_path.push(format!(".{}.tmp.{}", fname, suffix));
95
96    // 作成(既存不可)→ 書き込み → fsync
97    {
98        let mut file = OpenOptions::new()
99            .create_new(true)
100            .write(true)
101            .open(&tmp_path)
102            .await?;
103        file.write_all(content).await?;
104        // データの同期(失敗はそのままエラー伝播)
105        file.sync_all().await?;
106    }
107
108    // 原子的置換(Unixは既存を置換、Windowsは失敗しうるためフォールバック)
109    match fs::rename(&tmp_path, path).await {
110        Ok(_) => Ok(()),
111        Err(_) => {
112            let _ = fs::remove_file(path).await;
113            fs::rename(&tmp_path, path).await?;
114            Ok(())
115        }
116    }
117}