Simplify open_priv_sock

greenbone · Nov 13, 2024 · b6252f8 · b6252f8
1 parent 157d014
commit b6252f8
Showing 1 changed file with 50 additions and 54 deletions.
diff --git a/rust/src/nasl/builtin/network/socket.rs b/rust/src/nasl/builtin/network/socket.rs
@@ -4,6 +4,7 @@
 
 use std::{
     io::{self, BufRead, Read, Write},
+    net::IpAddr,
     sync::RwLock,
     thread::sleep,
     time::{Duration, SystemTime},
@@ -547,44 +548,48 @@ impl NaslSockets {
         Ok(NaslValue::Number(fd as i64))
     }
 
-    /// Open a privileged socket to the target host.
-    /// It takes three named integer arguments:
-    /// - dport is the destination port
-    /// - sport is the source port, which may be inferior to 1024. This argument is optional.
-    ///   If it is not set, the function will try to open a socket on any port from 1 to 1023.
-    /// - timeout: An integer with the timeout value in seconds.  The default timeout is controlled by a global value.
-    #[nasl_function(named(dport, sport))]
-    fn open_priv_sock_tcp(
+    fn open_priv_sock(
         &self,
-        context: &Context,
+        addr: IpAddr,
         dport: i64,
         sport: Option<i64>,
+        tcp: bool,
     ) -> Result<NaslValue, FunctionErrorKind> {
         let dport = verify_port(dport)?;
 
-        let addr = ipstr2ipaddr(context.target())?;
-
-        // TODO: set timeout to global recv timeout when available
-        let timeout = Duration::from_secs(10);
-
         if let Some(sport) = sport {
             let sport = verify_port(sport)?;
-            self.wait_before_next_probe();
-            let tcp = TcpConnection::connect_priv(addr, sport, dport, timeout)?;
-
-            let fd = self.add(NaslSocket::Tcp(Box::new(tcp)));
+            let fd = if tcp {
+                // TODO: set timeout to global recv timeout when available
+                let timeout = Duration::from_secs(10);
+                self.wait_before_next_probe();
+                let tcp = TcpConnection::connect_priv(addr, sport, dport, timeout)?;
+                self.add(NaslSocket::Tcp(Box::new(tcp)))
+            } else {
+                let udp = UdpConnection::new_priv(addr, sport, dport)?;
+                self.add(NaslSocket::Udp(udp))
+            };
             return Ok(NaslValue::Number(fd as i64));
         }
 
-        let mut sport = 1023;
-
-        while sport > 0 {
-            self.wait_before_next_probe();
-            if let Ok(tcp) = TcpConnection::connect_priv(addr, sport, dport, timeout) {
-                let fd = self.add(NaslSocket::Tcp(Box::new(tcp)));
-                return Ok(NaslValue::Number(fd as i64));
-            }
-            sport -= 1;
+        for sport in (1..=1023).rev() {
+            let fd = if tcp {
+                // TODO: set timeout to global recv timeout when available
+                let timeout = Duration::from_secs(10);
+                self.wait_before_next_probe();
+                if let Ok(tcp) = TcpConnection::connect_priv(addr, sport, dport, timeout) {
+                    self.add(NaslSocket::Tcp(Box::new(tcp)))
+                } else {
+                    continue;
+                }
+            } else {
+                if let Ok(udp) = UdpConnection::new_priv(addr, sport, dport) {
+                    self.add(NaslSocket::Udp(udp))
+                } else {
+                    continue;
+                }
+            };
+            return Ok(NaslValue::Number(fd as i64));
         }
         Err(FunctionErrorKind::Diagnostic(
             format!(
@@ -595,6 +600,23 @@ impl NaslSockets {
         ))
     }
 
+    /// Open a privileged socket to the target host.
+    /// It takes three named integer arguments:
+    /// - dport is the destination port
+    /// - sport is the source port, which may be inferior to 1024. This argument is optional.
+    ///   If it is not set, the function will try to open a socket on any port from 1 to 1023.
+    /// - timeout: An integer with the timeout value in seconds.  The default timeout is controlled by a global value.
+    #[nasl_function(named(dport, sport))]
+    fn open_priv_sock_tcp(
+        &self,
+        context: &Context,
+        dport: i64,
+        sport: Option<i64>,
+    ) -> Result<NaslValue, FunctionErrorKind> {
+        let addr = ipstr2ipaddr(context.target())?;
+        self.open_priv_sock(addr, dport, sport, true)
+    }
+
     /// Open a privileged UDP socket to the target host.
     /// It takes three named integer arguments:
     /// - dport is the destination port
@@ -607,34 +629,8 @@ impl NaslSockets {
         dport: i64,
         sport: Option<i64>,
     ) -> Result<NaslValue, FunctionErrorKind> {
-        let dport = verify_port(dport)?;
-
         let addr = ipstr2ipaddr(context.target())?;
-
-        if let Some(sport) = sport {
-            let sport = verify_port(sport)?;
-            let udp = UdpConnection::new_priv(addr, sport, dport)?;
-
-            let fd = self.add(NaslSocket::Udp(udp));
-            return Ok(NaslValue::Number(fd as i64));
-        }
-
-        let mut sport = 1023;
-
-        while sport > 0 {
-            if let Ok(udp) = UdpConnection::new_priv(addr, sport, dport) {
-                let fd = self.add(NaslSocket::Udp(udp));
-                return Ok(NaslValue::Number(fd as i64));
-            }
-            sport -= 1;
-        }
-        Err(FunctionErrorKind::Diagnostic(
-            format!(
-                "Unable to open priv socket to {} on any socket from 1-1023",
-                addr
-            ),
-            None,
-        ))
+        self.open_priv_sock(addr, dport, sport, false)
     }
 
     /// Get the source port of a open socket