perf: add back zero-copy decode for mai 1.50

.gitignore

@@ -3,4 +3,4 @@
 /*.txt
 
 /players.redb
-/players.json
+/players.json*

Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = ["sdgb-api", "sdgb-cli"]
 resolver = "3"
+default-members = ["sdgb-cli"]
 
 [workspace.dependencies]
 sdgb-api = { path = "./sdgb-api", default-features = false }
@@ -15,6 +16,7 @@ serde_json = "1.0.141"
 strum = { version = "0.27.2", features = ["derive"] }
 tokio = { version = "1", features = ["rt-multi-thread"] }
 compio = { version = "0.15.0", features = ["runtime"] }
+redb = "2.6.1"
 
 [profile.release]
 lto = true

sdgb-api/src/error.rs

@@ -1,6 +1,8 @@
 use aes::cipher::{block_padding::UnpadError, inout::PadError};
 use snafu::Snafu;
 
+use crate::title::model::LoginError;
+
 #[derive(Debug, Snafu)]
 pub enum ApiError {
     JoinError,
@@ -33,6 +35,12 @@ pub enum ApiError {
     Request {
         source: nyquest::Error,
     },
+
+    #[snafu(display("login error: {source}"))]
+    #[snafu(context(false))]
+    Login {
+        source: LoginError,
+    },
 }
 
 impl From<UnpadError> for ApiError {

sdgb-api/src/lib.rs

@@ -6,3 +6,6 @@ mod error;
 pub use error::ApiError;
 
 pub use bincode;
+
+#[cfg(all(feature = "compio", feature = "tokio"))]
+compile_error!("you must not enable both `compio` and `tokio`");

sdgb-api/src/title/encryption/mod.rs

@@ -12,8 +12,8 @@ use crate::error::ApiError;
 use crate::title::{MaiVersion, MaiVersionExt, Sdgb1_40, Sdgb1_50};
 
 impl MaiVersionExt for Sdgb1_40 {
-    fn decode(data: impl AsRef<[u8]>) -> Result<Vec<u8>, ApiError> {
-        let mut decompressed = decompress(data.as_ref());
+    fn decode(mut data: impl AsMut<[u8]>) -> Result<Vec<u8>, ApiError> {
+        let mut decompressed = decompress(data.as_mut());
         if decompressed.is_empty() {
             return Err(ApiError::EmptyResponse);
         }
@@ -41,19 +41,13 @@ impl MaiVersionExt for Sdgb1_40 {
 }
 
 impl MaiVersionExt for Sdgb1_50 {
-    fn decode(data: impl AsRef<[u8]>) -> Result<Vec<u8>, ApiError> {
-        let mut data = data.as_ref().to_vec();
-        if data.is_empty() {
+    fn decode(mut data: impl AsMut<[u8]>) -> Result<Vec<u8>, ApiError> {
+        if data.as_mut().is_empty() {
             return Err(ApiError::EmptyResponse);
         }
 
-        if data.len() % 16 != 0 {
-            let pad = 16 - (data.len() % 16);
-            data.resize(data.len() + pad, pad as _);
-        }
-
-        debug!("data size: {}", data.len());
-        let decrypted = decrypt_vec(&data, Self::AES_KEY, Self::AES_IV)?;
+        debug!("data size: {}", data.as_mut().len());
+        let decrypted = decrypt(&mut data, Self::AES_KEY, Self::AES_IV)?;
         Ok(decompress(decrypted))
     }
 
@@ -111,13 +105,6 @@ fn decrypt<'ct>(
     Ok(result)
 }
 
-fn decrypt_vec(data: impl AsRef<[u8]>, key: &[u8; 32], iv: &[u8; 16]) -> Result<Vec<u8>, ApiError> {
-    let key = GenericArray::from_slice(key);
-    let iv = GenericArray::from_slice(iv);
-    let decryptor = Aes256CbcDec::new(key, iv);
-    Ok(decryptor.decrypt_padded_vec_mut::<Pkcs7>(data.as_ref())?)
-}
-
 #[cfg(test)]
 mod _tests {
 

sdgb-api/src/title/mod.rs

@@ -25,7 +25,7 @@ pub trait MaiVersion {
 
 pub trait MaiVersionExt: MaiVersion {
     fn encode(data: impl AsRef<[u8]>) -> Result<Vec<u8>, ApiError>;
-    fn decode(data: impl AsRef<[u8]>) -> Result<Vec<u8>, ApiError>;
+    fn decode(data: impl AsMut<[u8]>) -> Result<Vec<u8>, ApiError>;
 
     fn api_hash(api: APIMethod) -> String {
         let api_name: &str = api.into();
@@ -74,6 +74,9 @@ pub trait MaiVersionExt: MaiVersion {
         #[cfg(feature = "tokio")]
         use tokio::task::spawn_blocking;
 
+        #[cfg(all(not(feature = "compio"), not(feature = "tokio")))]
+        compile_error!("you must enable one of `compio` or `tokio`");
+
         async {
             let req = spawn_blocking(move || Self::api_call(api, agent_extra, data))
                 .await

sdgb-api/src/title/model/mod.rs

@@ -8,7 +8,7 @@ mod user_logout_api;
 pub use user_logout_api::{UserLogoutApi, UserLogoutApiResp};
 
 mod user_login_api;
-pub use user_login_api::{UserLoginApi, UserLoginApiResp};
+pub use user_login_api::{LoginError, UserLoginApi, UserLoginApiResp};
 
 mod get_user_data_api;
 pub use get_user_data_api::{GetUserDataApi, GetUserDataApiResp};

sdgb-api/src/title/model/user_login_api/mod.rs

@@ -52,3 +52,28 @@ impl UserLoginApi {
         }
     }
 }
+
+impl UserLoginApiResp {
+    pub fn error(&self) -> Option<LoginError> {
+        match self.return_code {
+            1 => None,
+            100 => Some(LoginError::AlreadyLogged),
+            102 => Some(LoginError::QRCodeExpired),
+            103 => Some(LoginError::AccountUnregistered),
+            error @ _ => Some(LoginError::Unknown { error }),
+        }
+    }
+}
+
+#[derive(Debug, snafu::Snafu)]
+pub enum LoginError {
+    #[snafu(display("QRCode was expired"))]
+    QRCodeExpired,
+    #[snafu(display("You did not logout last session"))]
+    AlreadyLogged,
+    #[snafu(display("userId does not exist"))]
+    AccountUnregistered,
+
+    #[snafu(display("Unknown error: {error}"))]
+    Unknown { error: i32 },
+}

sdgb-cli/Cargo.toml

@@ -6,26 +6,26 @@ authors = ["mokurin000"]
 description = "CLI tool for SDGB protocol"
 
 [features]
-default = ["compio", "cache"]
+default = ["compio", "fetchall"]
 
 compio = ["dep:compio", "sdgb-api/compio"]
 tokio = ["dep:tokio", "sdgb-api/tokio"]
 
-cache = ["dep:redb"]
+fetchall = ["dep:redb", "dep:futures-util"]
 
 [dependencies]
 sdgb-api = { workspace = true, features = ["bincode"] }
-spdlog-rs = { workspace = true }
 
+spdlog-rs = { workspace = true }
 snafu = { workspace = true }
 serde_json = { workspace = true }
 strum = { workspace = true }
+redb = { workspace = true, optional = true }
 tokio = { workspace = true, features = ["macros"], optional = true }
 compio = { workspace = true, features = ["macros"], optional = true }
 
 nyquest-preset = { version = "0.2.0", features = ["async"] }
 
 palc = { version = "0.0.1", features = ["derive"] }
-futures-util = "0.3.31"
-redb = { version = "2.6.1", optional = true }
+futures-util = { version = "0.3.31", optional = true }
 ctrlc = { version = "3.4.7", features = ["termination"] }

sdgb-cli/src/cache/mod.rs

@@ -1,6 +1,6 @@
 use std::sync::LazyLock;
 
-use redb::{Table, TableDefinition, WriteTransaction};
+use redb::{ReadTransaction, Table, TableDefinition, WriteTransaction};
 
 static DATABASE: LazyLock<redb::Database> = LazyLock::new(|| {
     redb::Database::builder()
@@ -17,6 +17,16 @@ pub fn open_table(write: &WriteTransaction) -> Result<Table<'_, u32, Vec<u8>>, r
     Ok(write.open_table(DIFINITION)?)
 }
 
+pub fn read_txn() -> Result<ReadTransaction, redb::Error> {
+    Ok(DATABASE.begin_read()?)
+}
+
+pub fn open_table_read(
+    read: &ReadTransaction,
+) -> Result<redb::ReadOnlyTable<u32, Vec<u8>>, redb::Error> {
+    Ok(read.open_table(DIFINITION)?)
+}
+
 pub fn init_db() -> Result<(), redb::Error> {
     let write_txn = DATABASE.begin_write()?;
     write_txn.open_table(DIFINITION)?;

sdgb-cli/src/commands.rs

@@ -41,7 +41,13 @@ pub enum Commands {
         user_id: u32,
     },
 
-    ListAllUser,
+    #[cfg(feature = "fetchall")]
+    ListAllUser {
+        #[arg(short, long, default_value_t = 5)]
+        concurrency: usize,
+    },
+    #[cfg(feature = "fetchall")]
+    ListAllUserDump {},
 
     Logout {
         #[arg(short, long)]

sdgb-cli/src/main.rs

@@ -1,16 +1,10 @@
-use std::{
-    fs::OpenOptions,
-    io::{self, BufRead},
-    sync::atomic::{AtomicBool, Ordering},
-};
+use std::sync::atomic::{AtomicBool, Ordering};
 
-use futures_util::StreamExt;
 use nyquest_preset::nyquest::ClientBuilder;
 use palc::Parser;
 use spdlog::{Level, LevelFilter::MoreSevereEqual};
 
 use sdgb_api::{
-    ApiError,
     all_net::QRCode,
     auth_lite::{SDGB, SDHJ, delivery_raw},
     title::{
@@ -18,17 +12,18 @@ use sdgb_api::{
         methods::APIMethod,
         model::{
             GetUserDataApi, GetUserDataApiResp, GetUserPreviewApi, GetUserPreviewApiResp, Ping,
-            PingResp, UserLoginApi, UserLoginApiResp, UserLogoutApi, UserLogoutApiResp,
+            PingResp, UserLogoutApi, UserLogoutApiResp,
         },
     },
 };
 use spdlog::{error, info, warn};
 
-use crate::commands::Cli;
+use crate::{commands::Cli, utils::login_action};
 
-#[cfg(feature = "cache")]
+#[cfg(feature = "fetchall")]
 mod cache;
 mod commands;
+mod utils;
 
 static EARLY_QUIT: AtomicBool = AtomicBool::new(false);
 
@@ -44,16 +39,21 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
     }
 
     ctrlc::set_handler(|| {
-        warn!("received early-quit request! will abort soon");
-        EARLY_QUIT.store(true, Ordering::Relaxed);
+        if EARLY_QUIT.load(Ordering::Relaxed) {
+            error!("force-quit triggered!");
+            std::process::exit(1);
+        } else {
+            warn!("received early-quit request! will abort soon");
+            EARLY_QUIT.store(true, Ordering::Relaxed);
+        }
     })?;
 
-    let cmd = <Cli as Parser>::parse();
+    let Cli { command } = <Cli as Parser>::parse();
 
     let client = ClientBuilder::default().build_async().await?;
 
     // TODO: refactor via enum_dispatch
-    match cmd.command {
+    match command {
         commands::Commands::Logout { user_id } => {
             let logout: UserLogoutApiResp = Sdgb1_50::request(
                 &client,
@@ -107,43 +107,42 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
             println!("{}", String::from_utf8_lossy(&resp));
         }
 
-        commands::Commands::ListAllUser => {
-            let mut stdin = io::stdin().lock();
-            let mut buf = String::new();
+        #[cfg(feature = "fetchall")]
+        commands::Commands::ListAllUser { concurrency } => {
+            use futures_util::StreamExt;
+            use sdgb_api::bincode::borrow_decode_from_slice;
+            use std::io::{self, BufRead};
+
             let mut user_ids = Vec::new();
+            {
+                let mut stdin = io::stdin().lock();
+                let mut buf = String::new();
 
-            while stdin.read_line(&mut buf).is_ok_and(|size| size != 0) {
-                if buf.is_empty() {
-                    continue;
+                while stdin.read_line(&mut buf).is_ok_and(|size| size != 0) {
+                    if buf.is_empty() {
+                        continue;
+                    }
+
+                    let user_id: u32 = buf.trim().parse()?;
+                    buf.clear();
+                    user_ids.push(user_id);
                 }
-
-                let user_id: u32 = buf.trim().parse()?;
-                buf.clear();
-                user_ids.push(user_id);
             }
 
-            #[cfg(feature = "cache")]
             let _ = cache::init_db();
-            #[cfg(feature = "cache")]
+            let read = cache::read_txn()?;
             let write = cache::write_txn()?;
-            #[cfg(feature = "cache")]
             let config = sdgb_api::bincode::config::Configuration::<
                 sdgb_api::bincode::config::LittleEndian,
             >::default()
             .with_no_limit();
 
-            let players = futures_util::stream::iter(user_ids)
+            info!("number of user_id: {}", user_ids.len());
+
+            let collect = futures_util::stream::iter(user_ids)
                 .map(async |user_id| {
-                    if EARLY_QUIT.load(Ordering::Relaxed) {
-                        return Err("early skip due to ctrl-c")?;
-                    }
-
-                    #[cfg(feature = "cache")]
                     {
-                        use redb::ReadableTable;
-                        use sdgb_api::bincode::borrow_decode_from_slice;
-
-                        let cache_table = cache::open_table(&write)?;
+                        let cache_table = cache::open_table_read(&read)?;
                         let data = cache_table.get(user_id)?;
                         if let Some(data) = data {
                             let decoded: (GetUserPreviewApiResp, _) =
@@ -153,6 +152,10 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
                         }
                     }
 
+                    if EARLY_QUIT.load(Ordering::Relaxed) {
+                        return Err("early skip due to ctrl-c")?;
+                    }
+
                     let resp = Sdgb1_50::request::<_, GetUserPreviewApiResp>(
                         &client,
                         APIMethod::GetUserPreviewApi,
@@ -163,22 +166,17 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
 
                     match &resp {
                         Ok(resp) => {
-                            info!("preview: {user_id} succeed");
+                            use sdgb_api::bincode::encode_to_vec;
 
-                            #[cfg(feature = "cache")]
+                            info!("found: {user_id}");
+
+                            if let Ok(mut table) = cache::open_table(&write)
+                                && let Ok(encoded) = encode_to_vec(resp, config)
                             {
-                                use sdgb_api::bincode::encode_to_vec;
-
-                                if let Ok(mut table) = cache::open_table(&write)
-                                    && let Ok(encoded) = encode_to_vec(resp, config)
-                                {
-                                    _ = table.insert(resp.user_id, encoded);
-                                }
+                                _ = table.insert(resp.user_id, encoded);
                             }
                         }
-                        Err(ApiError::JSON { .. }) => {
-                            warn!("account unregistered: {user_id}");
-                        }
+                        Err(sdgb_api::ApiError::JSON { .. }) => {}
                         Err(e) => {
                             error!("preview failed: {e}");
                         }
@@ -186,64 +184,49 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
 
                     Result::<_, Box<dyn snafu::Error>>::Ok(resp?)
                 })
-                .buffer_unordered(20)
+                .buffer_unordered(concurrency) // slower to avoid being banned
                 .filter_map(async |r| r.ok())
                 .collect::<Vec<_>>()
                 .await;
+            drop(collect);
 
-            #[cfg(feature = "cache")]
             let _ = write.commit();
-
-            let output = OpenOptions::new()
-                .write(true)
-                .truncate(true)
-                .create(true)
-                .open("players.json")?;
-            serde_json::to_writer_pretty(output, &players)?;
         }
+        #[cfg(feature = "fetchall")]
+        commands::Commands::ListAllUserDump { .. } => {
+            use std::{fs::OpenOptions, io::BufWriter};
+
+            use redb::ReadableTable;
+            use sdgb_api::bincode::{self, borrow_decode_from_slice};
+
+            use crate::cache::{open_table_read, read_txn};
+
+            let txn = read_txn()?;
+            let table = open_table_read(&txn)?;
+
+            let config = bincode::config::Configuration::<bincode::config::LittleEndian>::default()
+                .with_no_limit();
+
+            let user_ids = table
+                .iter()?
+                .flatten()
+                .map(|d| borrow_decode_from_slice(&d.1.value(), config))
+                .flatten()
+                .map(|(value, _)| value)
+                .collect::<Vec<GetUserPreviewApiResp>>();
+
+            let file = OpenOptions::new()
+                .create(true)
+                .truncate(true)
+                .write(true)
+                .open("players.json")?;
+            file.lock()?;
+            let writer = BufWriter::new(file);
+            serde_json::to_writer(writer, &user_ids)?;
+        }
+
         commands::Commands::Userdata { user_id } => {
-            let login = UserLoginApi::new(user_id);
-            let date_time = login.date_time;
-            let Ok(login_resp): Result<UserLoginApiResp, _> =
-                Sdgb1_50::request(&client, APIMethod::UserLoginApi, user_id, login).await
-            else {
-                let logout_resp: UserLogoutApiResp = Sdgb1_50::request(
-                    &client,
-                    APIMethod::UserLogoutApi,
-                    user_id,
-                    UserLogoutApi {
-                        user_id,
-                        date_time,
-                        ..Default::default()
-                    },
-                )
-                .await?;
-
-                info!("logout: {logout_resp:?}");
-                return Ok(());
-            };
-
-            match login_resp.return_code {
-                1 => info!("login succeed"),
-                100 => {
-                    error!("user already logged");
-                    return Ok(());
-                }
-                102 => {
-                    error!("QRCode expired");
-                    return Ok(());
-                }
-                103 => {
-                    error!("Unregistered userId");
-                    return Ok(());
-                }
-                e @ _ => {
-                    error!("unknown login error: {e}");
-                    return Ok(());
-                }
-            }
-
-            match Sdgb1_50::request::<_, GetUserDataApiResp>(
+            let action = async |_| match Sdgb1_50::request::<_, GetUserDataApiResp>(
                 &client,
                 APIMethod::GetUserDataApi,
                 user_id,
@@ -257,23 +240,13 @@ async fn main() -> Result<(), Box<dyn snafu::Error>> {
                 Err(e) => {
                     error!("failed to get userdata: {e}");
                 }
-            }
-
-            let logout_resp: UserLogoutApiResp = Sdgb1_50::request(
-                &client,
-                APIMethod::UserLogoutApi,
-                user_id,
-                UserLogoutApi {
-                    user_id,
-                    date_time,
-                    ..Default::default()
-                },
-            )
-            .await?;
-
-            info!("logout: {logout_resp:?}");
+            };
+            login_action(&client, user_id, action).await?;
         }
     }
 
     Ok(())
 }
+
+#[cfg(all(feature = "compio", feature = "tokio"))]
+compile_error!("you must not enable both `compio` and `tokio`");

sdgb-cli/src/utils/mod.rs

@@ -0,0 +1,46 @@
+use nyquest_preset::nyquest::AsyncClient;
+use sdgb_api::{
+    ApiError,
+    title::{
+        MaiVersionExt as _, Sdgb1_50,
+        methods::APIMethod,
+        model::{UserLoginApi, UserLoginApiResp, UserLogoutApi, UserLogoutApiResp},
+    },
+};
+use spdlog::info;
+
+pub async fn login_action<R>(
+    client: &AsyncClient,
+    user_id: u32,
+    action: impl AsyncFnOnce(UserLoginApiResp) -> R,
+) -> Result<R, ApiError> {
+    let login = UserLoginApi::new(user_id);
+    let date_time = login.date_time;
+
+    info!("login unix timestamp: {date_time}");
+
+    let login_resp: UserLoginApiResp =
+        Sdgb1_50::request(&client, APIMethod::UserLoginApi, user_id, login).await?;
+
+    match login_resp.error() {
+        None => info!("login succeed"),
+        Some(e) => return Err(e)?,
+    }
+
+    let return_data = action(login_resp).await;
+
+    let logout_resp = Sdgb1_50::request::<_, UserLogoutApiResp>(
+        &client,
+        APIMethod::UserLogoutApi,
+        user_id,
+        UserLogoutApi {
+            user_id,
+            date_time,
+            ..Default::default()
+        },
+    )
+    .await;
+
+    info!("logout: {logout_resp:?}");
+    Ok(return_data)
+}