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 core::fmt::Debug;
15 use core::pin::Pin;
16 use core::task::{Context, Poll};
17 use std::io::{self, Read, Write};
18
19 use crate::runtime::{AsyncRead, AsyncWrite, ReadBuf};
20
21 #[derive(Debug)]
22 pub(crate) struct Wrapper<S> {
23 pub(crate) stream: S,
24 // Context of stream.
25 pub(crate) context: *mut (),
26 }
27
28 impl<S> Wrapper<S> {
29 /// Gets inner `Stream` and `Context` of `Stream`.
30 ///
31 /// # SAFETY
32 /// Must be called with `context` set to a valid pointer to a live `Context`
33 /// object, and the wrapper must be pinned in memory.
inner(&mut self) -> (Pin<&mut S>, &mut Context<'_>)34 unsafe fn inner(&mut self) -> (Pin<&mut S>, &mut Context<'_>) {
35 debug_assert!(!self.context.is_null());
36 let stream = Pin::new_unchecked(&mut self.stream);
37 let context = &mut *(self.context as *mut _);
38 (stream, context)
39 }
40 }
41
42 impl<S> Read for Wrapper<S>
43 where
44 S: AsyncRead,
45 {
read(&mut self, buf: &mut [u8]) -> io::Result<usize>46 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
47 let (stream, cx) = unsafe { self.inner() };
48 let mut buf = ReadBuf::new(buf);
49 match stream.poll_read(cx, &mut buf)? {
50 Poll::Ready(()) => Ok(buf.filled().len()),
51 Poll::Pending => Err(io::Error::from(io::ErrorKind::WouldBlock)),
52 }
53 }
54 }
55
56 impl<S> Write for Wrapper<S>
57 where
58 S: AsyncWrite,
59 {
write(&mut self, buf: &[u8]) -> io::Result<usize>60 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
61 let (stream, cx) = unsafe { self.inner() };
62 match stream.poll_write(cx, buf) {
63 Poll::Ready(r) => r,
64 Poll::Pending => Err(io::Error::from(io::ErrorKind::WouldBlock)),
65 }
66 }
67
flush(&mut self) -> io::Result<()>68 fn flush(&mut self) -> io::Result<()> {
69 let (stream, cx) = unsafe { self.inner() };
70 match stream.poll_flush(cx) {
71 Poll::Ready(r) => r,
72 Poll::Pending => Err(io::Error::from(io::ErrorKind::WouldBlock)),
73 }
74 }
75 }
76
77 // *mut () is not impl Send or Sync.
78 unsafe impl<S: Send> Send for Wrapper<S> {}
79 unsafe impl<S: Sync> Sync for Wrapper<S> {}
80
81 /// Checks `io::Result`.
check_io_to_poll<T>(r: io::Result<T>) -> Poll<io::Result<T>>82 pub(crate) fn check_io_to_poll<T>(r: io::Result<T>) -> Poll<io::Result<T>> {
83 match r {
84 Ok(t) => Poll::Ready(Ok(t)),
85 Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Poll::Pending,
86 Err(e) => Poll::Ready(Err(e)),
87 }
88 }
89