diff --git a/Cargo.lock b/Cargo.lock index 7c9c56b..59bf307 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -100,6 +100,30 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "bb8" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b10cf871f3ff2ce56432fddc2615ac7acc3aa22ca321f8fea800846fbb32f188" +dependencies = [ + "async-trait", + "futures-util", + "parking_lot", + "tokio", +] + +[[package]] +name = "bb8-postgres" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56ac82c42eb30889b5c4ee4763a24b8c566518171ebea648cd7e3bc532c60680" +dependencies = [ + "async-trait", + "bb8", + "tokio", + "tokio-postgres", +] + [[package]] name = "bitflags" version = "1.3.2" @@ -188,9 +212,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.1.7" +version = "1.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26a5c3fd7bfa1ce3897a3a3501d362b2d87b7f2583ebcb4a949ec25911025cbc" +checksum = "504bdec147f2cc13c8b57ed9401fd8a147cc66b67ad5cb241394244f2c947549" [[package]] name = "cfg-if" @@ -836,9 +860,9 @@ dependencies = [ [[package]] name = "hyper-util" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ab92f4f49ee4fb4f997c784b7a2e0fa70050211e0b6a287f898c3c9785ca956" +checksum = "cde7055719c54e36e95e8719f95883f22072a48ede39db7fc17a4e1d5281e9b9" dependencies = [ "bytes", "futures-channel", @@ -1106,9 +1130,9 @@ dependencies = [ [[package]] name = "object" -version = "0.36.2" +version = "0.36.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f203fa8daa7bb185f760ae12bd8e097f63d17041dcdcaf675ac54cdf863170e" +checksum = "27b64972346851a39438c60b341ebc01bba47464ae329e55cf343eb93964efd9" dependencies = [ "memchr", ] @@ -1569,11 +1593,13 @@ dependencies = [ name = "rustbot" version = "0.1.0" dependencies = [ + "bb8", + "bb8-postgres", "cargo_toml", "once_cell", "poise", + "regex", "sysinfo", - "tempfile", "tokenservice-client", "tokio", "tokio-postgres", @@ -1659,9 +1685,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.7.0" +version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "976295e77ce332211c0d24d92c0e83e50f5c5f046d11082cea19f3df13a3562d" +checksum = "fc0a2ce646f8655401bb81e7927b812614bd5d91dbc968696be50603510fcaf0" [[package]] name = "rustls-webpki" @@ -1768,9 +1794,9 @@ dependencies = [ [[package]] name = "serde" -version = "1.0.204" +version = "1.0.205" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bc76f558e0cbb2a839d37354c575f1dc3fdc6546b5be373ba43d95f231bf7c12" +checksum = "e33aedb1a7135da52b7c21791455563facbbcc43d0f0f66165b42c21b3dfb150" dependencies = [ "serde_derive", ] @@ -1786,9 +1812,9 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.204" +version = "1.0.205" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e0cd7e117be63d3c3678776753929474f3b04a43a080c744d6b0ae2a8c28e222" +checksum = "692d6f5ac90220161d6774db30c662202721e64aed9058d2c394f451261420c1" dependencies = [ "proc-macro2", "quote", @@ -2144,6 +2170,7 @@ dependencies = [ "bytes", "libc", "mio", + "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2", diff --git a/Cargo.toml b/Cargo.toml index 2599abf..17840d4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,16 +4,21 @@ version = "0.1.0" edition = "2021" [dependencies] +bb8 = "0.8.5" +bb8-postgres = "0.8.1" cargo_toml = "0.20.4" once_cell = "1.19.0" poise = "0.6.1" -sysinfo = "0.31.0" -tempfile = "3.10.1" +regex = "1.10.6" +sysinfo = "0.31.2" tokenservice-client = { version = "0.3.3", registry = "gitea" } -tokio = { version = "1.39.2", features = [ "macros", "signal", "rt-multi-thread" ] } +tokio = { version = "1.39.2", features = ["macros", "signal", "rt-multi-thread"] } tokio-postgres = "0.7.11" uptime_lib = "0.3.1" +[features] +production = [] + [[bin]] name = "rustbot" path = "src/main.rs" diff --git a/Dockerfile b/Dockerfile index d9ae3a2..bdaf274 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,23 +1,24 @@ -FROM rust:1.80-alpine3.19@sha256:b3ac1f65cf33390407c9b90558eb41e7a8311c47d836fca5800960f1aa2d11d5 AS chef -ENV RUSTFLAGS -C target-feature=-crt-static +FROM rust:1.80-alpine3.20 AS chef +ENV RUSTFLAGS="-C target-feature=-crt-static" +ARG GIT_HASH +ENV GIT_COMMIT_HASH=${GIT_HASH} RUN apk add --no-cache openssl-dev musl-dev RUN cargo install cargo-chef -WORKDIR /usr/src/rustbot +WORKDIR /builder FROM chef AS planner COPY . . -RUN mkdir -p .cargo && \ - printf '[registries.gitea]\nindex = "sparse+https://git.toast-server.net/api/packages/toast/cargo/"\ntoken = "Bearer %s"\n' "$CARGO_TOKEN" >> .cargo/config.toml RUN cargo chef prepare FROM chef AS builder -COPY --from=planner /usr/src/rustbot/recipe.json recipe.json +COPY --from=planner /builder/recipe.json recipe.json RUN cargo chef cook --release COPY . . -RUN cargo build -r +RUN cargo build --offline -rF production -FROM alpine:3.20@sha256:0a4eaa0eecf5f8c050e5bba433f58c052be7587ee8af3e8b3910ef9ab5fbe9f5 +FROM alpine:3.20 +LABEL org.opencontainers.image.source="https://git.toast-server.net/toast/Rustbot" RUN apk add --no-cache libgcc WORKDIR /rustbot -COPY --from=builder /usr/src/rustbot/target/release/rustbot . +COPY --from=builder /builder/target/release/rustbot . CMD ./rustbot diff --git a/run.sh b/run.sh index 97f9808..7646904 100755 --- a/run.sh +++ b/run.sh @@ -1,3 +1,3 @@ #!/bin/bash -export $(grep -v '^#' .env | xargs) && cargo run +clear && cargo run diff --git a/src/commands/mod.rs b/src/commands.rs similarity index 76% rename from src/commands/mod.rs rename to src/commands.rs index c4806b0..9e3a067 100644 --- a/src/commands/mod.rs +++ b/src/commands.rs @@ -1,4 +1,4 @@ +pub mod midi; pub mod ping; -pub mod eval; -pub mod uptime; pub mod sample; +pub mod uptime; diff --git a/src/commands/eval.rs b/src/commands/eval.rs deleted file mode 100644 index e03eedc..0000000 --- a/src/commands/eval.rs +++ /dev/null @@ -1,59 +0,0 @@ -use crate::Error; - -use poise::serenity_prelude::UserId; -use std::{ - io::Write, - process::Command, - os::unix::fs::PermissionsExt -}; - -const WHITELISTED_USERS: &[UserId] = &[ - UserId::new(190407856527376384) -]; - -/// Evaluate a piece of code -#[poise::command(slash_command)] -pub async fn eval( - ctx: poise::Context<'_, (), Error>, - #[description = "The Rust code to evaluate"] code: String -) -> Result<(), Error> { - if !WHITELISTED_USERS.contains(&ctx.author().id) { - ctx.reply("Whitelisted users can only use this command!").await?; - return Ok(()); - } - - // Create a temp directory - let dir = tempfile::tempdir()?; - let file_path = dir.path().join("temp.rs"); - - { - let mut file = std::fs::File::create(&file_path)?; - writeln!(file, "fn main() {{ {} }}", code)?; - } - - // Compile - let compiled_path = dir.path().join("temp"); - println!("Compiling {} -> {}", file_path.display(), compiled_path.display()); - let output = Command::new("rustc").arg(&file_path).arg("-o").arg(&compiled_path).output()?; - - if !output.status.success() { - ctx.reply(format!("Compilation failed:\n```{}```", String::from_utf8_lossy(&output.stderr))).await?; - return Ok(()); - } - - // Update binary's permissions before execution stage - let permissions = std::fs::Permissions::from_mode(0o755); - let compiled_path = dir.path().join("temp"); - std::fs::set_permissions(&compiled_path, permissions)?; - - // If success, run it. - let output = Command::new(compiled_path).output()?; - - if !output.status.success() { - ctx.reply(format!("Execution failed:\n```{}```", String::from_utf8_lossy(&output.stderr))).await?; - return Ok(()); - } - - ctx.reply(format!("Output:\n```{}```", String::from_utf8_lossy(&output.stdout))).await?; - Ok(()) -} diff --git a/src/commands/midi.rs b/src/commands/midi.rs new file mode 100644 index 0000000..57f7766 --- /dev/null +++ b/src/commands/midi.rs @@ -0,0 +1,101 @@ +use crate::{ + Error, + internals::utils::{ + mention_dev, + format_bytes + } +}; + +use regex::Regex; +use std::{ + os::unix::fs::MetadataExt, + fs::{ + write, + remove_file, + metadata + } +}; +use poise::{ + CreateReply, + serenity_prelude::CreateAttachment +}; + +/// Convert MIDI file to WAV +#[poise::command(context_menu_command = "MIDI -> WAV")] +pub async fn midi_to_wav( + ctx: poise::Context<'_, (), Error>, + #[description = "MIDI file to be converted"] message: poise::serenity_prelude::Message +) -> Result<(), Error> { + let re = Regex::new(r"(?i)\.mid$").unwrap(); + + if !message.embeds.is_empty() || message.attachments.is_empty() || !re.is_match(&message.attachments[0].filename) { + ctx.reply("That ain't a MIDI file! What are you even doing??").await?; + return Ok(()); + } + + ctx.defer().await?; + + let bytes = match message.attachments[0].download().await { + Ok(bytes) => bytes, + Err(y) => { + ctx.send(CreateReply::default() + .content(format!( + "Download failed, ask {} to check console for more information!", + mention_dev(ctx).unwrap_or_default() + )) + ) + .await.unwrap(); + + return Err(Error::from(format!("Failed to download the file: {}", y))) + } + }; + + let midi_path = &message.attachments[0].filename; + write(midi_path, bytes)?; + + let wav_path = re.replace(&midi_path, ".wav"); + + let sf2_path = "/tmp/FluidR3_GM.sf2"; + write(sf2_path, include_bytes!("../internals/assets/FluidR3_GM.sf2"))?; + + let output = std::process::Command::new("fluidsynth") + .args(&["-ni", sf2_path, midi_path, "-F", &wav_path]) + .output(); + + // Just to add an info to console to tell what the bot is doing when MIDI file is downloaded. + println!("Discord[{}:{}]: Processing MIDI file: \"{}\"", ctx.guild().unwrap().name, ctx.command().qualified_name, midi_path); + + match output { + Ok(_) => { + let reply = ctx.send(CreateReply::default() + .attachment(CreateAttachment::path(&*wav_path).await.unwrap()) + ).await; + + if reply.is_err() { + println!( + "Discord[{}:{}]: Processed file couldn't be uploaded back to Discord channel due to upload limit", + ctx.guild().unwrap().name, ctx.command().qualified_name + ); + + ctx.send(CreateReply::default() + .content(format!( + "Couldn't upload the processed file (`{}`, `{}`) due to upload limit", + &*wav_path, format_bytes(metadata(&*wav_path).unwrap().size()) + )) + ).await.unwrap(); + } else if reply.is_ok() { + remove_file(midi_path)?; + remove_file(&*wav_path)?; + } + }, + Err(y) => { + ctx.send(CreateReply::default() + .content("Command didn't execute successfully, check console for more information!") + ).await.unwrap(); + + return Err(Error::from(format!("Midi conversion failed: {}", y))) + } + } + + Ok(()) +} diff --git a/src/commands/uptime.rs b/src/commands/uptime.rs index c7928eb..1a17136 100644 --- a/src/commands/uptime.rs +++ b/src/commands/uptime.rs @@ -1,20 +1,52 @@ use crate::{ Error, + GIT_COMMIT_HASH, internals::utils::{ + BOT_VERSION, format_duration, - concat_message, - BOT_VERSION + concat_message } }; use sysinfo::System; use uptime_lib::get; -use std::time::{ - Duration, - SystemTime, - UNIX_EPOCH +use std::{ + fs::File, + path::Path, + time::{ + Duration, + SystemTime, + UNIX_EPOCH + }, + io::{ + BufRead, + BufReader + } }; +fn get_os_info() -> String { + let path = Path::new("/etc/os-release"); + let mut name = "BoringOS".to_string(); + let mut version = "v0.0".to_string(); + + if let Ok(file) = File::open(&path) { + let reader = BufReader::new(file); + for line in reader.lines() { + if let Ok(line) = line { + if line.starts_with("NAME=") { + name = line.split('=').nth(1).unwrap_or_default().trim_matches('"').to_string(); + } else if line.starts_with("VERSION=") { + version = line.split('=').nth(1).unwrap_or_default().trim_matches('"').to_string(); + } else if line.starts_with("VERSION_ID=") { + version = line.split('=').nth(1).unwrap_or_default().trim_matches('"').to_string(); + } + } + } + } + + format!("{} {}", name, version) +} + /// Retrieve host and bot uptimes #[poise::command(slash_command)] pub async fn uptime(ctx: poise::Context<'_, (), Error>) -> Result<(), Error> { @@ -25,6 +57,9 @@ pub async fn uptime(ctx: poise::Context<'_, (), Error>) -> Result<(), Error> { // Fetch system's uptime let sys_uptime = get().unwrap().as_secs(); + // Fetch system's processor + let cpu = sys.cpus(); + // Fetch bot's process uptime let curr_pid = sysinfo::get_current_pid().unwrap(); let now = SystemTime::now(); @@ -35,9 +70,11 @@ pub async fn uptime(ctx: poise::Context<'_, (), Error>) -> Result<(), Error> { } let stat_msg = vec![ - format!("**{} {}**", _bot.name, BOT_VERSION.as_str()), + format!("**{} {}** `{}`", _bot.name, BOT_VERSION.as_str(), GIT_COMMIT_HASH), format!(">>> System: `{}`", format_duration(sys_uptime)), - format!("Process: `{}`", format_duration(proc_uptime)) + format!("Process: `{}`", format_duration(proc_uptime)), + format!("CPU: `{}`", format!("{}", cpu[0].brand())), + format!("OS: `{}`", get_os_info()) ]; ctx.reply(concat_message(stat_msg)).await?; diff --git a/src/controllers/mod.rs b/src/controllers.rs similarity index 100% rename from src/controllers/mod.rs rename to src/controllers.rs diff --git a/src/controllers/database.rs b/src/controllers/database.rs index 34da831..cfb29c4 100644 --- a/src/controllers/database.rs +++ b/src/controllers/database.rs @@ -1,36 +1,84 @@ -use crate::internals; +use crate::internals::utils::token_path; -use poise::serenity_prelude::prelude::TypeMapKey; -use tokio_postgres::{Client, NoTls, Error}; +use bb8::{Pool, PooledConnection}; +use bb8_postgres::PostgresConnectionManager; +use tokio::time::{ + Duration, + sleep +}; +use tokio_postgres::{ + Client, + NoTls, + Error, + config::Config +}; +use std::{ + ops::Deref, + str::FromStr, + sync::{ + Mutex, + LazyLock + } +}; + +pub static DATABASE: LazyLock>> = LazyLock::new(|| Mutex::new(None)); pub struct DatabaseController { - pub client: Client -} - -impl TypeMapKey for DatabaseController { - type Value = DatabaseController; + pub pool: Pool> } impl DatabaseController { - pub async fn new() -> Result { - let (client, connection) = tokio_postgres::connect(&internals::utils::token_path().await.postgres_uri, NoTls).await?; + pub async fn new() -> Result<(), Error> { + let manager = PostgresConnectionManager::new(Config::from_str(token_path().await.postgres_uri.as_str())?, NoTls); + let pool = bb8::Pool::builder().build(manager).await?; + let err_name = "Postgres[Error]"; + let pool_clone = pool.clone(); tokio::spawn(async move { - if let Err(e) = connection.await { - eprintln!("Connection error: {}", e); + loop { + match Self::attempt_connect(&pool_clone).await { + Ok(conn) => { + println!("Postgres[Info]: Successfully connected"); + let client: &Client = conn.deref(); + + // Sample model + client.batch_execute(" + CREATE TABLE IF NOT EXISTS sample ( + id BIGSERIAL PRIMARY KEY, + text_val VARCHAR(255) NOT NULL, + int_val BIGINT NOT NULL, + boolean_val BOOLEAN NOT NULL + ); + ").await.unwrap(); + }, + Err(e) => { + eprintln!("{}: {}", err_name, e); + sleep(Duration::from_secs(5)).await; + } + } + break; } }); - // Sample model - client.batch_execute(" - CREATE TABLE IF NOT EXISTS sample ( - id BIGSERIAL PRIMARY KEY, - text_val VARCHAR(255) NOT NULL, - int_val BIGINT NOT NULL, - boolean_val BOOLEAN NOT NULL - ); - ").await?; + let controller = Self { pool }; + *DATABASE.lock().unwrap() = Some(controller); - Ok(DatabaseController { client }) + Ok(()) + } + + async fn attempt_connect<'a>(pool: &'a bb8::Pool>) -> Result>, bb8::RunError> { + let mut backoff = 1; + loop { + match pool.get().await { + Ok(conn) => return Ok(conn), + Err(e) => { + eprintln!("Postgres[ConnError]: {}, retrying in {} seconds", e, backoff); + sleep(Duration::from_secs(backoff)).await; + if backoff < 64 { + backoff *= 2; + } + } + } + } } } diff --git a/src/internals/mod.rs b/src/internals.rs similarity index 51% rename from src/internals/mod.rs rename to src/internals.rs index 3c59201..0ce8a43 100644 --- a/src/internals/mod.rs +++ b/src/internals.rs @@ -1,2 +1,4 @@ -pub mod utils; +pub mod config; +pub mod tasks; pub mod tsclient; +pub mod utils; diff --git a/src/internals/assets/FluidR3_GM.sf2 b/src/internals/assets/FluidR3_GM.sf2 new file mode 100755 index 0000000..443d42b Binary files /dev/null and b/src/internals/assets/FluidR3_GM.sf2 differ diff --git a/src/internals/config.rs b/src/internals/config.rs new file mode 100644 index 0000000..994498c --- /dev/null +++ b/src/internals/config.rs @@ -0,0 +1,62 @@ +use std::sync::LazyLock; + +pub struct ConfigMeta { + pub embed_color: i32, + pub ready_notify: u64, + pub rss_channel: u64, + pub rustbot_logs: u64, + pub deploy_commands: bool, + pub developers: Vec +} + +#[cfg(feature = "production")] +pub static BINARY_PROPERTIES: LazyLock = LazyLock::new(|| ConfigMeta::new()); + +#[cfg(not(feature = "production"))] +pub static BINARY_PROPERTIES: LazyLock = LazyLock::new(|| + ConfigMeta::new() + .embed_color(0xf1d63c) + .ready_notify(865673694184996888) + .rss_channel(865673694184996888) + .deploy_commands(false) +); + +impl ConfigMeta { + fn new() -> Self { + Self { + embed_color: 0x5a99c7, + ready_notify: 865673694184996888, + rss_channel: 865673694184996888, + rustbot_logs: 1268493237912604672, + deploy_commands: false, + developers: vec![ + 190407856527376384 // toast.ts + ] + } + } + + // Scalable functions below; + #[cfg(not(feature = "production"))] + fn embed_color(mut self, color: i32) -> Self { + self.embed_color = color; + self + } + + #[cfg(not(feature = "production"))] + fn ready_notify(mut self, channel_id: u64) -> Self { + self.ready_notify = channel_id; + self + } + + #[cfg(not(feature = "production"))] + fn rss_channel(mut self, channel_id: u64) -> Self { + self.rss_channel = channel_id; + self + } + + #[cfg(not(feature = "production"))] + fn deploy_commands(mut self, deploy: bool) -> Self { + self.deploy_commands = deploy; + self + } +} diff --git a/src/internals/tasks.rs b/src/internals/tasks.rs new file mode 100644 index 0000000..5063dc8 --- /dev/null +++ b/src/internals/tasks.rs @@ -0,0 +1,9 @@ +pub mod sample; + +fn task_info(name: &str, message: &str) { + println!("{}", format!("TaskScheduler[{}]: {}", name, message)); +} + +fn task_err(name: &str, message: &str) { + eprintln!("{}", format!("TaskScheduler[{}:Error]: {}", name, message)); +} diff --git a/src/internals/tasks/sample.rs b/src/internals/tasks/sample.rs new file mode 100644 index 0000000..2ca322c --- /dev/null +++ b/src/internals/tasks/sample.rs @@ -0,0 +1,43 @@ +use crate::Error; +use super::{ + super::config::BINARY_PROPERTIES, + task_info, + task_err +}; + +use std::sync::Arc; +use poise::serenity_prelude::{ + Context, + ChannelId, + CreateMessage +}; +use tokio::time::{ + Duration, + interval +}; + +pub async fn sample(ctx: Arc) -> Result<(), Error> { + let task_name = "SampleTask"; + let mut interval = interval(Duration::from_secs(10)); + task_info(&task_name, "Task loaded!"); + + loop { + interval.tick().await; + task_info(&task_name, "Task running!"); + + if BINARY_PROPERTIES.rss_channel == 0 { + task_err(&task_name, "RSS channel ID is not set!"); + ChannelId::new(BINARY_PROPERTIES.rustbot_logs).send_message( + &ctx.http, + CreateMessage::new().content("RSS channel ID is not set!") + ).await.unwrap(); + + continue; + } + + ChannelId::new(BINARY_PROPERTIES.rss_channel).send_message( + &ctx.http, + CreateMessage::new().content("This is a sample message executed by a task!") + ).await.unwrap(); + } +} diff --git a/src/internals/tsclient.rs b/src/internals/tsclient.rs index 3939193..857e062 100644 --- a/src/internals/tsclient.rs +++ b/src/internals/tsclient.rs @@ -1,19 +1,36 @@ -use tokenservice_client::{TokenService, TokenServiceApi}; +use std::sync::LazyLock; +use tokio::sync::RwLock; +use tokenservice_client::{ + TokenService, + TokenServiceApi +}; -pub struct TSClient { - client: TokenService -} +static TS_GLOBAL_CACHE: LazyLock>> = LazyLock::new(|| RwLock::new(None)); + +pub struct TSClient(TokenService); impl TSClient { pub fn new() -> Self { let args: Vec = std::env::args().collect(); - let service = if args.len() > 1 { args[1].as_str() } else { "pgbot" }; - TSClient { - client: TokenService::new(service) + let service = if args.len() > 1 { &args[1] } else { "pgbot" }; + Self(TokenService::new(service)) + } + + pub async fn get(&self) -> Result { + { + let cache = TS_GLOBAL_CACHE.read().await; + if let Some(ref api) = *cache { + return Ok(api.clone()); + } + } + + match self.0.connect().await { + Ok(api) => { + let mut cache = TS_GLOBAL_CACHE.write().await; + *cache = Some(api.clone()); + Ok(api) + } + Err(e) => Err(e) } } - pub async fn get(&self) -> Result> { - let api = self.client.connect().await.unwrap(); - Ok(api) - } } diff --git a/src/internals/utils.rs b/src/internals/utils.rs index 8e5272c..3f2990b 100644 --- a/src/internals/utils.rs +++ b/src/internals/utils.rs @@ -1,10 +1,11 @@ -use once_cell::sync::Lazy; +use poise::serenity_prelude::UserId; +use std::sync::LazyLock; +use tokio::sync::Mutex; use tokenservice_client::TokenServiceApi; +use super::tsclient::TSClient; -pub static EMBED_COLOR: i32 = 0xf1d63c; - -pub static BOT_VERSION: Lazy = Lazy::new(|| { - let cargo_version = cargo_toml::Manifest::from_path("./Cargo.toml") +pub static BOT_VERSION: LazyLock = LazyLock::new(|| { + let cargo_version = cargo_toml::Manifest::from_str(&include_str!("../../Cargo.toml")) .unwrap() .package .unwrap() @@ -13,15 +14,35 @@ pub static BOT_VERSION: Lazy = Lazy::new(|| { format!("v{}", cargo_version) }); +static TSCLIENT: LazyLock> = LazyLock::new(|| Mutex::new(TSClient::new())); + pub async fn token_path() -> TokenServiceApi { - let client = super::tsclient::TSClient::new().get().await.unwrap(); - client + TSCLIENT.lock().await.get().await.unwrap() } pub fn concat_message(messages: Vec) -> String { messages.join("\n") } +pub fn mention_dev(ctx: poise::Context<'_, (), crate::Error>) -> Option { + let devs = super::config::BINARY_PROPERTIES.developers.clone(); + let app_owners = ctx.framework().options().owners.clone(); + + let mut mentions = Vec::new(); + + for dev in devs { + if app_owners.contains(&UserId::new(dev)) { + mentions.push(format!("<@{}>", dev)); + } + } + + if mentions.is_empty() { + None + } else { + Some(mentions.join(", ")) + } +} + pub fn format_duration(secs: u64) -> String { let days = secs / 86400; let hours = (secs % 86400) / 3600; @@ -43,15 +64,23 @@ pub fn format_duration(secs: u64) -> String { formatted_string } -/* pub fn format_memory(bytes: u64) -> String { - let kb = 1024; - let mb = 1024 * 1024; - let gb = 1024 * 1024 * 1024; +pub fn format_bytes(bytes: u64) -> String { + let units = ["B", "KB", "MB", "GB", "TB", "PB"]; + let mut value = bytes as f64; + let mut unit = units[0]; - match bytes { - b if b >= gb => format!("{:.0} GB", (b as f64 / (1024.0 * 1024.0 * 1024.0)).ceil()), - b if b >= mb => format!("{:.0} MB", (b as f64 / (1024.0 * 1024.0)).ceil()), - b if b >= kb => format!("{:.0} KB", (b as f64 / 1024.0).ceil()), - _ => format!("{:.0} B", bytes), + for &u in &units[1..] { + if value < 1024.0 { + break; + } + + value /= 1024.0; + unit = u; } -} */ + + if unit == "B" { + format!("{}{}", value, unit) + } else { + format!("{:.2}{}", value, unit) + } +} diff --git a/src/main.rs b/src/main.rs index 7bcaca8..1d1c686 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,10 +2,30 @@ mod commands; mod controllers; mod models; mod internals; +// https://cdn.toast-server.net/RustFSHiearchy.png +// Using the new filesystem hierarchy + +use crate::{ + internals::{ + utils::{ + BOT_VERSION, + token_path, + mention_dev + }, + config::BINARY_PROPERTIES + }, + controllers::database::DatabaseController +}; use std::{ - env::var, - error + thread::current, + sync::{ + Arc, + atomic::{ + AtomicBool, + Ordering + } + } }; use poise::serenity_prelude::{ builder::{ @@ -13,45 +33,131 @@ use poise::serenity_prelude::{ CreateEmbed, CreateEmbedAuthor }, - Context, Ready, + Context, + FullEvent, ClientBuilder, ChannelId, Command, GatewayIntents }; -type Error = Box; +type Error = Box; +static TASK_RUNNING: AtomicBool = AtomicBool::new(false); -static BOT_READY_NOTIFY: u64 = 865673694184996888; +#[cfg(feature = "production")] +pub static GIT_COMMIT_HASH: &str = env!("GIT_COMMIT_HASH"); +#[cfg(not(feature = "production"))] +pub static GIT_COMMIT_HASH: &str = "devel"; async fn on_ready( ctx: &Context, ready: &Ready, framework: &poise::Framework<(), Error> ) -> Result<(), Error> { - println!("Connected to API as {}", ready.user.name); + #[cfg(not(feature = "production"))] + { + println!("Event[Ready][Notice]: Detected a non-production environment!"); + let gateway = ctx.http.get_bot_gateway().await?; + let session = gateway.session_start_limit; + println!("Event[Ready][Notice]: Session limit: {}/{}", session.remaining, session.total); + } + + println!("Event[Ready]: Build version: {} ({})", BOT_VERSION.to_string(), GIT_COMMIT_HASH); + println!("Event[Ready]: Connected to API as {}", ready.user.name); let message = CreateMessage::new(); let ready_embed = CreateEmbed::new() - .color(internals::utils::EMBED_COLOR) + .color(BINARY_PROPERTIES.embed_color) .thumbnail(ready.user.avatar_url().unwrap_or_default()) - .author(CreateEmbedAuthor::new(format!("{} is ready!", ready.user.name)).clone()); + .author(CreateEmbedAuthor::new(format!("{} is ready!", ready.user.name))); - ChannelId::new(BOT_READY_NOTIFY).send_message(&ctx.http, message.add_embed(ready_embed)).await?; + ChannelId::new(BINARY_PROPERTIES.ready_notify).send_message(&ctx.http, message.add_embed(ready_embed)).await?; - let register_commands = var("REGISTER_CMDS").unwrap_or_else(|_| String::from("true")).parse::().unwrap_or(true); - - if register_commands { + if BINARY_PROPERTIES.deploy_commands { let builder = poise::builtins::create_application_commands(&framework.options().commands); let commands = Command::set_global_commands(&ctx.http, builder).await; + let mut commands_deployed = std::collections::HashSet::new(); match commands { Ok(cmdmap) => for command in cmdmap.iter() { - println!("Registered command globally: {}", command.name); + commands_deployed.insert(command.name.clone()); }, - Err(why) => println!("Error registering commands: {:?}", why) + Err(y) => eprintln!("Error registering commands: {:?}", y) } + + if commands_deployed.len() > 0 { + println!("Event[Ready]: Deployed the commands globally:\n- {}", commands_deployed.into_iter().collect::>().join("\n- ")); + } + } + + Ok(()) +} + +async fn event_processor( + ctx: &Context, + event: &FullEvent, + framework: poise::FrameworkContext<'_, (), Error> +) -> Result<(), Error> { + match event { + FullEvent::Ratelimit { data } => { + println!("Event[Ratelimit]: {:#?}", data); + } + FullEvent::Message { new_message } => { + if new_message.author.bot || !new_message.guild_id.is_none() { + return Ok(()); + } + + if new_message.content.to_lowercase().starts_with("deploy") && new_message.author.id == BINARY_PROPERTIES.developers[0] { + let builder = poise::builtins::create_application_commands(&framework.options().commands); + let commands = Command::set_global_commands(&ctx.http, builder).await; + let mut commands_deployed = std::collections::HashSet::new(); + + match commands { + Ok(cmdmap) => for command in cmdmap.iter() { + commands_deployed.insert(command.name.clone()); + }, + Err(y) => { + eprintln!("Error registering commands: {:?}", y); + new_message.reply(&ctx.http, "Deployment failed, check console for more details!").await?; + } + } + + if commands_deployed.len() > 0 { + new_message.reply(&ctx.http, format!( + "Deployed the commands globally:\n- {}", + commands_deployed.into_iter().collect::>().join("\n- ") + )).await?; + } + } + } + FullEvent::Ready { .. } => { + let thread_id = format!("{:?}", current().id()); + let thread_num: String = thread_id.chars().filter(|c| c.is_digit(10)).collect(); + println!("Event[Ready]: Task Scheduler operating on thread {}", thread_num); + + let ctx = Arc::new(ctx.clone()); + + if !TASK_RUNNING.load(Ordering::SeqCst) { + TASK_RUNNING.store(true, Ordering::SeqCst); + + tokio::spawn(async move { + match internals::tasks::sample::sample(ctx).await { + Ok(_) => {}, + Err(y) => { + eprintln!("TaskScheduler[Main:Sample:Error]: Task execution failed: {}", y); + if let Some(source) = y.source() { + eprintln!("TaskScheduler[Main:Sample:Error]: Task execution failed caused by: {:#?}", source); + } + } + } + TASK_RUNNING.store(false, Ordering::SeqCst); + }); + } else { + println!("TaskScheduler[Main:Notice]: Another thread is already running, ignoring"); + } + } + _ => {} } Ok(()) @@ -59,47 +165,59 @@ async fn on_ready( #[tokio::main] async fn main() { - let db = controllers::database::DatabaseController::new().await.expect("Failed to connect to database"); + DatabaseController::new().await.expect("Error initializing database"); let framework = poise::Framework::builder() .options(poise::FrameworkOptions { commands: vec![ commands::ping::ping(), - commands::eval::eval(), - commands::uptime::uptime(), - commands::sample::sample() + commands::sample::sample(), + commands::midi::midi_to_wav(), + commands::uptime::uptime() ], pre_command: |ctx| Box::pin(async move { let get_guild_name = match ctx.guild() { Some(guild) => guild.name.clone(), - None => String::from("DM") + None => String::from("Direct Message") }; - println!("[{}] {} ran /{}", get_guild_name, ctx.author().name, ctx.command().qualified_name) + println!("Discord[{}]: {} ran /{}", get_guild_name, ctx.author().name, ctx.command().qualified_name); }), on_error: |error| Box::pin(async move { match error { poise::FrameworkError::Command { error, ctx, .. } => { println!("PoiseCommandError({}): {}", ctx.command().qualified_name, error); - } - other => println!("PoiseOtherError: {:?}", other) + ctx.reply(format!( + "Encountered an error during command execution, ask {} to check console for more details!", + mention_dev(ctx).unwrap_or_default() + )).await.expect("Error sending message"); + }, + poise::FrameworkError::EventHandler { error, event, .. } => println!("PoiseEventHandlerError({}): {}", event.snake_case_name(), error), + poise::FrameworkError::Setup { error, .. } => println!("PoiseSetupError: {}", error), + poise::FrameworkError::UnknownInteraction { interaction, .. } => println!( + "PoiseUnknownInteractionError: {} tried to execute an unknown interaction ({})", + interaction.user.name, + interaction.data.name + ), + other => println!("PoiseOtherError: {}", other) } }), initialize_owners: true, + event_handler: |ctx, event, framework, _| Box::pin(event_processor(ctx, event, framework)), ..Default::default() }) .setup(|ctx, ready, framework| Box::pin(on_ready(ctx, ready, framework))) .build(); - let mut client = ClientBuilder::new(internals::utils::token_path().await.main, GatewayIntents::GUILDS) - .framework(framework) - .await.expect("Error creating client"); - - { - let mut data = client.data.write().await; - data.insert::(db); - } + let mut client = ClientBuilder::new( + token_path().await.main, + GatewayIntents::GUILDS + | GatewayIntents::MESSAGE_CONTENT + | GatewayIntents::DIRECT_MESSAGES + ) + .framework(framework) + .await.expect("Error creating client"); if let Err(why) = client.start().await { - println!("Client error: {:?}", why); + println!("Error starting client: {:#?}", why); } } diff --git a/src/models/mod.rs b/src/models.rs similarity index 100% rename from src/models/mod.rs rename to src/models.rs diff --git a/src/models/sample.rs b/src/models/sample.rs index 3f2bdf6..43d3b04 100644 --- a/src/models/sample.rs +++ b/src/models/sample.rs @@ -1,4 +1,6 @@ -use crate::controllers::database::DatabaseController; +use crate::controllers::database::DATABASE; + +use std::ops::Deref; pub struct SampleData { pub id: i64, @@ -9,7 +11,13 @@ pub struct SampleData { impl SampleData { pub async fn list_data(id: u64) -> Result, tokio_postgres::Error> { - let client = DatabaseController::new().await?.client; + let pool = { + let db = DATABASE.lock().unwrap(); + let controller = db.as_ref().unwrap(); + controller.pool.clone() + }; + let conn = pool.get().await.unwrap(); + let client = conn.deref(); let rows = client.query(" SELECT * FROM sample WHERE id = $1 @@ -33,7 +41,13 @@ impl SampleData { int: i64, boolean: bool ) -> Result<(), tokio_postgres::Error> { - let client = DatabaseController::new().await?.client; + let pool = { + let db = DATABASE.lock().unwrap(); + let controller = db.as_ref().unwrap(); + controller.pool.clone() + }; + let conn = pool.get().await.unwrap(); + let client = conn.deref(); client.execute(" INSERT INTO sample (text_val, int_val, boolean_val) VALUES ($1, $2, $3) @@ -48,7 +62,13 @@ impl SampleData { int: i64, boolean: bool ) -> Result<(), tokio_postgres::Error> { - let client = DatabaseController::new().await?.client; + let pool = { + let db = DATABASE.lock().unwrap(); + let controller = db.as_ref().unwrap(); + controller.pool.clone() + }; + let conn = pool.get().await.unwrap(); + let client = conn.deref(); client.execute(" UPDATE sample SET text_val = $1, int_val = $2, boolean_val = $3 @@ -59,7 +79,13 @@ impl SampleData { } pub async fn delete_data(id: u64) -> Result<(), tokio_postgres::Error> { - let client = DatabaseController::new().await?.client; + let pool = { + let db = DATABASE.lock().unwrap(); + let controller = db.as_ref().unwrap(); + controller.pool.clone() + }; + let conn = pool.get().await.unwrap(); + let client = conn.deref(); client.execute(" DELETE FROM sample WHERE id = $1