initial commit

src/filters/mod.rs

@@ -0,0 +1,5 @@
+mod rate_limit;
+mod whitelist;
+
+pub use rate_limit::RateLimit;
+pub use whitelist::Whitelist;

src/filters/rate_limit.rs

@@ -0,0 +1,54 @@
+use crate::{Action, InputMessage, NoteFilter, OutputMessage};
+use serde::Deserialize;
+use std::collections::HashMap;
+use std::time::{Duration, Instant};
+
+pub struct RateInfo {
+    pub last_note: Instant,
+}
+
+#[derive(Deserialize, Default)]
+pub struct RateLimit {
+    pub notes_per_second: u64,
+    pub whitelist: Option<Vec<String>>,
+
+    #[serde(skip)]
+    pub sources: HashMap<String, RateInfo>,
+}
+
+impl NoteFilter for RateLimit {
+    fn filter_note(&mut self, msg: &InputMessage) -> OutputMessage {
+        if let Some(whitelist) = &self.whitelist {
+            if whitelist.contains(&msg.source_info) {
+                return OutputMessage::new(msg.event.id.clone(), Action::Accept, None);
+            }
+        }
+
+        if self.sources.contains_key(&msg.source_info) {
+            let now = Instant::now();
+            let entry = self.sources.get_mut(&msg.source_info).expect("impossiburu");
+            if now - entry.last_note < Duration::from_secs(self.notes_per_second) {
+                return OutputMessage::new(
+                    msg.event.id.clone(),
+                    Action::Reject,
+                    Some("rate-limited: you are noting too fast".to_string()),
+                );
+            } else {
+                entry.last_note = Instant::now();
+                return OutputMessage::new(msg.event.id.clone(), Action::Accept, None);
+            }
+        } else {
+            self.sources.insert(
+                msg.source_info.to_owned(),
+                RateInfo {
+                    last_note: Instant::now(),
+                },
+            );
+            return OutputMessage::new(msg.event.id.clone(), Action::Accept, None);
+        }
+    }
+
+    fn name(&self) -> &'static str {
+        "ratelimit"
+    }
+}

src/filters/whitelist.rs

@@ -0,0 +1,26 @@
+use crate::{Action, InputMessage, NoteFilter, OutputMessage};
+use serde::Deserialize;
+
+#[derive(Deserialize)]
+pub struct Whitelist {
+    pub pubkeys: Vec<String>,
+    pub ips: Vec<String>,
+}
+
+impl NoteFilter for Whitelist {
+    fn filter_note(&mut self, msg: &InputMessage) -> OutputMessage {
+        if self.pubkeys.contains(&msg.event.pubkey) || self.ips.contains(&msg.source_info) {
+            OutputMessage::new(msg.event.id.clone(), Action::Accept, None)
+        } else {
+            OutputMessage::new(
+                msg.event.id.clone(),
+                Action::Reject,
+                Some("blocked: pubkey not on the whitelist".to_string()),
+            )
+        }
+    }
+
+    fn name(&self) -> &'static str {
+        "whitelist"
+    }
+}

src/lib.rs

@@ -0,0 +1,6 @@
+pub mod filters;
+mod messages;
+mod note_filter;
+
+pub use messages::{Action, InputMessage, OutputMessage};
+pub use note_filter::{Note, NoteFilter};

src/main.rs

@@ -0,0 +1,143 @@
+use noteguard::filters::RateLimit;
+use noteguard::{Action, InputMessage, NoteFilter, OutputMessage};
+use serde::de::DeserializeOwned;
+use serde::Deserialize;
+use std::collections::HashMap;
+use std::io::{self, BufRead, Read, Write};
+
+#[derive(Deserialize)]
+struct Config {
+    pipeline: Vec<String>,
+    filters: HashMap<String, toml::Value>,
+}
+
+type ConstructFilter = Box<fn(toml::Value) -> Result<Box<dyn NoteFilter>, toml::de::Error>>;
+
+#[derive(Default)]
+struct Noteguard {
+    registered_filters: HashMap<String, ConstructFilter>,
+    loaded_filters: Vec<Box<dyn NoteFilter>>,
+}
+
+impl Noteguard {
+    pub fn new() -> Self {
+        let mut noteguard = Noteguard::default();
+        noteguard.register_builtin_filters();
+        noteguard
+    }
+
+    pub fn register_filter<F: NoteFilter + 'static + Default + DeserializeOwned>(&mut self) {
+        self.registered_filters.insert(
+            F::name(&F::default()).to_string(),
+            Box::new(|filter_config| {
+                filter_config
+                    .try_into()
+                    .map(|filter: F| Box::new(filter) as Box<dyn NoteFilter>)
+            }),
+        );
+    }
+
+    /// All builtin filters are registered here, and are made available with
+    /// every new instance of [`Noteguard`]
+    fn register_builtin_filters(&mut self) {
+        self.register_filter::<RateLimit>();
+    }
+
+    /// Run the loaded filters. You must call `load_config` before calling this, otherwise
+    /// not filters will be run.
+    fn run(&mut self, input: InputMessage) -> OutputMessage {
+        let mut mout: Option<OutputMessage> = None;
+
+        let id = input.event.id.clone();
+        for filter in &mut self.loaded_filters {
+            let out = filter.filter_note(&input);
+            match out.action {
+                Action::Accept => {
+                    mout = Some(out);
+                    continue;
+                }
+                Action::Reject => {
+                    return out;
+                }
+                Action::ShadowReject => {
+                    return out;
+                }
+            }
+        }
+
+        mout.unwrap_or_else(|| OutputMessage::new(id, Action::Accept, None))
+    }
+
+    /// Initializes a noteguard config. If it finds any filter configurations
+    /// matching the registered filters, it loads those into our filter pipeline.
+    fn load_config(&mut self, config: &Config) -> Result<(), toml::de::Error> {
+        self.loaded_filters.clear();
+
+        for (name, config_value) in &config.filters {
+            if let Some(constructor) = self.registered_filters.get(name.as_str()) {
+                let filter = constructor(config_value.clone())?;
+                self.loaded_filters.push(filter);
+            } else {
+                panic!("Found config settings with no matching filter: {}", name);
+            }
+        }
+
+        Ok(())
+    }
+}
+
+fn main() {
+    let config_path = "noteguard.toml";
+    let mut noteguard = Noteguard::new();
+
+    let config: Config = {
+        let mut file = std::fs::File::open(config_path).expect("Failed to open config file");
+        let mut contents = String::new();
+        file.read_to_string(&mut contents)
+            .expect("Failed to read config file");
+        toml::from_str(&contents).expect("Failed to parse config file")
+    };
+
+    noteguard
+        .load_config(&config)
+        .expect("Expected filter config to be loaded ok");
+
+    let stdin = io::stdin();
+    let stdout = io::stdout();
+    let handle = stdout.lock();
+    let mut writer = io::BufWriter::new(handle);
+
+    for line in stdin.lock().lines() {
+        match line {
+            Ok(input) => {
+                let input_message: InputMessage = match serde_json::from_str(&input) {
+                    Ok(msg) => msg,
+                    Err(e) => {
+                        eprintln!("Failed to parse input: {}", e);
+                        continue;
+                    }
+                };
+
+                if input_message.message_type != "new" {
+                    eprintln!("Unexpected request type");
+                    continue;
+                }
+
+                let output_message = noteguard.run(input_message);
+
+                match serde_json::to_string(&output_message) {
+                    Ok(json) => {
+                        writeln!(writer, "{}", json).unwrap();
+                        writer.flush().unwrap();
+                    }
+                    Err(e) => {
+                        eprintln!("Failed to serialize output: {}", e);
+                    }
+                }
+            }
+            Err(e) => {
+                eprintln!("Failed to read line: {}", e);
+            }
+        }
+    }
+}

src/messages.rs

@@ -0,0 +1,37 @@
+use crate::Note;
+use serde::{Deserialize, Serialize};
+
+#[derive(Deserialize)]
+pub struct InputMessage {
+    #[serde(rename = "type")]
+    pub message_type: String,
+    pub event: Note,
+    #[serde(rename = "receivedAt")]
+    pub received_at: u64,
+    #[serde(rename = "sourceType")]
+    pub source_type: String,
+    #[serde(rename = "sourceInfo")]
+    pub source_info: String,
+}
+
+#[derive(Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub enum Action {
+    Accept,
+    Reject,
+    ShadowReject,
+}
+
+#[derive(Serialize)]
+pub struct OutputMessage {
+    pub id: String,
+    pub action: Action,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub msg: Option<String>,
+}
+
+impl OutputMessage {
+    pub fn new(id: String, action: Action, msg: Option<String>) -> Self {
+        OutputMessage { id, action, msg }
+    }
+}

src/note_filter.rs

@@ -0,0 +1,20 @@
+use crate::{InputMessage, OutputMessage};
+use serde::Deserialize;
+
+#[derive(Deserialize)]
+pub struct Note {
+    pub id: String,
+    pub pubkey: String,
+    pub content: String,
+    pub created_at: i64,
+    pub kind: i64,
+    pub tags: Vec<Vec<String>>,
+    pub sig: String,
+}
+
+pub trait NoteFilter {
+    fn filter_note(&mut self, msg: &InputMessage) -> OutputMessage;
+
+    /// A key corresponding to an entry in the noteguard.toml file.
+    fn name(&self) -> &'static str;
+}