feat: add --addr flag to daft-dashboard cli (#5444)

Author: VOID001

src/daft-cli/src/python.rs

@@ -1,11 +1,17 @@
-use std::str::FromStr;
+use std::{
+    net::{IpAddr, SocketAddr},
+    str::FromStr,
+};
 
 use clap::{Args, Parser, Subcommand, arg};
 use pyo3::prelude::*;
 use tracing_subscriber::{self, filter::Directive, layer::SubscriberExt, util::SubscriberInitExt};
 
 #[derive(Debug, Args)]
 struct DashboardArgs {
+    /// The address to launch the dashboard on
+    #[arg(short, long, default_value = "0.0.0.0")]
+    addr: IpAddr,
     #[arg(short, long, default_value_t = 80)]
     /// The port to launch the dashboard on
     port: u16,
@@ -35,6 +41,15 @@ fn run_dashboard(py: Python, args: DashboardArgs) {
     let filter = Directive::from_str(if args.verbose { "INFO" } else { "ERROR" })
         .expect("Failed to parse tracing filter");
 
+    if args.addr.is_unspecified() {
+        println!("{}", console::style(format!(
+            "⚠️  Listening on all network interfaces ({})! This is not recommended in production.",
+            args.addr
+        )).yellow().bold());
+    }
+
+    let socket_addr = SocketAddr::from((args.addr, args.port));
+
     // Set the subscriber for the detached run
     tracing_subscriber::registry()
         .with(
@@ -56,26 +71,25 @@ fn run_dashboard(py: Python, args: DashboardArgs) {
             "{}  To get started, run your Daft script with env `{}`",
             console::style("█").magenta(),
             console::style(format!(
-                "DAFT_DASHBOARD_URL=\"http://{}:{}\" python ...",
-                daft_dashboard::DEFAULT_SERVER_ADDR,
-                args.port
+                "DAFT_DASHBOARD_URL=\"http://{}\" python ...",
+                socket_addr
             ))
             .bold(),
         );
         println!(
             "✨ View the dashboard at {}. Press Ctrl+C to shutdown",
-            console::style(format!(
-                "http://{}:{}",
-                daft_dashboard::DEFAULT_SERVER_ADDR,
-                args.port
-            ))
-            .bold()
-            .magenta()
-            .underlined(),
+            console::style(format!("http://{}", socket_addr))
+                .bold()
+                .magenta()
+                .underlined(),
         );
-        daft_dashboard::launch_server(args.port, async move { shutdown_rx.await.unwrap() })
-            .await
-            .expect("Failed to launch dashboard server");
+        daft_dashboard::launch_server(
+            args.addr,
+            args.port,
+            async move { shutdown_rx.await.unwrap() },
+        )
+        .await
+        .expect("Failed to launch dashboard server");
     });
 
     loop {

src/daft-dashboard/src/lib.rs

@@ -5,7 +5,10 @@ pub mod engine;
 pub mod python;
 pub(crate) mod state;
 
-use std::{net::Ipv4Addr, sync::Arc};
+use std::{
+    net::{IpAddr, Ipv4Addr},
+    sync::Arc,
+};
 
 use axum::{
     Json, Router,
@@ -28,6 +31,7 @@ use tracing::Level;
 
 use crate::state::{DashboardState, GLOBAL_DASHBOARD_STATE};
 
+// NOTE(void001): default listen to all ipv4 address, which pose a security risk in production environment
 pub const DEFAULT_SERVER_ADDR: Ipv4Addr = Ipv4Addr::UNSPECIFIED;
 pub const DEFAULT_SERVER_PORT: u16 = 3238;
 
@@ -274,6 +278,7 @@ async fn ping() -> StatusCode {
 }
 
 pub async fn launch_server(
+    addr: IpAddr,
     port: u16,
     shutdown_fn: impl Future<Output = ()> + Send + 'static,
 ) -> std::io::Result<()> {
@@ -303,7 +308,7 @@ pub async fn launch_server(
         .with_state(GLOBAL_DASHBOARD_STATE.clone());
 
     // Start the server
-    let listener = TcpListener::bind((DEFAULT_SERVER_ADDR, port)).await?;
+    let listener = TcpListener::bind((addr, port)).await?;
     axum::serve(listener, app)
         .with_graceful_shutdown(shutdown_fn)
         .await

src/daft-dashboard/src/python.rs

@@ -83,7 +83,7 @@ pub fn launch(noop_if_initialized: bool) -> PyResult<ConnectionHandle> {
         }
     }
 
-    let port = super::DEFAULT_SERVER_PORT; // TODO: Make configurable
+    let port = super::DEFAULT_SERVER_PORT;
     let (send, recv) = oneshot::channel::<()>();
 
     let handle = ConnectionHandle {
@@ -93,7 +93,12 @@ pub fn launch(noop_if_initialized: bool) -> PyResult<ConnectionHandle> {
     let _ = std::thread::spawn(move || {
         DASHBOARD_ENABLED.store(true, Ordering::SeqCst);
         let res = tokio_runtime().block_on(async {
-            super::launch_server(port, async move { recv.await.unwrap() }).await
+            super::launch_server(
+                std::net::IpAddr::V4(super::DEFAULT_SERVER_ADDR),
+                port,
+                async move { recv.await.unwrap() },
+            )
+            .await
         });
         DASHBOARD_ENABLED.store(false, Ordering::SeqCst);
         res