Move ssl-related PySocket code to ssl module

youknowone · youknowone · commit b8d8930b9ffc · 2021-09-27T00:38:38.000+09:00
diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs
@@ -5,7 +5,7 @@ use crate::{
     exceptions::{IntoPyException, PyBaseExceptionRef},
     function::{FuncArgs, OptionalArg, OptionalOption},
     utils::{Either, ToCString},
-    IntoPyObject, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TryFromBorrowedObject,
+    IntoPyObject, PyClassImpl, PyObjectRef, PyResult, PyValue, TryFromBorrowedObject,
     TryFromObject, TypeProtocol, VirtualMachine,
 };
 use crossbeam_utils::atomic::AtomicCell;
@@ -158,8 +158,6 @@ impl Default for PySocket {
     }
 }
 
-pub type PySocketRef = PyRef<PySocket>;
-
 #[cfg(windows)]
 const CLOSED_ERR: i32 = c::WSAENOTSOCK;
 #[cfg(unix)]
@@ -170,7 +168,7 @@ impl PySocket {
         PyRwLockReadGuard::try_map(self.sock.read(), |sock| sock.get()).ok()
     }
 
-    fn sock_io(&self) -> io::Result<PyMappedRwLockReadGuard<'_, Socket>> {
+    pub(crate) fn sock_io(&self) -> io::Result<PyMappedRwLockReadGuard<'_, Socket>> {
         self.sock_opt()
             .ok_or_else(|| io::Error::from_raw_os_error(CLOSED_ERR))
     }
@@ -919,20 +917,6 @@ impl PySocket {
     }
 }
 
-impl io::Read for PySocketRef {
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        <&Socket as io::Read>::read(&mut &*self.sock_io()?, buf)
-    }
-}
-impl io::Write for PySocketRef {
-    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        <&Socket as io::Write>::write(&mut &*self.sock_io()?, buf)
-    }
-    fn flush(&mut self) -> io::Result<()> {
-        <&Socket as io::Write>::flush(&mut &*self.sock_io()?)
-    }
-}
-
 struct Address {
     host: PyStrRef,
     port: u16,
diff --git a/vm/src/stdlib/ssl.rs b/vm/src/stdlib/ssl.rs
@@ -1,4 +1,4 @@
-use super::socket::{self, PySocketRef};
+use super::socket::{self, PySocket};
 use crate::common::lock::{PyRwLock, PyRwLockWriteGuard};
 use crate::{
     builtins::{pytype, weakref::PyWeak, PyStrRef, PyTypeRef},
@@ -592,7 +592,7 @@ impl PySslContext {
             }
         }
 
-        let stream = ssl::SslStream::new(ssl, args.sock.clone())
+        let stream = ssl::SslStream::new(ssl, SocketStream(args.sock.clone()))
             .map_err(|e| convert_openssl_error(vm, e))?;
 
         // TODO: use this
@@ -611,7 +611,7 @@ impl PySslContext {
 #[derive(FromArgs)]
 struct WrapSocketArgs {
     #[pyarg(any)]
-    sock: PySocketRef,
+    sock: PyRef<PySocket>,
     #[pyarg(any)]
     server_side: bool,
     #[pyarg(any, default)]
@@ -647,8 +647,8 @@ struct SocketTimeout {
     deadline: Result<Instant, bool>,
 }
 impl SocketTimeout {
-    fn get(s: &socket::PySocket) -> Self {
-        let deadline = s.get_timeout().map(|d| Instant::now() + d);
+    fn get(s: &SocketStream) -> Self {
+        let deadline = s.0.get_timeout().map(|d| Instant::now() + d);
         Self { deadline }
     }
 }
@@ -659,8 +659,8 @@ enum SelectRet {
     Closed,
     Ok,
 }
-fn ssl_select(sock: &socket::PySocket, needs: SslNeeds, timeout: &SocketTimeout) -> SelectRet {
-    let sock = match sock.sock_opt() {
+fn ssl_select(sock: &SocketStream, needs: SslNeeds, timeout: &SocketTimeout) -> SelectRet {
+    let sock = match sock.0.sock_opt() {
         Some(s) => s,
         None => return SelectRet::Closed,
     };
@@ -693,7 +693,7 @@ enum SslNeeds {
 
 fn socket_needs(
     err: &ssl::Error,
-    sock: &socket::PySocket,
+    sock: &SocketStream,
     timeout: &SocketTimeout,
 ) -> (Option<SslNeeds>, SelectRet) {
     let needs = match err.code() {
@@ -716,7 +716,7 @@ fn socket_closed_error(vm: &VirtualMachine) -> PyBaseExceptionRef {
 #[derive(PyValue)]
 struct PySslSocket {
     ctx: PyRef<PySslContext>,
-    stream: PyRwLock<ssl::SslStream<PySocketRef>>,
+    stream: PyRwLock<ssl::SslStream<SocketStream>>,
     socket_type: SslServerOrClient,
     server_hostname: Option<PyStrRef>,
     owner: PyRwLock<Option<PyWeak>>,
@@ -819,7 +819,7 @@ impl PySslSocket {
                 Ok(()) => return Ok(()),
                 Err(e) => e,
             };
-            let (needs, state) = socket_needs(&err, stream.get_ref(), &timeout);
+            let (needs, state) = socket_needs(&err, &stream.get_ref(), &timeout);
             match state {
                 SelectRet::TimedOut => {
                     return Err(socket::timeout_error_msg(
@@ -1237,3 +1237,19 @@ fn extend_module_platform_specific(module: &PyObjectRef, vm: &VirtualMachine) {
 
 #[cfg(not(windows))]
 fn extend_module_platform_specific(_module: &PyObjectRef, _vm: &VirtualMachine) {}
+
+struct SocketStream(PyRef<PySocket>);
+
+impl std::io::Read for SocketStream {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        <&socket2::Socket as std::io::Read>::read(&mut &*self.0.sock_io()?, buf)
+    }
+}
+impl std::io::Write for SocketStream {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        <&socket2::Socket as std::io::Write>::write(&mut &*self.0.sock_io()?, buf)
+    }
+    fn flush(&mut self) -> std::io::Result<()> {
+        <&socket2::Socket as std::io::Write>::flush(&mut &*self.0.sock_io()?)
+    }
+}
