1 // Copyright (c) 2023 Huawei Device Co., Ltd. 2 // Licensed under the Apache License, Version 2.0 (the "License"); 3 // you may not use this file except in compliance with the License. 4 // You may obtain a copy of the License at 5 // 6 // http://www.apache.org/licenses/LICENSE-2.0 7 // 8 // Unless required by applicable law or agreed to in writing, software 9 // distributed under the License is distributed on an "AS IS" BASIS, 10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 11 // See the License for the specific language governing permissions and 12 // limitations under the License. 13 14 use std::net::SocketAddr; 15 use std::os::windows::io::{AsRawSocket, FromRawSocket, RawSocket}; 16 use std::{io, mem, net}; 17 18 use windows_sys::Win32::Networking::WinSock::{ 19 self, closesocket, ioctlsocket, socket, ADDRESS_FAMILY, AF_INET, AF_INET6, FIONBIO, 20 INVALID_SOCKET, SOCKET, SOCKET_ERROR, SOCK_STREAM, 21 }; 22 23 use crate::sys::windows::net::init; 24 use crate::sys::windows::socket_addr::socket_addr_trans; 25 26 pub(crate) struct TcpSocket { 27 socket: SOCKET, 28 } 29 30 impl TcpSocket { 31 /// Gets new socket new_socket(addr: SocketAddr) -> io::Result<TcpSocket>32 pub(crate) fn new_socket(addr: SocketAddr) -> io::Result<TcpSocket> { 33 if addr.is_ipv4() { 34 Self::create_socket(AF_INET, SOCK_STREAM) 35 } else { 36 Self::create_socket(AF_INET6, SOCK_STREAM) 37 } 38 } 39 create_socket(domain: ADDRESS_FAMILY, socket_type: u16) -> io::Result<TcpSocket>40 fn create_socket(domain: ADDRESS_FAMILY, socket_type: u16) -> io::Result<TcpSocket> { 41 init(); 42 43 let socket = socket_syscall!( 44 socket(domain as i32, socket_type as i32, 0), 45 PartialEq::eq, 46 INVALID_SOCKET 47 )?; 48 49 match socket_syscall!(ioctlsocket(socket, FIONBIO, &mut 1), PartialEq::ne, 0) { 50 Err(err) => { 51 let _ = unsafe { closesocket(socket) }; 52 Err(err) 53 } 54 Ok(_) => Ok(TcpSocket { 55 socket: socket as SOCKET, 56 }), 57 } 58 } 59 60 /// System call to bind Socket. bind(&self, addr: SocketAddr) -> io::Result<()>61 pub(crate) fn bind(&self, addr: SocketAddr) -> io::Result<()> { 62 use WinSock::bind; 63 64 let (raw_addr, raw_addr_length) = socket_addr_trans(&addr); 65 socket_syscall!( 66 bind(self.socket as _, raw_addr.as_ptr(), raw_addr_length), 67 PartialEq::eq, 68 SOCKET_ERROR 69 )?; 70 Ok(()) 71 } 72 73 /// System call to listen. listen(self, backlog: u32) -> io::Result<()>74 pub(crate) fn listen(self, backlog: u32) -> io::Result<()> { 75 use std::convert::TryInto; 76 77 use WinSock::listen; 78 79 let backlog = backlog.try_into().unwrap_or(i32::MAX); 80 socket_syscall!( 81 listen(self.socket as _, backlog), 82 PartialEq::eq, 83 SOCKET_ERROR 84 )?; 85 mem::forget(self); 86 Ok(()) 87 } 88 89 /// System call to connect. connect(self, addr: SocketAddr) -> io::Result<()>90 pub(crate) fn connect(self, addr: SocketAddr) -> io::Result<()> { 91 use WinSock::connect; 92 93 let (socket_addr, socket_addr_length) = socket_addr_trans(&addr); 94 let res = socket_syscall!( 95 connect(self.socket as _, socket_addr.as_ptr(), socket_addr_length), 96 PartialEq::eq, 97 SOCKET_ERROR 98 ); 99 100 match res { 101 Err(e) if e.kind() != io::ErrorKind::WouldBlock => Err(e), 102 _ => { 103 mem::forget(self); 104 Ok(()) 105 } 106 } 107 } 108 109 /// Closes Socket close(&self)110 pub(crate) fn close(&self) { 111 let _ = unsafe { net::TcpStream::from_raw_socket(self.socket as RawSocket) }; 112 } 113 } 114 115 impl AsRawSocket for TcpSocket { as_raw_socket(&self) -> RawSocket116 fn as_raw_socket(&self) -> RawSocket { 117 self.socket as RawSocket 118 } 119 } 120 121 impl FromRawSocket for TcpSocket { from_raw_socket(sock: RawSocket) -> Self122 unsafe fn from_raw_socket(sock: RawSocket) -> Self { 123 TcpSocket { 124 socket: sock as SOCKET, 125 } 126 } 127 } 128 129 impl Drop for TcpSocket { drop(&mut self)130 fn drop(&mut self) { 131 self.close(); 132 } 133 } 134