From 6c2f411d00fc62bea908cca95b75eb143b20bbb8 Mon Sep 17 00:00:00 2001
From: LoveSy <shana@zju.edu.cn>
Date: Thu, 23 May 2024 00:05:20 +0800
Subject: [PATCH] Bind interface for udp proxy

---
 Cargo.toml   |  3 +--
 src/main.rs  | 14 +++++++++-----
 src/proxy.rs | 27 +++++++++++++++++++++++----
 3 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index bea0e25..a514471 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -31,12 +31,11 @@ futures-util = "0.3"
 retina = { git = "https://github.com/yujincheng08/retina.git", branch = "iptv", default-features = false }
 rtp-rs = "0.6"
 tokio-util = { version = "0.7.0", features = ["codec", "net"] }
+local-ip-address = "0.6"
 
 [target.'cfg(target_os = "windows")'.dependencies]
 socket2 = "0.5"
 
-[target.'cfg(not(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))'.dependencies]
-local-ip-address = "0.6"
 
 [features]
 http2 = ["reqwest/http2"]
diff --git a/src/main.rs b/src/main.rs
index 5b8a495..1b7e12c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -193,7 +193,7 @@ async fn logo(args: Data<Args>, path: Path<String>) -> impl Responder {
     debug!("Get logo");
     match get_icon(&args, &path).await {
         Ok(icon) => HttpResponse::Ok().content_type("image/png").body(icon),
-        Err(e) => HttpResponse::NotFound().body(format!("Error getting channels: {}", e))
+        Err(e) => HttpResponse::NotFound().body(format!("Error getting channels: {}", e)),
     }
 }
 
@@ -205,7 +205,9 @@ async fn playlist(args: Data<Args>, req: HttpRequest) -> impl Responder {
     match get_channels(&args, false, &scheme, &host).await {
         Err(e) => {
             if let Some(old_playlist) = OLD_PLAYLIST.try_lock().ok().and_then(|f| f.to_owned()) {
-                HttpResponse::Ok().content_type("application/vnd.apple.mpegurl").body(old_playlist)
+                HttpResponse::Ok()
+                    .content_type("application/vnd.apple.mpegurl")
+                    .body(old_playlist)
             } else {
                 HttpResponse::InternalServerError().body(format!("Error getting channels: {}", e))
             }
@@ -238,7 +240,9 @@ async fn playlist(args: Data<Args>, req: HttpRequest) -> impl Responder {
             if let Ok(mut old_playlist) = OLD_PLAYLIST.try_lock() {
                 *old_playlist = Some(playlist.clone());
             }
-            HttpResponse::Ok().content_type("application/vnd.apple.mpegurl").body(playlist)
+            HttpResponse::Ok()
+                .content_type("application/vnd.apple.mpegurl")
+                .body(playlist)
         }
     }
 }
@@ -261,13 +265,13 @@ async fn rtsp(
 }
 
 #[get("/udp/{addr}")]
-async fn udp(addr: Path<String>) -> impl Responder {
+async fn udp(args: Data<Args>, addr: Path<String>) -> impl Responder {
     let addr = &*addr;
     let addr = match SocketAddrV4::from_str(addr) {
         Ok(addr) => addr,
         Err(e) => return HttpResponse::BadRequest().body(format!("Error: {}", e)),
     };
-    HttpResponse::Ok().streaming(proxy::udp(addr))
+    HttpResponse::Ok().streaming(proxy::udp(addr, args.interface.clone()))
 }
 
 fn usage(cmd: &str) -> std::io::Result<()> {
diff --git a/src/proxy.rs b/src/proxy.rs
index 6e5ec1a..c4bbe98 100644
--- a/src/proxy.rs
+++ b/src/proxy.rs
@@ -1,11 +1,10 @@
-use std::net::{Ipv4Addr, SocketAddrV4};
+use std::net::{IpAddr, Ipv4Addr, SocketAddrV4};
 
 use actix_web::web::Bytes;
 use anyhow::Result;
 use async_stream::stream;
 use futures_core::stream::Stream;
 use futures_util::stream::StreamExt;
-#[cfg(not(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))]
 use local_ip_address::list_afinet_netifas;
 use log::{error, info};
 use reqwest::Url;
@@ -91,7 +90,10 @@ pub(crate) fn rtsp(url: String, if_name: Option<String>) -> impl Stream<Item = R
     }
 }
 
-pub(crate) fn udp(multi_addr: SocketAddrV4) -> impl Stream<Item = Result<Bytes>> {
+pub(crate) fn udp(
+    multi_addr: SocketAddrV4,
+    if_name: Option<String>,
+) -> impl Stream<Item = Result<Bytes>> {
     stream! {
         #[cfg(target_os = "windows")]
         let socket =  {
@@ -108,11 +110,28 @@ pub(crate) fn udp(multi_addr: SocketAddrV4) -> impl Stream<Item = Result<Bytes>>
         let socket = {
             UdpSocket::bind(multi_addr).await?
         };
+
+        let mut interface = Ipv4Addr::new(0, 0, 0, 0);
+        if let Some(ref i) = if_name {
+            use log::debug;
+            let network_interfaces = list_afinet_netifas()?;
+            for (name, ip) in network_interfaces.iter() {
+                debug!("{}: {}", name, ip);
+                if name != i {
+                    continue;
+                }
+                if let IpAddr::V4(ip) = ip {
+                    interface = *ip;
+                    break;
+                }
+            }
+        }
+
         socket.set_multicast_loop_v4(true)?;
 
         socket.join_multicast_v4(
             *multi_addr.ip(),
-            Ipv4Addr::new(0, 0, 0, 0),
+            interface,
         )?;
 
         info!("Udp proxy joined {}", multi_addr);