Address PR feedback

rust/src/cli/l2cap/server_bridge.rs

@@ -0,0 +1,205 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// L2CAP CoC server bridge: waits for a peer to connect an L2CAP CoC channel
+/// on a specified PSM. When the connection is made, the bridge connects a TCP
+/// socket to a remote host and bridges the data in both directions, with flow
+/// control.
+/// When the L2CAP CoC channel is closed, the bridge disconnects the TCP socket
+/// and waits for a new L2CAP CoC channel to be connected.
+/// When the TCP connection is closed by the TCP server, the L2CAP connection is closed as well.
+use crate::cli::l2cap::{
+    proxy_l2cap_rx_to_tcp_tx, proxy_tcp_rx_to_l2cap_tx, run_future_with_current_task_locals,
+    BridgeData,
+};
+use bumble::wrapper::{device::Device, hci::HciConstant, l2cap::LeConnectionOrientedChannel};
+use futures::executor::block_on;
+use owo_colors::OwoColorize;
+use pyo3::{PyResult, Python};
+use std::{sync::Arc, time::Duration};
+use tokio::{
+    join,
+    net::TcpStream,
+    select,
+    sync::{mpsc, Mutex},
+};
+
+pub struct Args {
+    pub psm: u16,
+    pub max_credits: Option<u16>,
+    pub mtu: Option<u16>,
+    pub mps: Option<u16>,
+    pub tcp_host: String,
+    pub tcp_port: u16,
+}
+
+pub async fn start(args: &Args, device: &mut Device) -> PyResult<()> {
+    let host = args.tcp_host.clone();
+    let port = args.tcp_port;
+    device.register_l2cap_channel_server(
+        args.psm,
+        move |_py, l2cap_channel| {
+            let channel_info = l2cap_channel
+                .debug_string()
+                .unwrap_or_else(|e| format!("failed to get l2cap channel info ({e})"));
+            println!("{} {channel_info}", "*** L2CAP channel:".cyan());
+
+            let host = host.clone();
+            // Ensure Python event loop is available to l2cap `disconnect`
+            let _ = run_future_with_current_task_locals(proxy_data_between_l2cap_and_tcp(
+                l2cap_channel,
+                host,
+                port,
+            ));
+            Ok(())
+        },
+        args.max_credits,
+        args.mtu,
+        args.mps,
+    )?;
+
+    println!(
+        "{}",
+        format!("### Listening for CoC connection on PSM {}", args.psm).yellow()
+    );
+
+    device.on_connection(|_py, mut connection| {
+        let connection_info = connection
+            .debug_string()
+            .unwrap_or_else(|e| format!("failed to get connection info ({e})"));
+        println!(
+            "{} {}",
+            "@@@ Bluetooth connection: ".green(),
+            connection_info,
+        );
+        connection.on_disconnection(|_py, reason| {
+            let disconnection_info = match HciConstant::error_name(reason) {
+                Ok(info_string) => info_string,
+                Err(py_err) => format!("failed to get disconnection error name ({})", py_err),
+            };
+            println!(
+                "{} {}",
+                "@@@ Bluetooth disconnection: ".red(),
+                disconnection_info,
+            );
+            Ok(())
+        })?;
+        Ok(())
+    })?;
+
+    device.start_advertising(false).await?;
+
+    Ok(())
+}
+
+async fn proxy_data_between_l2cap_and_tcp(
+    mut l2cap_channel: LeConnectionOrientedChannel,
+    tcp_host: String,
+    tcp_port: u16,
+) -> PyResult<()> {
+    let (l2cap_to_tcp_tx, mut l2cap_to_tcp_rx) = mpsc::channel::<BridgeData>(10);
+
+    // Set callback (`set_sink`) for when l2cap data is received.
+    let l2cap_to_tcp_tx_clone = l2cap_to_tcp_tx.clone();
+    l2cap_channel
+        .set_sink(move |_py, sdu| {
+            block_on(l2cap_to_tcp_tx_clone.send(BridgeData::Data(sdu.into())))
+                .expect("failed to channel data to tcp");
+            Ok(())
+        })
+        .expect("failed to set sink for l2cap connection");
+
+    // Set l2cap callback for when the channel is closed.
+    l2cap_channel
+        .on_close(move |_py| {
+            println!("{}", "*** L2CAP channel closed".red());
+            block_on(l2cap_to_tcp_tx.send(BridgeData::CloseSignal))
+                .expect("failed to channel close signal to tcp");
+            Ok(())
+        })
+        .expect("failed to set on_close callback for l2cap channel");
+
+    println!(
+        "{}",
+        format!("### Connecting to TCP {tcp_host}:{tcp_port}...").yellow()
+    );
+
+    let l2cap_channel = Arc::new(Mutex::new(Some(l2cap_channel)));
+    let tcp_stream = match TcpStream::connect(format!("{tcp_host}:{tcp_port}")).await {
+        Ok(stream) => {
+            println!("{}", "### Connected".green());
+            Some(stream)
+        }
+        Err(err) => {
+            println!("{}", format!("!!! Connection failed: {err}").red());
+            if let Some(mut channel) = l2cap_channel.lock().await.take() {
+                // Bumble might enter an invalid state if disconnection request is received from
+                // l2cap client before receiving a disconnection response from the same client,
+                // blocking this async call from returning.
+                // See: https://github.com/google/bumble/issues/257
+                select! {
+                    res = channel.disconnect() => {
+                        let _ = res.map_err(|e| eprintln!("Failed to call disconnect on l2cap channel: {e}"));
+                    },
+                    _ = tokio::time::sleep(Duration::from_secs(1)) => eprintln!("Timed out while calling disconnect on l2cap channel."),
+                }
+            }
+            None
+        }
+    };
+
+    match tcp_stream {
+        None => {
+            while let Some(bridge_data) = l2cap_to_tcp_rx.recv().await {
+                match bridge_data {
+                    BridgeData::Data(sdu) => {
+                        println!("{}", format!("<<< [L2CAP SDU]: {} bytes", sdu.len()).cyan());
+                        println!("{}", "!!! TCP socket not open, dropping".red())
+                    }
+                    BridgeData::CloseSignal => break,
+                }
+            }
+        }
+        Some(tcp_stream) => {
+            let (tcp_reader, tcp_writer) = tcp_stream.into_split();
+
+            // Do tcp stuff when something happens on the l2cap channel.
+            let handle_l2cap_data_future =
+                proxy_l2cap_rx_to_tcp_tx(l2cap_to_tcp_rx, tcp_writer, l2cap_channel.clone());
+
+            // Do l2cap stuff when something happens on tcp.
+            let handle_tcp_data_future =
+                proxy_tcp_rx_to_l2cap_tx(tcp_reader, l2cap_channel.clone(), false);
+
+            let (handle_l2cap_result, handle_tcp_result) =
+                join!(handle_l2cap_data_future, handle_tcp_data_future);
+
+            if let Err(e) = handle_l2cap_result {
+                println!("!!! Error: {e}");
+            }
+
+            if let Err(e) = handle_tcp_result {
+                println!("!!! Error: {e}");
+            }
+        }
+    };
+
+    Python::with_gil(|_| {
+        // Must hold GIL at least once while/after dropping for Python heap object to ensure
+        // de-allocation.
+        drop(l2cap_channel);
+    });
+
+    Ok(())
+}