diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index aa403d213f..3c69688c9d 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -355,6 +355,7 @@ pub fn create_default_context(cwd: impl AsRef) -> EngineState { UrlPath, UrlQuery, UrlScheme, + Port, } // Random diff --git a/crates/nu-command/src/network/mod.rs b/crates/nu-command/src/network/mod.rs index 543e6e1c1c..330e1f17c8 100644 --- a/crates/nu-command/src/network/mod.rs +++ b/crates/nu-command/src/network/mod.rs @@ -1,7 +1,9 @@ mod fetch; +mod port; mod post; mod url; pub use self::url::*; pub use fetch::SubCommand as Fetch; +pub use port::SubCommand as Port; pub use post::SubCommand as Post; diff --git a/crates/nu-command/src/network/port.rs b/crates/nu-command/src/network/port.rs new file mode 100644 index 0000000000..3b38988f53 --- /dev/null +++ b/crates/nu-command/src/network/port.rs @@ -0,0 +1,109 @@ +use nu_engine::CallExt; +use nu_protocol::ast::Call; +use nu_protocol::engine::{Command, EngineState, Stack}; +use nu_protocol::IntoPipelineData; +use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, TcpListener}; + +use nu_protocol::{ + Category, Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Value, +}; + +#[derive(Clone)] +pub struct SubCommand; + +impl Command for SubCommand { + fn name(&self) -> &str { + "port" + } + + fn signature(&self) -> Signature { + Signature::build("post") + .optional( + "start", + SyntaxShape::Int, + "The start port to scan (inclusive)", + ) + .optional("end", SyntaxShape::Int, "The end port to scan (inclusive)") + .category(Category::Network) + } + + fn usage(&self) -> &str { + "Get a free port from system" + } + + fn search_terms(&self) -> Vec<&str> { + vec!["network", "http", "port"] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + _input: PipelineData, + ) -> Result { + get_free_port(engine_state, stack, call) + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "get free port between 3121 and 4000", + example: "port 3121 4000", + result: Some(Value::Int { + val: 3121, + span: Span::test_data(), + }), + }, + Example { + description: "get free port from system", + example: "port", + result: None, + }, + ] + } +} + +fn get_free_port( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, +) -> Result { + let start_port: Option = call.opt(engine_state, stack, 0)?; + let end_port: Option = call.opt(engine_state, stack, 1)?; + + let listener = if start_port.is_none() && end_port.is_none() { + // get free port from system. + TcpListener::bind("127.0.0.1:0")? + } else { + let start_port = start_port.unwrap_or(1024); + let end_port = end_port.unwrap_or(65535); + + // check input range valid. + if start_port > end_port { + return Err(ShellError::InvalidRange( + start_port.to_string(), + end_port.to_string(), + call.head, + )); + } + + // try given port one by one. + let addrs: Vec = (start_port..=end_port) + .map(|current| { + SocketAddr::V4(SocketAddrV4::new( + Ipv4Addr::new(127, 0, 0, 1), + current as u16, + )) + }) + .collect(); + TcpListener::bind(addrs.as_slice())? + }; + + let free_port = listener.local_addr()?.port(); + Ok(Value::Int { + val: free_port as i64, + span: call.head, + } + .into_pipeline_data()) +} diff --git a/crates/nu-command/tests/commands/mod.rs b/crates/nu-command/tests/commands/mod.rs index f006f29050..662ff627cd 100644 --- a/crates/nu-command/tests/commands/mod.rs +++ b/crates/nu-command/tests/commands/mod.rs @@ -36,6 +36,7 @@ mod math; mod merge; mod mkdir; mod move_; +mod network; mod open; mod parse; mod path; diff --git a/crates/nu-command/tests/commands/network/mod.rs b/crates/nu-command/tests/commands/network/mod.rs new file mode 100644 index 0000000000..87ed03c5a4 --- /dev/null +++ b/crates/nu-command/tests/commands/network/mod.rs @@ -0,0 +1 @@ +mod port; diff --git a/crates/nu-command/tests/commands/network/port.rs b/crates/nu-command/tests/commands/network/port.rs new file mode 100644 index 0000000000..3161d82ca5 --- /dev/null +++ b/crates/nu-command/tests/commands/network/port.rs @@ -0,0 +1,52 @@ +use nu_test_support::{nu, pipeline}; +use std::net::TcpListener; +use std::sync::mpsc; + +#[test] +fn port_with_invalid_range() { + let actual = nu!( + cwd: ".", pipeline( + r#" + port 4000 3999 + "# + )); + + assert!(actual.err.contains("Invalid range")) +} + +#[test] +fn port_with_already_usage() { + let (tx, rx) = mpsc::sync_channel(0); + + // let system pick a free port for us. + let free_port = { + let listener = TcpListener::bind("127.0.0.1:0").expect("failed to pick a port"); + listener.local_addr().unwrap().port() + }; + let handler = std::thread::spawn(move || { + let _listener = TcpListener::bind(format!("127.0.0.1:{free_port}")); + let _ = rx.recv(); + }); + let actual = nu!( + cwd: ".", pipeline(&format!("port {free_port} {free_port}")) + ); + let _ = tx.send(true); + // make sure that the thread is closed and we release the port. + handler.join().unwrap(); + + // check for error kind str. + assert!(actual.err.contains("AddrInUse")) +} + +#[test] +fn port_from_system_given() { + let actual = nu!( + cwd: ".", pipeline( + r#" + port + "# + )); + + // check that we can get an integer port from system. + assert!(actual.out.parse::().unwrap() > 0) +}