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.
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.
38     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.
registernull45     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     /// Deregister the socket
deregisternull61     pub fn deregister(&mut self) -> io::Result<()> {
62         match self.inner.as_mut() {
63             Some(state) => {
64                 {
65                     let mut sock_state = state.state.lock().unwrap();
66                     sock_state.start_drop();
67                 }
68                 self.inner = None;
69                 Ok(())
70             }
71             None => Err(io::ErrorKind::NotFound.into()),
72         }
73     }
74 
75     /// The IO operation does not really report an error when Err(WouldBlock)
76     /// occurs. We need to re-register the current IO operation.
77     pub(crate) fn try_io<T, F, R>(&self, task: F, io: &T) -> io::Result<R>
78     where
79         F: FnOnce(&T) -> io::Result<R>,
80     {
81         let result = task(io);
82         if let Err(ref e) = result {
83             if e.kind() == io::ErrorKind::WouldBlock {
84                 self.inner.as_ref().map_or(Ok(()), |net_inner| {
85                     net_inner.selector.reregister(
86                         net_inner.state.clone(),
87                         net_inner.token,
88                         net_inner.interests,
89                     )
90                 })?;
91             }
92         }
93         result
94     }
95 }
96 
97 /// This structure used to re-register the socket when Err(WouldBlock) occurs
98 #[derive(Clone)]
99 pub(crate) struct NetInner {
100     selector: Arc<SelectorInner>,
101     token: Token,
102     interests: Interest,
103     state: Pin<Arc<Mutex<SockState>>>,
104 }
105 
106 impl NetInner {
107     pub(crate) fn new(
108         selector: Arc<SelectorInner>,
109         token: Token,
110         interests: Interest,
111         state: Pin<Arc<Mutex<SockState>>>,
112     ) -> NetInner {
113         NetInner {
114             selector,
115             token,
116             interests,
117             state,
118         }
119     }
120 }
121 
122 impl Drop for NetInner {
dropnull123     fn drop(&mut self) {
124         let mut sock_state = self.state.lock().unwrap();
125         sock_state.start_drop();
126     }
127 }
128