Add tab completion to sn0int repl

src/repl/complete.rs

@@ -1,3 +1,4 @@
+use crate::repl::tokenize::{self, Token};
 use rustyline::Context;
 use rustyline::completion::Completer;
 use rustyline::highlight::Highlighter;
@@ -5,7 +6,16 @@ use rustyline::hint::Hinter;
 use std::borrow::Cow;
 
 
-pub struct ReplCompleter;
+#[derive(Default)]
+pub struct ReplCompleter {
+    globals: Vec<String>,
+}
+
+impl ReplCompleter {
+    pub fn set(&mut self, globals: Vec<String>) {
+        self.globals = globals;
+    }
+}
 
 impl rustyline::Helper for ReplCompleter {}
 
@@ -13,13 +23,25 @@ impl Completer for ReplCompleter {
     type Candidate = String;
 
     #[inline]
-    fn complete(&self, _line: &str, pos: usize, _ctx: &Context<'_>) -> rustyline::Result<(usize, Vec<String>)> {
+    fn complete(&self, line: &str, pos: usize, _ctx: &Context<'_>) -> rustyline::Result<(usize, Vec<String>)> {
         if pos == 0 {
             Ok((0, vec![
                 String::from("return "),
             ]))
         } else {
-            Ok((0, vec![]))
+            let filter = match tokenize::parse_last(&line[..pos]) {
+                Token::Name(name) => name,
+                Token::Empty => String::new(),
+                _ => return Ok((0, vec![])),
+            };
+
+            let mut options = Vec::new();
+            for g in &self.globals {
+                if g.starts_with(&filter) {
+                    options.push(g.to_string());
+                }
+            }
+            Ok((pos - filter.len(), options))
         }
     }
 }

src/repl/mod.rs

@@ -1,18 +1,82 @@
 use crate::config::Config;
 use crate::errors::*;
 use crate::engine::{ctx, Environment, DummyReporter};
-use crate::engine::ctx::State;
+use crate::engine::ctx::{State, LuaState};
 use crate::geoip::{Maxmind, AsnDB, GeoIP};
-use crate::hlua::AnyLuaValue;
+use crate::hlua::{Lua, AnyLuaValue};
 use crate::psl::PslReader;
 use crate::shell::readline::Readline;
 use crate::runtime::format_lua;
 use chrootable_https::Resolver;
 use std::collections::HashMap;
+use std::sync::Arc;
 
 mod complete;
 use self::complete::ReplCompleter;
+mod tokenize;
 
+pub struct Repl<'a> {
+    rl: Readline<ReplCompleter>,
+    lua: Lua<'a>,
+    state: Arc<LuaState>,
+}
+
+impl<'a> Repl<'a> {
+    pub fn new(lua: Lua<'a>, state: Arc<LuaState>) -> Repl<'a> {
+        let rl = Readline::with(ReplCompleter::default());
+        Repl {
+            rl,
+            lua,
+            state,
+        }
+    }
+
+    fn update_globals(&mut self) {
+        let mut globals = Vec::new();
+        for item in self.lua.globals_table().iter::<String, AnyLuaValue>() {
+            if let Some((k, _)) = item {
+                globals.push(k);
+            }
+        }
+        if let Some(helper) = self.rl.helper_mut() {
+            debug!("updating globals: {:?}", globals);
+            helper.set(globals);
+        }
+    }
+
+    pub fn run(&mut self) {
+        loop {
+            self.update_globals();
+
+            match self.rl.readline("> ") {
+                Ok(line) => {
+                    self.rl.add_history_entry(line.as_str());
+                    self.exec(&line);
+                },
+                Err(_) => break,
+            }
+        }
+    }
+
+    pub fn exec(&mut self, line: &str) {
+        match self.lua.execute::<AnyLuaValue>(line) {
+            Ok(val) => {
+                if val != AnyLuaValue::LuaNil {
+                    let mut out = String::new();
+                    format_lua(&mut out, &val);
+                    println!("{}", out);
+                }
+                if let Some(err) = self.state.last_error() {
+                    println!("Error: {}", err);
+                    self.state.clear_error();
+                }
+            },
+            Err(err) => {
+                println!("Fatal: {}", err);
+            }
+        }
+    }
+}
 
 pub fn run(config: &Config) -> Result<()> {
     let keyring = Vec::new();
@@ -35,39 +99,15 @@ pub fn run(config: &Config) -> Result<()> {
     };
 
     let tx = DummyReporter::new();
-    let (mut lua, state) = ctx::ctx(env, tx);
+    let (lua, state) = ctx::ctx(env, tx);
+    let mut repl = Repl::new(lua, state);
 
     println!(r#":: sn0int v{} lua repl
 Assign variables with `a = sn0int_version()` and `return a` to print
 Read the docs at https://sn0int.readthedocs.io/en/stable/reference.html
 "#, env!("CARGO_PKG_VERSION"));
 
-    let mut rl = Readline::with(ReplCompleter);
-    loop {
-        match rl.readline("> ") {
-            Ok(line) => {
-                rl.add_history_entry(line.as_str());
-
-                match lua.execute::<AnyLuaValue>(&line) {
-                    Ok(val) => {
-                        if val != AnyLuaValue::LuaNil {
-                            let mut out = String::new();
-                            format_lua(&mut out, &val);
-                            println!("{}", out);
-                        }
-                        if let Some(err) = state.last_error() {
-                            println!("Error: {}", err);
-                            state.clear_error();
-                        }
-                    },
-                    Err(err) => {
-                        println!("Fatal: {}", err);
-                    }
-                }
-            },
-            Err(_) => break,
-        }
-    }
+    repl.run();
 
     Ok(())
 }

src/repl/tokenize.rs

@@ -0,0 +1,119 @@
+#[derive(Debug, PartialEq)]
+pub enum Token {
+    Empty,
+    Name(String),
+    InString(StringState),
+    Value(String),
+}
+
+#[derive(Debug, PartialEq)]
+pub struct StringState {
+    quote: char,
+    escape: bool,
+    buf: String,
+}
+
+impl StringState {
+    pub fn new(quote: char) -> StringState {
+        StringState {
+            quote,
+            escape: false,
+            buf: String::new(),
+        }
+    }
+}
+
+pub fn parse_last(line: &str) -> Token {
+    let mut token = Token::Empty;
+
+    for c in line.chars() {
+        match &mut token {
+            Token::Empty => {
+                match c {
+                    'a'..='z' | 'A'..='Z' | '_' => {
+                        token = Token::Name(c.to_string());
+                    },
+                    '0'..='9' => {
+                        token = Token::Value(c.to_string());
+                    },
+                    '\'' | '"' => {
+                        token = Token::InString(StringState::new(c));
+                    },
+                    // ignore operators
+                    _ => (),
+                }
+            },
+            Token::Name(s) => {
+                match c {
+                    'a'..='z' | 'A'..='Z' | '0'..='9' | '_' => {
+                        s.push(c);
+                    },
+                    // ignore operators
+                    _ => {
+                        token = Token::Empty;
+                    },
+                }
+            },
+            Token::InString(s) => {
+                if s.escape {
+                    s.buf.push(c);
+                    s.escape = false;
+                } else if c == '\\' {
+                    s.escape = true;
+                } else if c == s.quote {
+                    // done
+                    token = Token::Empty;
+                } else {
+                    s.buf.push(c);
+                }
+            },
+            Token::Value(s) => {
+                match c {
+                    '0'..='9' => {
+                        s.push(c);
+                    },
+                    // ignore operators
+                    // a name can't directly follow a value
+                    _ => {
+                        token = Token::Empty;
+                    },
+                }
+            },
+        }
+    }
+
+    token
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    pub fn test_empty() {
+        let token = parse_last("");
+        assert_eq!(token, Token::Empty);
+    }
+
+    #[test]
+    pub fn test_abc() {
+        let token = parse_last("abc");
+        assert_eq!(token, Token::Name("abc".into()));
+    }
+
+    #[test]
+    pub fn test_in_string() {
+        let token = parse_last("return url_encode(\"asdf");
+        assert_eq!(token, Token::InString(StringState {
+            quote: '"',
+            escape: false,
+            buf: "asdf".into(),
+        }));
+    }
+
+    #[test]
+    pub fn test_in_func() {
+        let token = parse_last("return url_encode(\"asdf\") .. url_deco");
+        assert_eq!(token, Token::Name("url_deco".into()));
+    }
+}