forge/socket/
tcp_stream.rs

1use core::{
2    ffi::c_void,
3    fmt,
4    net::{SocketAddr, SocketAddrV4, SocketAddrV6},
5};
6use log::error;
7use sys::socket::*;
8
9pub struct TcpStream {
10    pub(super) socket: Socket,
11}
12
13pub enum Shutdown {
14    Read = 0,
15    Write = 1,
16    Both = 2,
17}
18
19impl TcpStream {
20    pub fn connect(addr: SocketAddr) -> Option<Self> {
21        match addr {
22            SocketAddr::V4(v4) => Self::connect_v4(v4),
23            SocketAddr::V6(v6) => Self::connect_v6(v6),
24        }
25    }
26
27    pub(super) fn from_socket(socket: Socket) -> Self {
28        Self { socket }
29    }
30
31    pub fn close(&mut self) {
32        unsafe { forge_socket_destroy(&mut self.socket) }
33    }
34
35    pub fn read(&self, buf: &mut [u8]) -> Option<usize> {
36        let n = unsafe { forge_socket_recv(&self.socket, buf.as_mut_ptr() as *mut c_void, buf.len(), 0) };
37        if n < 0 {
38            error!("Failed to receive data");
39            return None;
40        }
41        Some(n as usize)
42    }
43
44    pub fn write(&self, buf: &[u8]) -> Option<usize> {
45        let n = unsafe { forge_socket_send(&self.socket, buf.as_ptr() as *const c_void, buf.len(), 0) };
46        if n < 0 {
47            error!("Failed to send data");
48            return None;
49        }
50        Some(n as usize)
51    }
52
53    pub fn peek(&self, buf: &mut [u8]) -> Option<usize> {
54        let n = unsafe {
55            forge_socket_recv(
56                &self.socket,
57                buf.as_mut_ptr() as *mut c_void,
58                buf.len(),
59                super::consts::MSG_PEEK,
60            )
61        };
62        if n < 0 {
63            error!("Failed to peek data");
64            return None;
65        }
66        Some(n as usize)
67    }
68
69    pub fn shutdown(&self, how: Shutdown) -> bool {
70        let result = unsafe { forge_socket_shutdown(&self.socket, how as i32) };
71        if result == -1 {
72            error!("Failed to shutdown socket");
73        }
74        result == 0
75    }
76
77    pub fn set_nodelay(&self, nodelay: bool) -> bool {
78        let val: i32 = nodelay as i32;
79        let result = unsafe {
80            forge_socket_setSockOpt(
81                &self.socket,
82                super::consts::IPPROTO_TCP,
83                super::consts::TCP_NODELAY,
84                &val as *const i32 as *const c_void,
85                size_of::<i32>() as u32,
86            )
87        };
88        if result == -1 {
89            error!("Failed to set TCP_NODELAY");
90        }
91        result == 0
92    }
93
94    fn connect_v4(addr: SocketAddrV4) -> Option<Self> {
95        let mut socket = unsafe { forge_socket_create(super::consts::AF_INET, super::consts::SOCK_STREAM, 0) };
96        if socket.fd == -1 {
97            error!("Failed to create socket");
98            return None;
99        }
100
101        let in_addr = SockAddrIn {
102            sin_len: 0,
103            sin_family: super::consts::AF_INET as u8,
104            sin_port: addr.port().to_be(),
105            sin_addr: addr.ip().octets(),
106            sin_zero: [0; 8],
107        };
108
109        let result = unsafe { forge_socket_connect(&socket, &in_addr, size_of::<SockAddrIn>()) };
110        if result == -1 {
111            unsafe { forge_socket_destroy(&mut socket) };
112            error!("Failed to connect to address");
113            return None;
114        }
115
116        Some(Self { socket })
117    }
118
119    fn connect_v6(_addr: SocketAddrV6) -> Option<Self> {
120        error!("IPv6 not currently supported");
121        None
122    }
123}
124
125impl Drop for TcpStream {
126    fn drop(&mut self) {
127        unsafe { forge_socket_destroy(&mut self.socket) }
128    }
129}
130
131impl fmt::Write for TcpStream {
132    fn write_str(&mut self, s: &str) -> fmt::Result {
133        match self.write(s.as_bytes()) {
134            Some(_) => Ok(()),
135            None => Err(fmt::Error {}),
136        }
137    }
138}
forge/socket/tcp_stream.rs

forge/socket/
tcp_stream.rs
