Skip to content

Commit

Permalink
feat: add an allow list for access control
Browse files Browse the repository at this point in the history
Implementation contains a known bug, in that if the allow_list includes the name of the bot it will respond to it's own messages.
  • Loading branch information
arcuru committed Mar 19, 2024
1 parent caf145a commit 5f0c442
Show file tree
Hide file tree
Showing 3 changed files with 35 additions and 2 deletions.
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 3 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,4 +21,6 @@ matrix-sdk = "0.7.1"
ollama-rs = "0.1.0"
serde = { version = "1.0", features = ["derive"] }
serde_yaml = "0.9"
clap = { version = "4.2.1", features = ["derive"] }
clap = { version = "4.2.1", features = ["derive"] }
lazy_static = "1.4.0"
regex = "1.10.3"
31 changes: 30 additions & 1 deletion src/main.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use clap::Parser;
use lazy_static::lazy_static;
use matrix_sdk::{
config::SyncSettings,
ruma::events::room::{
Expand All @@ -8,10 +9,12 @@ use matrix_sdk::{
Client, Room, RoomState,
};
use ollama_rs::{generation::completion::request::GenerationRequest, Ollama};
use regex::Regex;
use serde::Deserialize;
use std::fs::File;
use std::io::Read;
use std::path::PathBuf;
use std::sync::Mutex;
use tokio::time::{sleep, Duration};

#[derive(Parser)]
Expand All @@ -22,11 +25,17 @@ struct HeadJackArgs {
config: PathBuf,
}

#[derive(Debug, Deserialize)]
#[derive(Debug, Deserialize, Clone)]
struct Config {
homeserver_url: String,
username: String,
password: String,
/// Allow list of which accounts we will respond to
allow_list: Option<String>,
}

lazy_static! {
static ref GLOBAL_CONFIG: Mutex<Option<Config>> = Mutex::new(None);
}

/// This is the starting point of the app. `main` is called by rust binaries to
Expand All @@ -45,12 +54,24 @@ async fn main() -> anyhow::Result<()> {
file.read_to_string(&mut contents)?;

let config: Config = serde_yaml::from_str(&contents)?;
*GLOBAL_CONFIG.lock().unwrap() = Some(config.clone());

// our actual runner
login_and_sync(config.homeserver_url, &config.username, &config.password).await?;
Ok(())
}

/// Verify if the sender is on the allow_list
fn is_allowed(sender: &str) -> bool {
let config = GLOBAL_CONFIG.lock().unwrap().clone().unwrap();
// FIXME: Check to see if it's from ourselves, in which case we should do nothing
if let Some(allow_list) = config.allow_list {
let regex = Regex::new(&allow_list).expect("Invalid regular expression");
return regex.is_match(sender);
}
false
}

// The core sync loop we have running.
async fn login_and_sync(
homeserver_url: String,
Expand Down Expand Up @@ -118,6 +139,10 @@ async fn on_stripped_state_member(
// the invite we've seen isn't for us, but for someone else. ignore
return;
}
if !is_allowed(room_member.sender.as_str()) {
// Sender is not on the allowlist
return;
}

// The event handlers are called before the next sync begins, but
// methods that change the state of a room (joining, leaving a room)
Expand Down Expand Up @@ -162,6 +187,10 @@ async fn on_room_message(event: OriginalSyncRoomMessageEvent, room: Room) {
let MessageType::Text(text_content) = event.content.msgtype else {
return;
};
if !is_allowed(event.sender.as_str()) {
// Sender is not on the allowlist
return;
}

// If we start with a single '!', interpret as a command
let text = text_content.body.trim_start();
Expand Down

0 comments on commit 5f0c442

Please sign in to comment.