diff --git a/core/rust.templating/src/lang_lua/ctx.rs b/core/rust.templating/src/lang_lua/ctx.rs new file mode 100644 index 00000000..ec0f1da6 --- /dev/null +++ b/core/rust.templating/src/lang_lua/ctx.rs @@ -0,0 +1,47 @@ +use mlua::prelude::*; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; + +#[derive(Serialize, Deserialize)] +pub struct TemplateContext { + pub template_data: Arc, + + #[serde(skip)] + #[serde(default = "std::sync::Mutex::default")] + /// The cached serialized value of the template data + cached_template_data: std::sync::Mutex>, +} + +impl TemplateContext { + pub fn new(template_data: super::state::TemplateData) -> Self { + Self { + template_data: Arc::new(template_data), + cached_template_data: std::sync::Mutex::default(), + } + } +} + +pub type TemplateContextRef = LuaUserDataRef; + +impl LuaUserData for TemplateContext { + fn add_fields>(fields: &mut F) { + fields.add_field_method_get("template_data", |lua, this| { + // Check for cached serialized data + let mut cached_data = this + .cached_template_data + .lock() + .map_err(|e| LuaError::external(e.to_string()))?; + + if let Some(v) = cached_data.as_ref() { + return Ok(v.clone()); + } + + log::debug!("TemplateContext: Serializing data"); + let v = lua.to_value(&this.template_data)?; + + *cached_data = Some(v.clone()); + + Ok(v) + }); + } +} diff --git a/core/rust.templating/src/lang_lua/handler.rs b/core/rust.templating/src/lang_lua/handler.rs index fc072503..4f66a6d7 100644 --- a/core/rust.templating/src/lang_lua/handler.rs +++ b/core/rust.templating/src/lang_lua/handler.rs @@ -1,4 +1,3 @@ -use super::state; use super::{resolve_template_to_bytecode, ArLuaThreadInnerState, LuaVmAction, LuaVmResult}; use mlua::prelude::*; @@ -30,34 +29,27 @@ pub async fn handle_event(action: LuaVmAction, tis_ref: &ArLuaThreadInnerState) } }; - let token = match state::add_template( - &tis_ref.lua, - match template { - crate::Template::Raw(_) => "".to_string(), - crate::Template::Named(ref name) => name.clone(), - }, - template.clone(), - pragma, - ) { - Ok(token) => token, - Err(e) => { - return LuaVmResult::LuaError { - err: LuaError::external(e), - }; - } - }; - let exec_name = match template { crate::Template::Raw(_) => "script".to_string(), crate::Template::Named(ref name) => name.to_string(), }; + // Now, create the template context that should be passed to the template + let template_context = super::ctx::TemplateContext::new(super::state::TemplateData { + path: match template { + crate::Template::Raw(_) => "".to_string(), + crate::Template::Named(ref name) => name.clone(), + }, + template, + pragma, + }); + let v: LuaValue = match tis_ref .lua .load(&template_bytecode) .set_name(&exec_name) .set_mode(mlua::ChunkMode::Binary) // Ensure auto-detection never selects binary mode - .call_async((event, token.clone())) + .call_async((event, template_context)) .await { Ok(f) => f, @@ -73,26 +65,10 @@ pub async fn handle_event(action: LuaVmAction, tis_ref: &ArLuaThreadInnerState) _ => {} } - while let Err(e) = state::remove_template(&tis_ref.lua, &token) { - log::error!( - "Could not remove template: {}. Trying again in 300 milliseconds", - e - ); - tokio::time::sleep(std::time::Duration::from_millis(300)).await; - } - return LuaVmResult::LuaError { err: e }; } }; - while let Err(e) = state::remove_template(&tis_ref.lua, &token) { - log::error!( - "Could not remove template: {}. Trying again in 300 milliseconds", - e - ); - tokio::time::sleep(std::time::Duration::from_millis(300)).await; - } - match tis_ref.lua.from_value::(v) { Ok(v) => { return LuaVmResult::Ok { result_val: v }; diff --git a/core/rust.templating/src/lang_lua/mod.rs b/core/rust.templating/src/lang_lua/mod.rs index 5960d41b..c4244bdb 100644 --- a/core/rust.templating/src/lang_lua/mod.rs +++ b/core/rust.templating/src/lang_lua/mod.rs @@ -1,4 +1,4 @@ -// Work in progress +pub mod ctx; pub mod event; pub mod primitives_docs; pub mod samples; @@ -199,7 +199,6 @@ async fn create_lua_vm( serenity_context, reqwest_client, kv_constraints: state::LuaKVConstraints::default(), - per_template: scc::HashMap::new(), kv_ratelimits: Arc::new( state::LuaRatelimits::new_kv_rl().map_err(|e| LuaError::external(e.to_string()))?, ), diff --git a/core/rust.templating/src/lang_lua/plugins/discord.rs b/core/rust.templating/src/lang_lua/plugins/discord.rs index f88a092a..f820c36f 100644 --- a/core/rust.templating/src/lang_lua/plugins/discord.rs +++ b/core/rust.templating/src/lang_lua/plugins/discord.rs @@ -629,7 +629,7 @@ pub fn plugin_docs() -> templating_docgen::Plugin { } ) .method_mut("new", |mut m| { - m.parameter("token", |p| p.typ("string").description("The token of the template to use.")) + m.parameter("token", |p| p.typ("TemplateContext").description("The token of the template to use.")) .return_("executor", |r| r.typ("DiscordExecutor").description("A discord executor.")) }) } @@ -1976,18 +1976,13 @@ pub fn init_plugin(lua: &Lua) -> LuaResult { module.set( "new", - lua.create_function(|lua, (token,): (String,)| { + lua.create_function(|lua, (token,): (crate::TemplateContextRef,)| { let Some(data) = lua.app_data_ref::() else { return Err(LuaError::external("No app data found")); }; - let template_data = data - .per_template - .get(&token) - .ok_or_else(|| LuaError::external("Template not found"))?; - let executor = DiscordActionExecutor { - template_data: template_data.clone(), + template_data: token.template_data.clone(), guild_id: data.guild_id, cache_http: botox::cache::CacheHttpImpl::from_ctx(&data.serenity_context), serenity_context: data.serenity_context.clone(), diff --git a/core/rust.templating/src/lang_lua/plugins/interop.rs b/core/rust.templating/src/lang_lua/plugins/interop.rs index 9e6b434e..2e918d3c 100644 --- a/core/rust.templating/src/lang_lua/plugins/interop.rs +++ b/core/rust.templating/src/lang_lua/plugins/interop.rs @@ -43,47 +43,6 @@ pub fn plugin_docs() -> templating_docgen::Plugin { r.typ("string").description("The current guild ID.") }) }) - .type_mut( - "TemplatePragma", - "`TemplatePragma` contains the pragma of the template. Note that the list of fields below in non-exhaustive as templates can define extra fields on the pragma as well", - |t| { - t - .example(std::sync::Arc::new(crate::TemplatePragma::default())) - .field("lang", |f| { - f.typ("string").description("The language of the template.") - }) - .field("allowed_caps", |f| { - f.typ("{string}").description("The allowed capabilities provided to the template.") - }) - }, - ) - .type_mut( - "TemplateData", - "`TemplateData` is a struct that represents the data associated with a template token. It is used to store the path and pragma of a template token.", - |t| { - t - .example(std::sync::Arc::new(crate::lang_lua::state::TemplateData { - path: "test".to_string(), - pragma: crate::TemplatePragma::default(), - template: crate::Template::Named("foo".to_string()), - })) - .field("path", |f| { - f.typ("string").description("The path of the template token.") - }) - .field("pragma", |f| { - f.typ("TemplatePragma").description("The pragma of the template.") - }) - }, - ) - .method_mut("gettemplatedata", |m| { - m.description("Returns the data associated with a template token.") - .parameter("token", |p| { - p.typ("string").description("The token of the template to retrieve data for.") - }) - .return_("data", |r| { - r.typ("TemplateData?").description("The data associated with the template token, or `null` if no data is found.") - }) - }) .method_mut("current_user", |m| { m.description("Returns the current user of the Lua VM.") .return_("user", |r| { @@ -119,23 +78,6 @@ pub fn init_plugin(lua: &Lua) -> LuaResult { })?, )?; - module.set( - "gettemplatedata", - lua.create_function(|lua, token: String| { - let Some(data) = lua.app_data_ref::() else { - return Err(LuaError::external("No app data found")); - }; - - match data.per_template.read(&token, |_, x| x.clone()) { - Some(data) => { - let v = lua.to_value(&data)?; - Ok(v) - } - None => Ok(lua.null()), - } - })?, - )?; - module.set( "current_user", lua.create_function(|lua, _: ()| { diff --git a/core/rust.templating/src/lang_lua/plugins/kv.rs b/core/rust.templating/src/lang_lua/plugins/kv.rs index 82be2793..b27494fe 100644 --- a/core/rust.templating/src/lang_lua/plugins/kv.rs +++ b/core/rust.templating/src/lang_lua/plugins/kv.rs @@ -107,7 +107,7 @@ pub fn plugin_docs() -> templating_docgen::Plugin { }, ) .method_mut("new", |mut m| { - m.parameter("token", |p| p.typ("string").description("The token of the template to use.")) + m.parameter("token", |p| p.typ("TemplateContext").description("The token of the template to use.")) .return_("executor", |r| r.typ("KvExecutor").description("A key-value executor.")) }) } @@ -301,18 +301,13 @@ pub fn init_plugin(lua: &Lua) -> LuaResult { module.set( "new", - lua.create_function(|lua, (token,): (String,)| { + lua.create_function(|lua, (token,): (crate::TemplateContextRef,)| { let Some(data) = lua.app_data_ref::() else { return Err(LuaError::external("No app data found")); }; - let template_data = data - .per_template - .get(&token) - .ok_or_else(|| LuaError::external("Template not found"))?; - let executor = KvExecutor { - template_data: template_data.clone(), + template_data: token.template_data.clone(), guild_id: data.guild_id, pool: data.pool.clone(), ratelimits: data.kv_ratelimits.clone(), diff --git a/core/rust.templating/src/lang_lua/plugins/mod.rs b/core/rust.templating/src/lang_lua/plugins/mod.rs index 08f1bfc2..c5ed9b03 100644 --- a/core/rust.templating/src/lang_lua/plugins/mod.rs +++ b/core/rust.templating/src/lang_lua/plugins/mod.rs @@ -44,7 +44,6 @@ pub struct RequirePluginArgs { #[derive(serde::Serialize, serde::Deserialize, Default)] pub struct RequireTemplateImportArgs { - pub token: Option, pub current_path: Option, pub custom_prefix: Option, } @@ -55,7 +54,7 @@ pub async fn require(lua: Lua, (plugin_name, args): (String, LuaValue)) -> LuaRe || plugin_name.starts_with("../") || plugin_name.starts_with("$shop/") { - let (pool, guild_id, compiler, vm_bytecode_cache, per_template) = { + let (pool, guild_id, compiler, vm_bytecode_cache) = { let Some(data) = lua.app_data_ref::() else { return Err(LuaError::external("No app data found")); }; @@ -65,7 +64,6 @@ pub async fn require(lua: Lua, (plugin_name, args): (String, LuaValue)) -> LuaRe data.guild_id, data.compiler.clone(), data.vm_bytecode_cache.clone(), - data.per_template.clone(), ) }; @@ -75,14 +73,7 @@ pub async fn require(lua: Lua, (plugin_name, args): (String, LuaValue)) -> LuaRe // Get the current path if token is specified let current_path = { - if let Some(token) = args.token { - // Get the current path from the token - let template_data = per_template - .get(&token) - .ok_or_else(|| LuaError::external("Template not found"))?; - - template_data.path.clone() - } else if let Some(current_path) = args.current_path { + if let Some(current_path) = args.current_path { current_path } else { // Root is the current path diff --git a/core/rust.templating/src/lang_lua/plugins/page.rs b/core/rust.templating/src/lang_lua/plugins/page.rs index c29bfd98..3ccb1d1a 100644 --- a/core/rust.templating/src/lang_lua/plugins/page.rs +++ b/core/rust.templating/src/lang_lua/plugins/page.rs @@ -385,7 +385,7 @@ pub fn plugin_docs() -> templating_docgen::Plugin { .description("Create a page dedicated to your template on a server.") .method_mut("create_page", |mut m| { m.parameter("token", |p| { - p.typ("string") + p.typ("TemplateContext") .description("The token of the template to use.") }) .return_("create_page", |r| { @@ -682,21 +682,16 @@ pub fn init_plugin(lua: &Lua) -> LuaResult { module.set( "create_page", - lua.create_function(|lua, (token,): (String,)| { + lua.create_function(|lua, (token,): (crate::TemplateContextRef,)| { let Some(data) = lua.app_data_ref::() else { return Err(LuaError::external("No app data found")); }; - let template_data = data - .per_template - .get(&token) - .ok_or_else(|| LuaError::external("Template not found"))?; - let page = CreatePage { page_id: sqlx::types::Uuid::new_v4().to_string(), guild_id: data.guild_id, - template: template_data.template.clone(), - title: template_data.path.clone(), + template: token.template_data.template.clone(), + title: token.template_data.path.clone(), description: "Missing description".to_string(), settings: vec![], is_created: false, diff --git a/core/rust.templating/src/lang_lua/plugins/stings.rs b/core/rust.templating/src/lang_lua/plugins/stings.rs index 8240d7f2..92223f42 100644 --- a/core/rust.templating/src/lang_lua/plugins/stings.rs +++ b/core/rust.templating/src/lang_lua/plugins/stings.rs @@ -82,18 +82,13 @@ pub fn init_plugin(lua: &Lua) -> LuaResult { module.set( "new", - lua.create_function(|lua, (token,): (String,)| { + lua.create_function(|lua, (token,): (crate::TemplateContextRef,)| { let Some(data) = lua.app_data_ref::() else { return Err(LuaError::external("No app data found")); }; - let template_data = data - .per_template - .get(&token) - .ok_or_else(|| LuaError::external("Template not found"))?; - let executor = StingExecutor { - template_data: template_data.clone(), + template_data: token.template_data.clone(), guild_id: data.guild_id, serenity_context: data.serenity_context.clone(), ratelimits: data.sting_ratelimits.clone(), diff --git a/core/rust.templating/src/lang_lua/primitives_docs.rs b/core/rust.templating/src/lang_lua/primitives_docs.rs index c99e87cc..ddaaa1b0 100644 --- a/core/rust.templating/src/lang_lua/primitives_docs.rs +++ b/core/rust.templating/src/lang_lua/primitives_docs.rs @@ -133,4 +133,48 @@ pub fn document_primitives() -> templating_docgen::PrimitiveListBuilder { .description("The unique identifier ID of the event. Will be guaranteed to be unique at a per-guild level.") }) }) + .type_mut( + "TemplatePragma", + "`TemplatePragma` contains the pragma of the template. Note that the list of fields below in non-exhaustive as templates can define extra fields on the pragma as well", + |t| { + t + .example(std::sync::Arc::new(crate::TemplatePragma::default())) + .field("lang", |f| { + f.typ("string").description("The language of the template.") + }) + .field("allowed_caps", |f| { + f.typ("{string}").description("The allowed capabilities provided to the template.") + }) + }, + ) + .type_mut( + "TemplateData", + "`TemplateData` is a struct that represents the data associated with a template token. It is used to store the path and pragma of a template token.", + |t| { + t + .example(std::sync::Arc::new(crate::lang_lua::state::TemplateData { + path: "test".to_string(), + pragma: crate::TemplatePragma::default(), + template: crate::Template::Named("foo".to_string()), + })) + .field("path", |f| { + f.typ("string").description("The path of the template token.") + }) + .field("pragma", |f| { + f.typ("TemplatePragma").description("The pragma of the template.") + }) + }, + ) + .type_mut( + "TemplateContext", + "`TemplateContext` is a struct that represents the context of a template. Stores data including the templates data, pragma and what capabilities it should have access to. Passing a TemplateContext is often required when using AntiRaid plugins for security purposes.", + |mut t| { + t + .field("template_data", |f| { + f + .typ("TemplateData") + .description("The data associated with the template.") + }) + }, + ) } diff --git a/core/rust.templating/src/lang_lua/state.rs b/core/rust.templating/src/lang_lua/state.rs index ca69e546..1e4f82da 100644 --- a/core/rust.templating/src/lang_lua/state.rs +++ b/core/rust.templating/src/lang_lua/state.rs @@ -170,11 +170,6 @@ pub struct LuaUserData { pub reqwest_client: reqwest::Client, pub kv_constraints: LuaKVConstraints, - /// Stores a list of tokens to template data - /// - /// Used by actions and other things which use pragma - pub per_template: scc::HashMap>, - /// Stores the lua actions ratelimiters pub actions_ratelimits: Arc, @@ -190,41 +185,3 @@ pub struct LuaUserData { /// Stores the luau compiler pub compiler: Arc, } - -pub fn add_template( - lua: &mlua::Lua, - path: String, - template: crate::Template, - pragma: crate::TemplatePragma, -) -> Result { - let token = botox::crypto::gen_random(32); - - let data = TemplateData { - path, - pragma, - template, - }; - - let data = Arc::new(data); - - let app_data = lua - .app_data_ref::() - .ok_or("Failed to get user data")?; - - app_data - .per_template - .insert(token.clone(), data) - .map_err(|_| "Failed to insert template token")?; - - Ok(token) -} - -pub fn remove_template(lua: &mlua::Lua, token: &str) -> Result<(), crate::Error> { - let app_data = lua - .app_data_ref::() - .ok_or("Failed to get user data")?; - - app_data.per_template.remove(token); - - Ok(()) -} diff --git a/core/rust.templating/src/lib.rs b/core/rust.templating/src/lib.rs index 84cec2a9..ddd7157b 100644 --- a/core/rust.templating/src/lib.rs +++ b/core/rust.templating/src/lib.rs @@ -8,6 +8,7 @@ pub use core::templating_core::{ create_shop_template, parse_shop_template, GuildTemplate, Template, TemplateLanguage, TemplatePragma, }; +pub use lang_lua::ctx::TemplateContextRef; pub use lang_lua::event; pub use lang_lua::primitives_docs; pub use lang_lua::samples;