From f4a91d4f6071e815d3067be5d5c33a4d36016908 Mon Sep 17 00:00:00 2001
From: b1n <b1n@b1n.io>
Date: Sat, 22 Jun 2024 15:38:49 +0800
Subject: [PATCH] feat(bincode): add 500MiB limit to plugin

---
 src/bin/maker.rs |  8 +++-----
 src/plugin.rs    | 36 +++++++++++++++++++++++-------------
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/src/bin/maker.rs b/src/bin/maker.rs
index 2629bf6..b9736f3 100644
--- a/src/bin/maker.rs
+++ b/src/bin/maker.rs
@@ -1,6 +1,5 @@
 use std::{fs, path::PathBuf, usize};
 
-use bincode::encode_to_vec;
 use dirs::{desktop_dir, home_dir};
 use iced::{
     advanced::Application,
@@ -419,9 +418,6 @@ impl Application for Maker {
 
                 let plugin_name = self.plugin_name().to_owned();
                 let make_plugin = async move {
-                    let buf = encode_to_vec(plugin, bincode::config::standard())
-                        .map_err(|_| "Encode plugin failed.".to_string())?;
-
                     let file = AsyncFileDialog::new()
                         .set_directory(desktop_dir().unwrap_or(".".into()))
                         .set_file_name(format!("{}.b1n", plugin_name))
@@ -431,7 +427,9 @@ impl Application for Maker {
                         .await
                         .ok_or("Canceled plugin saving.".to_string())?;
 
-                    fs::write(file.path(), buf).map_err(|_| "Write plugin failed.".to_string())?;
+                    plugin
+                        .write_plugin(file.path())
+                        .map_err(|_| "Write plugin failed.".to_string())?;
 
                     Ok(())
                 };
diff --git a/src/plugin.rs b/src/plugin.rs
index 000d71e..4f53213 100644
--- a/src/plugin.rs
+++ b/src/plugin.rs
@@ -2,9 +2,18 @@ use std::{collections::HashMap, fmt::Display, fs, ops::Not, path::Path};
 
 use aes_gcm::{aead::Aead, Aes256Gcm, KeyInit, Nonce};
 use anyhow::anyhow;
-use bincode::{decode_from_reader, decode_from_slice, encode_to_vec, Decode, Encode};
+use bincode::{decode_from_slice, encode_to_vec, Decode, Encode};
 use dirs::data_dir;
 
+// 500 MiB
+const LIMIT: usize = 1024 * 1024 * 500;
+pub const BINCODE_PLUGIN_CONFIG: bincode::config::Configuration<
+    bincode::config::LittleEndian,
+    bincode::config::Varint,
+    bincode::config::Limit<LIMIT>,
+> = bincode::config::standard().with_limit();
+const BINCODE_PLUGINS_CONFIG: bincode::config::Configuration = bincode::config::standard();
+
 #[derive(Debug, Clone, Encode, Decode, PartialEq, Eq)]
 pub struct AesGcmPass {
     key_holder: Vec<u8>,
@@ -168,12 +177,16 @@ impl Plugin {
 
 impl Plugin {
     pub fn reade_plugin(path: &Path) -> anyhow::Result<Plugin> {
-        let config = bincode::config::standard();
+        let buf = fs::read(path)?;
+        let (plugin, _) = decode_from_slice(buf.as_slice(), BINCODE_PLUGIN_CONFIG)?;
+        Ok(plugin)
+    }
 
-        let file = fs::File::open(path)?;
-        let reader = std::io::BufReader::new(file);
-        let decoded_plugin = decode_from_reader(reader, config)?;
-        Ok(decoded_plugin)
+    pub fn write_plugin(&self, path: &Path) -> anyhow::Result<()> {
+        let buf = encode_to_vec(self, BINCODE_PLUGIN_CONFIG)?;
+        fs::write(path, buf.as_slice())?;
+
+        Ok(())
     }
 }
 
@@ -182,24 +195,21 @@ pub struct Plugins(pub HashMap<String, Plugin>);
 
 impl Plugins {
     pub fn reade_plugins() -> anyhow::Result<Plugins> {
-        let config = bincode::config::standard();
-
         let mut plugins_path = data_dir().ok_or(anyhow::anyhow!("data_dir is none."))?;
         plugins_path.push("pumpbin");
         plugins_path.push("plugins");
 
         if plugins_path.exists() && plugins_path.is_file() {
-            let plugins = fs::read(plugins_path)?;
-            let (decoded_plugins, _) = decode_from_slice(plugins.as_slice(), config)?;
-            Ok(decoded_plugins)
+            let buf = fs::read(plugins_path)?;
+            let (plugins, _) = decode_from_slice(buf.as_slice(), BINCODE_PLUGINS_CONFIG)?;
+            Ok(plugins)
         } else {
             anyhow::bail!("file not exists.")
         }
     }
 
     pub fn uptade_plugins(&self) -> anyhow::Result<()> {
-        let config = bincode::config::standard();
-        let buf = encode_to_vec(&self.0, config)?;
+        let buf = encode_to_vec(self, BINCODE_PLUGINS_CONFIG)?;
 
         let mut plugins_path = data_dir().ok_or(anyhow::anyhow!("data_dir is none."))?;
         plugins_path.push("pumpbin");