فهرست منبع

feat: add `--genkey`

Yujia Qiao 4 سال پیش
والد
کامیت
46ec6400e4
4فایلهای تغییر یافته به همراه61 افزوده شده و 11 حذف شده
  1. 23 6
      src/cli.rs
  2. 33 2
      src/lib.rs
  3. 1 1
      src/transport/noise.rs
  4. 4 2
      tests/common/mod.rs

+ 23 - 6
src/cli.rs

@@ -1,20 +1,37 @@
-use clap::{AppSettings, Parser};
+use clap::{AppSettings, ArgGroup, Parser};
 
-#[derive(Parser, Debug)]
+#[derive(clap::ArgEnum, Clone, Debug, Copy)]
+pub enum KeypairType {
+    X25519,
+    X448,
+}
+
+#[derive(Parser, Debug, Default)]
 #[clap(about, version, setting(AppSettings::DeriveDisplayOrder))]
+#[clap(group(
+            ArgGroup::new("cmds")
+                .required(true)
+                .args(&["CONFIG", "genkey"]),
+        ))]
 pub struct Cli {
     /// The path to the configuration file
     ///
     /// Running as a client or a server is automatically determined
     /// according to the configuration file.
-    #[clap(parse(from_os_str), name = "config")]
-    pub config_path: std::path::PathBuf,
+    #[clap(parse(from_os_str), name = "CONFIG")]
+    pub config_path: Option<std::path::PathBuf>,
 
     /// Run as a server
-    #[clap(long, short)]
+    #[clap(long, short, group = "mode")]
     pub server: bool,
 
     /// Run as a client
-    #[clap(long, short)]
+    #[clap(long, short, group = "mode")]
     pub client: bool,
+
+    /// Generate a keypair for the use of the noise protocol
+    ///
+    /// The DH function to use is x25519
+    #[clap(long, arg_enum, value_name = "CURVE")]
+    pub genkey: Option<Option<KeypairType>>,
 }

+ 33 - 2
src/lib.rs

@@ -7,6 +7,7 @@ mod protocol;
 mod transport;
 
 pub use cli::Cli;
+use cli::KeypairType;
 pub use config::Config;
 pub use constants::UDP_BUFFER_SIZE;
 
@@ -24,8 +25,37 @@ mod server;
 #[cfg(feature = "server")]
 use server::run_server;
 
+const DEFAULT_CURVE: KeypairType = KeypairType::X25519;
+
+fn get_str_from_keypair_type(curve: KeypairType) -> &'static str {
+    match curve {
+        KeypairType::X25519 => "25519",
+        KeypairType::X448 => "448",
+    }
+}
+
+fn genkey(curve: Option<KeypairType>) -> Result<()> {
+    let curve = curve.unwrap_or(DEFAULT_CURVE);
+    let builder = snowstorm::Builder::new(
+        format!(
+            "Noise_KK_{}_ChaChaPoly_BLAKE2s",
+            get_str_from_keypair_type(curve)
+        )
+        .parse()?,
+    );
+    let keypair = builder.generate_keypair()?;
+
+    println!("Private Key:\n{}\n", base64::encode(keypair.private));
+    println!("Public Key:\n{}", base64::encode(keypair.public));
+    Ok(())
+}
+
 pub async fn run(args: &Cli, shutdown_rx: broadcast::Receiver<bool>) -> Result<()> {
-    let config = Config::from_file(&args.config_path).await?;
+    if args.genkey.is_some() {
+        return genkey(args.genkey.unwrap());
+    }
+
+    let config = Config::from_file(args.config_path.as_ref().unwrap()).await?;
 
     debug!("{:?}", config);
 
@@ -155,9 +185,10 @@ mod tests {
             };
 
             let args = Cli {
-                config_path: std::path::PathBuf::new(),
+                config_path: Some(std::path::PathBuf::new()),
                 server: t.arg_s,
                 client: t.arg_c,
+                ..Default::default()
             };
 
             assert_eq!(determine_run_mode(&config, &args), t.run_mode);

+ 1 - 1
src/transport/noise.rs

@@ -54,7 +54,7 @@ impl Transport for NoiseTransport {
             None => builder.generate_keypair()?.private,
         };
 
-        let params: NoiseParams = config.pattern.clone().parse()?;
+        let params: NoiseParams = config.pattern.parse()?;
 
         Ok(NoiseTransport {
             config,

+ 4 - 2
tests/common/mod.rs

@@ -15,9 +15,10 @@ pub async fn run_rathole_server(
     shutdown_rx: broadcast::Receiver<bool>,
 ) -> Result<()> {
     let cli = rathole::Cli {
-        config_path: PathBuf::from(config_path),
+        config_path: Some(PathBuf::from(config_path)),
         server: true,
         client: false,
+        ..Default::default()
     };
     rathole::run(&cli, shutdown_rx).await
 }
@@ -27,9 +28,10 @@ pub async fn run_rathole_client(
     shutdown_rx: broadcast::Receiver<bool>,
 ) -> Result<()> {
     let cli = rathole::Cli {
-        config_path: PathBuf::from(config_path),
+        config_path: Some(PathBuf::from(config_path)),
         server: false,
         client: true,
+        ..Default::default()
     };
     rathole::run(&cli, shutdown_rx).await
 }