Use connected UDP sockets

gufolabs · Jan 10, 2024 · 053aaac · 053aaac
1 parent 5df3ce9
commit 053aaac
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 7 deletions.
diff --git a/src/error.rs b/src/error.rs
@@ -7,7 +7,9 @@
 
 use pyo3::{
     create_exception,
-    exceptions::{PyBlockingIOError, PyException, PyNotImplementedError, PyOSError},
+    exceptions::{
+        PyBlockingIOError, PyException, PyNotImplementedError, PyOSError, PyTimeoutError,
+    },
     PyErr,
 };
 
@@ -41,6 +43,8 @@ pub enum SnmpError {
     SocketError(String),
     /// Blocking operation
     WouldBlock,
+    /// Connection refused
+    ConnectionRefused,
 }
 
 impl From<nom::Err<SnmpError>> for SnmpError {
@@ -105,6 +109,7 @@ impl From<SnmpError> for PyErr {
             SnmpError::NoSuchInstance => PyNoSuchInstance::new_err("no such instance"),
             SnmpError::WouldBlock => PyBlockingIOError::new_err("blocked"),
             SnmpError::SocketError(x) => PyOSError::new_err(x),
+            SnmpError::ConnectionRefused => PyTimeoutError::new_err("connection refused"),
         }
     }
 }
diff --git a/src/socket/transport.rs b/src/socket/transport.rs
@@ -8,13 +8,12 @@
 use crate::ber::BerEncoder;
 use crate::buf::Buffer;
 use crate::error::SnmpError;
-use socket2::{Domain, Protocol, SockAddr, Socket, Type};
+use socket2::{Domain, Protocol, Socket, Type};
 use std::net::SocketAddr;
 use std::os::fd::{AsRawFd, RawFd};
 
 pub struct SnmpTransport {
     io: Socket,
-    addr: SockAddr,
     buf: Buffer,
 }
 
@@ -52,10 +51,12 @@ impl SnmpTransport {
         if recv_buffer_size > 0 {
             Self::set_recv_buffer_size(&io, recv_buffer_size)?;
         }
+        // Make socket connected
+        io.connect(&sock_addr.into())
+            .map_err(|e| SnmpError::SocketError(e.to_string()))?;
         //
         Ok(Self {
             io,
-            addr: sock_addr.into(),
             buf: Buffer::default(),
         })
     }
@@ -68,7 +69,7 @@ impl SnmpTransport {
         msg.push_ber(&mut self.buf)?;
         // Send
         self.io
-            .send_to(self.buf.data(), &self.addr)
+            .send(self.buf.data())
             .map_err(|e| SnmpError::SocketError(e.to_string()))?;
         Ok(())
     }
@@ -78,11 +79,14 @@ impl SnmpTransport {
         T: TryFrom<&'a [u8], Error = SnmpError>,
         'b: 'a,
     {
-        let size = match self.io.recv_from(self.buf.as_mut()) {
-            Ok((s, _)) => s,
+        let size = match self.io.recv(self.buf.as_mut()) {
+            Ok(s) => s,
             Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                 return Err(SnmpError::WouldBlock)
             }
+            Err(e) if e.kind() == std::io::ErrorKind::ConnectionRefused => {
+                return Err(SnmpError::ConnectionRefused)
+            }
             Err(e) => return Err(SnmpError::SocketError(e.to_string())),
         };
         // Parse response