diff --git a/Cargo.lock b/Cargo.lock index 13cfdcfd..83ba3656 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1399,6 +1399,12 @@ dependencies = [ "slab", ] +[[package]] +name = "fuzzt" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a15f3d0fa42283a765e5fb609683ddab4ee4ff245d8db66a24d926c05e518c6" + [[package]] name = "generic-array" version = "0.14.7" @@ -1780,6 +1786,7 @@ dependencies = [ "derive_more", "eyre", "futures", + "fuzzt", "getrandom", "hex", "hkdf", diff --git a/Cargo.toml b/Cargo.toml index 5ede0dd7..a2077480 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,7 @@ homepage = "http://magic-wormhole.io/" repository = "https://github.com/magic-wormhole/magic-wormhole.rs/tree/main/cli" license = "EUPL-1.2" -rust-version = "1.75" +rust-version = "1.81" edition = "2021" [workspace.dependencies] @@ -38,7 +38,6 @@ futures = "0.3.12" hex = "0.4.2" hkdf = "0.12.2" indicatif = "0.17.0" -log = "0.4.13" noise-protocol = "0.2" noise-rust-crypto = "0.6.0-rc.1" number_prefix = "0.4.0" @@ -107,6 +106,8 @@ zxcvbn = { workspace = true, optional = true } tracing = { workspace = true, features = ["log", "log-always"] } +fuzzt = { version = "0.3.1", optional = true } + # Transit dependencies @@ -162,7 +163,7 @@ transit = [ "dep:async-trait", ] forwarding = ["transit", "dep:rmp-serde"] -default = ["transit", "transfer"] +default = ["transit", "transfer", "fuzzy-complete"] all = ["default", "forwarding"] # TLS implementations for websocket connections via async-tungstenite @@ -173,6 +174,7 @@ native-tls = ["async-tungstenite/async-native-tls"] # By enabling this option you are opting out of semver stability. experimental-transfer-v2 = [] experimental = ["experimental-transfer-v2"] +fuzzy-complete = ["fuzzt"] [profile.release] overflow-checks = true diff --git a/cli/Cargo.toml b/cli/Cargo.toml index c203bdf8..19947599 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -33,7 +33,7 @@ clap_complete = { workspace = true } env_logger = { workspace = true } console = { workspace = true } indicatif = { workspace = true } -dialoguer = { workspace = true } +dialoguer = { workspace = true, features = ["completion"] } color-eyre = { workspace = true } number_prefix = { workspace = true } ctrlc = { workspace = true } diff --git a/cli/src/completer.rs b/cli/src/completer.rs new file mode 100644 index 00000000..8ed0b0b0 --- /dev/null +++ b/cli/src/completer.rs @@ -0,0 +1,31 @@ +use std::sync::LazyLock; + +use color_eyre::eyre; +use dialoguer::{Completion, Input}; +use magic_wormhole::wordlist::{default_wordlist, Wordlist}; + +static WORDLIST: LazyLock = LazyLock::new(|| default_wordlist(2)); + +struct CustomCompletion {} + +impl CustomCompletion { + pub fn default() -> Self { + CustomCompletion {} + } +} + +impl Completion for CustomCompletion { + fn get(&self, input: &str) -> Option { + WORDLIST.get_completions(input).first().cloned() + } +} + +pub fn enter_code() -> eyre::Result { + let custom_completion = CustomCompletion::default(); + + Input::new() + .with_prompt("Wormhole Code") + .completion_with(&custom_completion) + .interact_text() + .map_err(From::from) +} diff --git a/cli/src/main.rs b/cli/src/main.rs index 35284632..6cd67779 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -1,4 +1,5 @@ #![allow(clippy::too_many_arguments)] +mod completer; mod util; use std::{ @@ -9,6 +10,7 @@ use std::{ use async_std::sync::Arc; use clap::{Args, CommandFactory, Parser, Subcommand}; use color_eyre::{eyre, eyre::Context}; +use completer::enter_code; use console::{style, Term}; use futures::{future::Either, Future, FutureExt}; use indicatif::{MultiProgress, ProgressBar}; @@ -750,15 +752,6 @@ fn create_progress_handler(pb: ProgressBar) -> impl FnMut(u64, u64) { } } -fn enter_code() -> eyre::Result { - use dialoguer::Input; - - Input::new() - .with_prompt("Enter code") - .interact_text() - .map_err(From::from) -} - fn print_welcome(term: &mut Term, welcome: Option<&str>) -> eyre::Result<()> { if let Some(welcome) = &welcome { writeln!(term, "Got welcome from server: {}", welcome)?; diff --git a/src/core.rs b/src/core.rs index 59f463cb..6d660665 100644 --- a/src/core.rs +++ b/src/core.rs @@ -5,7 +5,9 @@ pub mod rendezvous; mod server_messages; #[cfg(test)] mod test; -mod wordlist; + +/// Module for wormhole code generation and completion. +pub mod wordlist; use serde_derive::{Deserialize, Serialize}; use std::{borrow::Cow, str::FromStr}; diff --git a/src/core/wordlist.rs b/src/core/wordlist.rs index 9ba082ce..f059f5d4 100644 --- a/src/core/wordlist.rs +++ b/src/core/wordlist.rs @@ -1,12 +1,17 @@ +///! Wordlist generation and wormhole code utilities use rand::{rngs::OsRng, seq::SliceRandom}; use serde_json::{self, Value}; use std::fmt; use super::Password; +/// Represents a list of words used to generate and complete wormhole codes. +/// A wormhole code is a sequence of words used for secure communication or identification. #[derive(PartialEq)] pub struct Wordlist { - pub num_words: usize, + /// Number of words in a wormhole code + num_words: usize, + /// Odd and even wordlist words: Vec>, } @@ -18,45 +23,67 @@ impl fmt::Debug for Wordlist { impl Wordlist { #[cfg(test)] + #[doc(hidden)] pub fn new(num_words: usize, words: Vec>) -> Wordlist { Wordlist { num_words, words } } - #[allow(dead_code)] // TODO make this API public one day + /// Completes a wormhole code + /// + /// Completion can be done either with fuzzy search (approximate string matching) + /// or simple `starts_with` matching. pub fn get_completions(&self, prefix: &str) -> Vec { - let count_dashes = prefix.matches('-').count(); - let mut completions = Vec::new(); - let words = &self.words[count_dashes % self.words.len()]; - - let last_partial_word = prefix.split('-').last(); - let lp = if let Some(w) = last_partial_word { - w.len() - } else { - 0 - }; - - for word in words { - let mut suffix: String = prefix.to_owned(); - if word.starts_with(last_partial_word.unwrap()) { - if lp == 0 { - suffix.push_str(word); - } else { - let p = prefix.len() - lp; - suffix.truncate(p); - suffix.push_str(word); + let words = self.get_wordlist(prefix); + + let (prefix_without_last, last_partial) = prefix.rsplit_once('-').unwrap_or(("", prefix)); + + #[cfg(feature = "fuzzy-complete")] + let matches = self.fuzzy_complete(last_partial, words); + #[cfg(not(feature = "fuzzy-complete"))] + let matches = self.normal_complete(last_partial, words); + + matches + .into_iter() + .map(|word| { + let mut completion = String::new(); + completion.push_str(prefix_without_last); + if !prefix_without_last.is_empty() { + completion.push('-'); } + completion.push_str(&word); + completion + }) + .collect() + } - if count_dashes + 1 < self.num_words { - suffix.push('-'); - } + fn get_wordlist(&self, prefix: &str) -> &Vec { + let count_dashes = prefix.matches('-').count(); + &self.words[count_dashes % self.words.len()] + } - completions.push(suffix); - } - } - completions.sort(); - completions + #[cfg(feature = "fuzzy-complete")] + fn fuzzy_complete(&self, partial: &str, words: &[String]) -> Vec { + // We use Jaro-Winkler algorithm because it emphasizes the beginning of a word + use fuzzt::algorithms::JaroWinkler; + + let words = words.iter().map(|w| w.as_str()).collect::>(); + + fuzzt::get_top_n(partial, &words, None, None, None, Some(&JaroWinkler)) + .into_iter() + .map(|s| s.to_string()) + .collect() } + #[allow(unused)] + fn normal_complete(&self, partial: &str, words: &[String]) -> Vec { + words + .iter() + .filter(|word| word.starts_with(partial)) + .cloned() + .collect() + } + + /// Choose wormhole code word pub fn choose_words(&self) -> Password { let mut rng = OsRng; let components: Vec = self @@ -106,6 +133,7 @@ fn load_pgpwords() -> Vec> { vec![even_words, odd_words] } +/// Construct Wordlist struct with given number of words in a wormhole code pub fn default_wordlist(num_words: usize) -> Wordlist { Wordlist { num_words, @@ -155,8 +183,9 @@ mod test { ]; let w = Wordlist::new(2, words); - assert_eq!(w.get_completions(""), vec!["green-", "purple-", "yellow-"]); - assert_eq!(w.get_completions("pur"), vec!["purple-"]); + assert_eq!(w.get_completions(""), Vec::::new()); + assert_eq!(w.get_completions("9"), Vec::::new()); + assert_eq!(w.get_completions("pur"), vec!["purple"]); assert_eq!(w.get_completions("blu"), Vec::::new()); assert_eq!(w.get_completions("purple-sa"), vec!["purple-sausages"]); } @@ -197,45 +226,59 @@ mod test { } #[test] - fn test_default_completions() { - let w = default_wordlist(2); - let c = w.get_completions("ar"); - assert_eq!(c.len(), 2); - assert!(c.contains(&String::from("article-"))); - assert!(c.contains(&String::from("armistice-"))); + #[cfg(feature = "fuzzy-complete")] + fn test_wormhole_code_fuzzy_completions() { + let list = default_wordlist(2); + + assert_eq!(list.get_completions("22"), Vec::::new()); + assert_eq!(list.get_completions("22-"), Vec::::new()); - let c = w.get_completions("armis"); - assert_eq!(c.len(), 1); - assert!(c.contains(&String::from("armistice-"))); + // Invalid wormhole code check + assert_eq!(list.get_completions("trj"), Vec::::new()); - let c = w.get_completions("armistice-"); - assert_eq!(c.len(), 256); + assert_eq!( + list.get_completions("22-chisel"), + ["22-chisel", "22-chairlift", "22-christmas"] + ); - let c = w.get_completions("armistice-ba"); assert_eq!( - c, - vec![ - "armistice-baboon", - "armistice-backfield", - "armistice-backward", - "armistice-banjo", - ] + list.get_completions("22-chle"), + ["22-chisel", "22-chatter", "22-checkup"] ); - let w = default_wordlist(3); - let c = w.get_completions("armistice-ba"); + assert_eq!(list.get_completions("22-chisel-tba"), ["22-chisel-tobacco"]); + } + + #[test] + #[cfg(feature = "fuzzy-complete")] + fn test_completion_fuzzy() { + let wl = default_wordlist(2); + let list = wl.get_wordlist("22-"); + + assert_eq!(wl.fuzzy_complete("chck", list), ["checkup", "choking"]); + assert_eq!(wl.fuzzy_complete("checkp", list), ["checkup"]); assert_eq!( - c, - vec![ - "armistice-baboon-", - "armistice-backfield-", - "armistice-backward-", - "armistice-banjo-", - ] + wl.fuzzy_complete("checkup", list), + ["checkup", "lockup", "cleanup"] ); + } - let w = default_wordlist(4); - let c = w.get_completions("armistice-baboon"); - assert_eq!(c, vec!["armistice-baboon-"]); + #[test] + fn test_completion_normal() { + let wl = default_wordlist(2); + let list = wl.get_wordlist("22-"); + + assert_eq!(wl.normal_complete("che", list), ["checkup"]); + } + + #[test] + fn test_full_wormhole_completion() { + let wl = default_wordlist(2); + + assert_eq!(wl.get_completions("22-chec").first().unwrap(), "22-checkup"); + assert_eq!( + wl.get_completions("22-checkup-t").first().unwrap(), + "22-checkup-tobacco" + ); } } diff --git a/src/lib.rs b/src/lib.rs index bd0438e1..03384d6b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,6 +28,7 @@ #[macro_use] mod util; mod core; +pub use core::wordlist; #[cfg(feature = "forwarding")] pub mod forwarding; #[cfg(feature = "transfer")]