raw_socket/tokio/
socket.rs

1// Copyright (C) 2020 - Will Glozer. All rights reserved.
2
3use std::io::{IoSlice, IoSliceMut, Result};
4use std::net::{SocketAddr, ToSocketAddrs};
5use tokio::io::unix::AsyncFd;
6use crate::{Domain, Type, Protocol};
7use crate::option::{Level, Name, Opt};
8
9pub struct RawSocket {
10    io: AsyncFd<crate::RawSocket>,
11}
12
13impl RawSocket {
14    pub fn new(domain: Domain, kind: Type, protocol: Option<Protocol>) -> Result<Self> {
15        let sys = crate::RawSocket::new(domain, kind, protocol)?;
16        sys.set_nonblocking(true)?;
17        let io  = AsyncFd::new(sys)?;
18        Ok(RawSocket { io })
19    }
20
21    pub async fn bind<A: ToSocketAddrs>(&self, addr: A) -> Result<()> {
22        self.io.get_ref().bind(addr)
23    }
24
25    pub fn local_addr(&self) -> Result<SocketAddr> {
26        self.io.get_ref().local_addr()
27    }
28
29    pub async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, SocketAddr)> {
30        self.read(|s| s.recv_from(buf)).await
31    }
32
33    pub async fn recv_msg(
34        &self,
35        data: &[IoSliceMut<'_>],
36        ctrl: Option<&mut [u8]>
37    ) -> Result<(usize, SocketAddr)> {
38        let ctrl = ctrl.unwrap_or(&mut []);
39        self.read(|s| s.recv_msg(data, ctrl)).await
40    }
41
42    pub async fn send_to<A: ToSocketAddrs>(&self, buf: &[u8], addr: A) -> Result<usize> {
43        self.write(|s| s.send_to(buf, &addr)).await
44    }
45
46    pub async fn send_msg<A: ToSocketAddrs>(
47        &self,
48        addr: A,
49        data: &[IoSlice<'_>],
50        ctrl: Option<&[u8]>
51    ) -> Result<usize> {
52        let ctrl = ctrl.unwrap_or(&[]);
53        self.write(|s| s.send_msg(&addr, data, ctrl)).await
54    }
55
56    pub fn get_sockopt<O: Opt>(&self, level: Level, name: Name) -> Result<O> {
57        self.io.get_ref().get_sockopt(level, name)
58    }
59
60    pub fn set_sockopt<O: Opt>(&self, level: Level, name: Name, value: &O) -> Result<()> {
61        self.io.get_ref().set_sockopt(level, name, value)
62    }
63
64    async fn read<F: FnMut(&crate::RawSocket) -> Result<R>, R>(&self, mut f: F) -> Result<R> {
65        loop {
66            let mut guard = self.io.readable().await?;
67            match guard.try_io(|inner| f(inner.get_ref())) {
68                Ok(r)  => return r,
69                Err(_) => continue,
70            }
71        }
72    }
73
74    async fn write<F: FnMut(&crate::RawSocket) -> Result<R>, R>(&self, mut f: F) -> Result<R> {
75        loop {
76            let mut guard = self.io.writable().await?;
77            match guard.try_io(|inner| f(inner.get_ref())) {
78                Ok(r)  => return r,
79                Err(_) => continue,
80            }
81        }
82    }
83}