[GR-47856] Fix using recv_into with buffers without internal byte array

msimacek · msimacek · commit 45dfc3d8fb40 · 2023-08-10T09:26:37.000Z
PullRequest: graalpython/2909
diff --git a/graalpython/com.oracle.graal.python.test/src/tests/test_socket.py b/graalpython/com.oracle.graal.python.test/src/tests/test_socket.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2020, 2021, Oracle and/or its affiliates. All rights reserved.
+# Copyright (c) 2020, 2023, Oracle and/or its affiliates. All rights reserved.
 # DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 #
 # The Universal Permissive License (UPL), Version 1.0
@@ -36,10 +36,10 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
-
-import unittest
-
 import socket
+import sys
+import threading
+import unittest
 
 
 def test_inet_aton():
@@ -58,10 +58,43 @@ def test_inet_aton_errs(self):
         self.assertRaises(OSError, lambda : socket.inet_aton('255.255.256.1'))
         self.assertRaises(TypeError, lambda : socket.inet_aton(255))
 
-def test_get_name_info():    
+def test_get_name_info():
     import socket
     try :
         socket.getnameinfo((1, 0, 0, 0), 0)
     except TypeError:
         raised = True
     assert raised
+
+
+def test_recv_into():
+    port = None
+    event = threading.Event()
+    def server():
+        nonlocal port
+        with socket.create_server(('localhost', 0)) as sock:
+            port = sock.getsockname()[1]
+            event.set()
+            conn, addr = sock.accept()
+            conn.send(b'123')
+            conn.close()
+    thread = threading.Thread(target=server)
+    thread.start()
+    event.wait()
+    with socket.create_connection(('localhost', port)) as sock:
+        # Byte buffer with direct access to internal array
+        b = bytearray(b'aaa')
+        sock.recv_into(b, 1)
+        assert b == b'1aa'
+        # Byte buffer with offset, this currently doesn't have internal array acces, but we might implement it later
+        buffer = memoryview(b)[1:]
+        sock.recv_into(buffer, 1)
+        assert b == b'12a'
+        if sys.implementation.name == 'graalpy':
+            assert hasattr(__graalpython__, 'storage_to_native'), "Needs to be run with --python.EnableDebuggingBuiltins"
+            __graalpython__.storage_to_native(b)
+        # Native buffer, no internal array
+        buffer = memoryview(buffer)[1:]
+        sock.recv_into(buffer, 1)
+        assert b == b'123'
+    thread.join()
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/GraalPythonModuleBuiltins.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/GraalPythonModuleBuiltins.java
@@ -49,6 +49,7 @@
 import static com.oracle.graal.python.nodes.SpecialMethodNames.T_INSERT;
 import static com.oracle.graal.python.nodes.StringLiterals.J_LLVM_LANGUAGE;
 import static com.oracle.graal.python.nodes.StringLiterals.T_COLON;
+import static com.oracle.graal.python.nodes.StringLiterals.T_EMPTY_STRING;
 import static com.oracle.graal.python.nodes.StringLiterals.T_PATH;
 import static com.oracle.graal.python.nodes.StringLiterals.T_STRICT;
 import static com.oracle.graal.python.nodes.StringLiterals.T_SURROGATEESCAPE;
@@ -90,7 +91,10 @@
 import com.oracle.graal.python.builtins.modules.GraalPythonModuleBuiltinsFactory.DebugNodeFactory;
 import com.oracle.graal.python.builtins.objects.PNone;
 import com.oracle.graal.python.builtins.objects.bytes.PBytes;
+import com.oracle.graal.python.builtins.objects.bytes.PBytesLike;
 import com.oracle.graal.python.builtins.objects.cext.PythonAbstractNativeObject;
+import com.oracle.graal.python.builtins.objects.cext.capi.CApiContext;
+import com.oracle.graal.python.builtins.objects.cext.capi.PySequenceArrayWrapper.ToNativeStorageNode;
 import com.oracle.graal.python.builtins.objects.code.CodeNodes;
 import com.oracle.graal.python.builtins.objects.code.PCode;
 import com.oracle.graal.python.builtins.objects.common.DynamicObjectStorage;
@@ -140,6 +144,8 @@
 import com.oracle.graal.python.runtime.exception.PException;
 import com.oracle.graal.python.runtime.exception.PythonExitException;
 import com.oracle.graal.python.runtime.object.PythonObjectFactory;
+import com.oracle.graal.python.runtime.sequence.PSequence;
+import com.oracle.graal.python.runtime.sequence.storage.NativeSequenceStorage;
 import com.oracle.graal.python.runtime.sequence.storage.SequenceStorage;
 import com.oracle.graal.python.util.PythonUtils;
 import com.oracle.truffle.api.CallTarget;
@@ -253,6 +259,7 @@ public void postInitialize(Python3Core core) {
             mod.setAttribute(tsLiteral("dump_truffle_ast"), PNone.NO_VALUE);
             mod.setAttribute(tsLiteral("tdebug"), PNone.NO_VALUE);
             mod.setAttribute(tsLiteral("set_storage_strategy"), PNone.NO_VALUE);
+            mod.setAttribute(tsLiteral("storage_to_native"), PNone.NO_VALUE);
             mod.setAttribute(tsLiteral("dump_heap"), PNone.NO_VALUE);
             mod.setAttribute(tsLiteral("is_native_object"), PNone.NO_VALUE);
         }
@@ -779,6 +786,36 @@ private void validate(HashingStorage dictStorage) {
         }
     }
 
+    @Builtin(name = "storage_to_native", minNumOfPositionalArgs = 1)
+    @GenerateNodeFactory
+    abstract static class StorageToNative extends PythonUnaryBuiltinNode {
+        @Specialization
+        @TruffleBoundary
+        Object toNative(PBytesLike bytes) {
+            ensureCapi();
+            NativeSequenceStorage newStorage = ToNativeStorageNode.getUncached().execute(bytes.getSequenceStorage(), true);
+            bytes.setSequenceStorage(newStorage);
+            return bytes;
+        }
+
+        @Specialization
+        @TruffleBoundary
+        Object toNative(PSequence sequence) {
+            ensureCapi();
+            NativeSequenceStorage newStorage = ToNativeStorageNode.getUncached().execute(sequence.getSequenceStorage(), false);
+            sequence.setSequenceStorage(newStorage);
+            return sequence;
+        }
+
+        private void ensureCapi() {
+            try {
+                CApiContext.ensureCapiWasLoaded(null, getContext(), T_EMPTY_STRING, T_EMPTY_STRING);
+            } catch (Exception e) {
+                throw CompilerDirectives.shouldNotReachHere(e);
+            }
+        }
+    }
+
     @Builtin(name = J_EXTEND, minNumOfPositionalArgs = 1, doc = "Extends Java class and return HostAdapterCLass")
     @GenerateNodeFactory
     public abstract static class JavaExtendNode extends PythonUnaryBuiltinNode {
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/PySequenceArrayWrapper.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/PySequenceArrayWrapper.java
@@ -62,7 +62,7 @@
 public final class PySequenceArrayWrapper {
 
     @GenerateUncached
-    abstract static class ToNativeStorageNode extends Node {
+    public abstract static class ToNativeStorageNode extends Node {
 
         public abstract NativeSequenceStorage execute(SequenceStorage object, boolean isBytesLike);
 
@@ -104,6 +104,10 @@ static NativeSequenceStorage doEmptyStorage(@SuppressWarnings("unused") EmptySeq
         protected static boolean isNative(SequenceStorage s) {
             return s instanceof NativeSequenceStorage;
         }
+
+        public static ToNativeStorageNode getUncached() {
+            return ToNativeStorageNodeGen.getUncached();
+        }
     }
 
     @TruffleBoundary
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/socket/SocketBuiltins.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/socket/SocketBuiltins.java
@@ -680,7 +680,7 @@ Object recvInto(VirtualFrame frame, PSocket socket, Object bufferObj, int recvle
                                     () -> posixLib.recv(getPosixSupport(), socket.getFd(), bytes, 0, len, flags),
                                     false, false);
                     if (!directWrite) {
-                        bufferLib.readIntoByteArray(buffer, 0, bytes, 0, outlen);
+                        bufferLib.writeFromByteArray(buffer, 0, bytes, 0, outlen);
                     }
                     return outlen;
                 } catch (PosixException e) {
@@ -743,7 +743,7 @@ Object recvFromInto(VirtualFrame frame, PSocket socket, Object bufferObj, int re
                                     () -> posixLib.recvfrom(getPosixSupport(), socket.getFd(), bytes, 0, bytes.length, flags),
                                     false, false);
                     if (!directWrite) {
-                        bufferLib.readIntoByteArray(buffer, 0, bytes, 0, result.readBytes);
+                        bufferLib.writeFromByteArray(buffer, 0, bytes, 0, result.readBytes);
                     }
                     return factory().createTuple(new Object[]{result.readBytes, makeSockAddrNode.execute(frame, result.sockAddr)});
                 } catch (PosixException e) {
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/ssl/SSLSocketBuiltins.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/ssl/SSLSocketBuiltins.java
@@ -135,7 +135,7 @@ Object readInto(VirtualFrame frame, PSSLSocket self, int len, Object bufferObj,
                 PythonUtils.flipBuffer(output);
                 int readBytes = PythonUtils.getBufferRemaining(output);
                 if (!directWrite) {
-                    bufferLib.readIntoByteArray(buffer, 0, bytes, 0, readBytes);
+                    bufferLib.writeFromByteArray(buffer, 0, bytes, 0, readBytes);
                 }
                 return readBytes;
             } finally {

Original file line number	Diff line number	Diff line change
`@@ -135,7 +135,7 @@ Object readInto(VirtualFrame frame, PSSLSocket self, int len, Object bufferObj,`
`135`	`135`	`PythonUtils.flipBuffer(output);`
`136`	`136`	`int readBytes = PythonUtils.getBufferRemaining(output);`
`137`	`137`	`if (!directWrite) {`
`138`		`- bufferLib.readIntoByteArray(buffer, 0, bytes, 0, readBytes);`
	`138`	`+ bufferLib.writeFromByteArray(buffer, 0, bytes, 0, readBytes);`
`139`	`139`	`}`
`140`	`140`	`return readBytes;`
`141`	`141`	`} finally {`