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::os::windows::io::RawSocket;
15 use std::pin::Pin;
16 use std::sync::{Arc, Mutex, Once};
17 use std::{io, net};
18
19 use crate::sys::windows::selector::{SelectorInner, SockState};
20 use crate::{Interest, Selector, Token};
21
22 /// Initialise the network stack for Windows.
init()23 pub(crate) fn init() {
24 static INIT: Once = Once::new();
25 INIT.call_once(|| {
26 drop(net::UdpSocket::bind("127.0.0.1:0"));
27 });
28 }
29
30 #[derive(Clone)]
31 pub(crate) struct NetState {
32 /// State is None if the socket has not been Registered.
33 inner: Option<Box<NetInner>>,
34 }
35
36 impl NetState {
37 /// Creates a new `NetState` with None.
new() -> NetState38 pub(crate) fn new() -> NetState {
39 NetState { inner: None }
40 }
41
42 /// Register the socket to [`Selector`]
43 /// If inner is Some, this function returns Err(AlreadyExists).
44 /// If register success, Set the inner to Some.
register( &mut self, selector: &Selector, token: Token, interests: Interest, socket: RawSocket, ) -> io::Result<()>45 pub fn register(
46 &mut self,
47 selector: &Selector,
48 token: Token,
49 interests: Interest,
50 socket: RawSocket,
51 ) -> io::Result<()> {
52 match self.inner {
53 Some(_) => Err(io::ErrorKind::AlreadyExists.into()),
54 None => selector.register(socket, token, interests).map(|state| {
55 self.inner = Some(Box::new(state));
56 }),
57 }
58 }
59
60 /// Reregister the socket
reregister( &mut self, selector: &Selector, token: Token, interests: Interest, ) -> io::Result<()>61 pub fn reregister(
62 &mut self,
63 selector: &Selector,
64 token: Token,
65 interests: Interest,
66 ) -> io::Result<()> {
67 match self.inner.as_mut() {
68 Some(state) => selector
69 .reregister(state.state.clone(), token, interests)
70 .map(|_| {
71 state.token = token;
72 state.interests = interests;
73 }),
74 None => Err(io::ErrorKind::NotFound.into()),
75 }
76 }
77
78 /// Deregister the socket
deregister(&mut self) -> io::Result<()>79 pub fn deregister(&mut self) -> io::Result<()> {
80 match self.inner.as_mut() {
81 Some(state) => {
82 {
83 let mut sock_state = state.state.lock().unwrap();
84 sock_state.start_drop();
85 }
86 self.inner = None;
87 Ok(())
88 }
89 None => Err(io::ErrorKind::NotFound.into()),
90 }
91 }
92
93 /// The IO operation does not really report an error when Err(WouldBlock)
94 /// occurs. We need to re-register the current IO operation.
try_io<T, F, R>(&self, task: F, io: &T) -> io::Result<R> where F: FnOnce(&T) -> io::Result<R>,95 pub(crate) fn try_io<T, F, R>(&self, task: F, io: &T) -> io::Result<R>
96 where
97 F: FnOnce(&T) -> io::Result<R>,
98 {
99 let result = task(io);
100 if let Err(ref e) = result {
101 if e.kind() == io::ErrorKind::WouldBlock {
102 self.inner.as_ref().map_or(Ok(()), |net_inner| {
103 net_inner.selector.reregister(
104 net_inner.state.clone(),
105 net_inner.token,
106 net_inner.interests,
107 )
108 })?;
109 }
110 }
111 result
112 }
113 }
114
115 /// This structure used to re-register the socket when Err(WouldBlock) occurs
116 #[derive(Clone)]
117 pub(crate) struct NetInner {
118 selector: Arc<SelectorInner>,
119 token: Token,
120 interests: Interest,
121 state: Pin<Arc<Mutex<SockState>>>,
122 }
123
124 impl NetInner {
new( selector: Arc<SelectorInner>, token: Token, interests: Interest, state: Pin<Arc<Mutex<SockState>>>, ) -> NetInner125 pub(crate) fn new(
126 selector: Arc<SelectorInner>,
127 token: Token,
128 interests: Interest,
129 state: Pin<Arc<Mutex<SockState>>>,
130 ) -> NetInner {
131 NetInner {
132 selector,
133 token,
134 interests,
135 state,
136 }
137 }
138 }
139
140 impl Drop for NetInner {
drop(&mut self)141 fn drop(&mut self) {
142 let mut sock_state = self.state.lock().unwrap();
143 sock_state.start_drop();
144 }
145 }
146