8362169: Pointer passed to upcall may get wrong scope

src/java.base/share/classes/jdk/internal/foreign/abi/BindingSpecializer.java

@@ -297,7 +297,7 @@ private void specialize() {
         if (callingSequence.allocationSize() != 0) {
             cb.loadConstant(callingSequence.allocationSize())
               .invokestatic(CD_SharedUtils, "newBoundedArena", MTD_NEW_BOUNDED_ARENA);
-        } else if (callingSequence.forUpcall() && needsSession()) {
+        } else if (callingSequence.forUpcall() && anyArgNeedsScope()) {
             cb.invokestatic(CD_SharedUtils, "newEmptyArena", MTD_NEW_EMPTY_ARENA);
         } else {
             cb.getstatic(CD_SharedUtils, "DUMMY_ARENA", CD_Arena);
@@ -437,7 +437,7 @@ private void specialize() {
         cb.exceptionCatchAll(tryStart, tryEnd, catchStart);
     }
 
-    private boolean needsSession() {
+    private boolean anyArgNeedsScope() {
         return callingSequence.argumentBindings()
                 .filter(BoxAddress.class::isInstance)
                 .map(BoxAddress.class::cast)
@@ -590,7 +590,7 @@ private void emitBoxAddress(BoxAddress boxAddress) {
         popType(long.class);
         cb.loadConstant(boxAddress.size())
           .loadConstant(boxAddress.align());
-        if (needsSession()) {
+        if (boxAddress.needsScope()) {
             emitLoadInternalSession();
             cb.invokestatic(CD_Utils, "longToAddress", MTD_LONG_TO_ADDRESS_SCOPE);
         } else {

test/jdk/java/foreign/TestUpcallStructScope.java

@@ -43,31 +43,48 @@
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 
 import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.assertEquals;
 
 public class TestUpcallStructScope extends NativeTestHelper {
     static final MethodHandle MH_do_upcall;
+    static final MethodHandle MH_do_upcall_ptr;
     static final MethodHandle MH_Consumer_accept;
+    static final MethodHandle MH_BiConsumer_accept;
 
     static {
         System.loadLibrary("TestUpcallStructScope");
         MH_do_upcall = LINKER.downcallHandle(
                 findNativeOrThrow("do_upcall"),
                 FunctionDescriptor.ofVoid(C_POINTER, S_PDI_LAYOUT)
         );
+        MH_do_upcall_ptr = LINKER.downcallHandle(
+                findNativeOrThrow("do_upcall_ptr"),
+                FunctionDescriptor.ofVoid(C_POINTER, S_PDI_LAYOUT, C_POINTER)
+        );
 
         try {
             MH_Consumer_accept = MethodHandles.publicLookup().findVirtual(Consumer.class, "accept",
                     MethodType.methodType(void.class, Object.class));
+            MH_BiConsumer_accept = MethodHandles.publicLookup().findVirtual(BiConsumer.class, "accept",
+                    MethodType.methodType(void.class, Object.class, Object.class));
         } catch (NoSuchMethodException | IllegalAccessException e) {
             throw new RuntimeException(e);
         }
     }
 
-    private static MethodHandle methodHandle (Consumer<MemorySegment> callback) {
-        return MH_Consumer_accept.bindTo(callback).asType(MethodType.methodType(void.class, MemorySegment.class));
+    private static MethodHandle methodHandle(Consumer<MemorySegment> callback) {
+        return MH_Consumer_accept.bindTo(callback)
+                .asType(MethodType.methodType(void.class, MemorySegment.class));
+    }
+
+    private static MethodHandle methodHandle(BiConsumer<MemorySegment, MemorySegment> callback) {
+        return MH_BiConsumer_accept.bindTo(callback)
+                .asType(MethodType.methodType(void.class, MemorySegment.class, MemorySegment.class));
     }
 
     @Test
@@ -85,4 +102,22 @@ public void testUpcall() throws Throwable {
         assertFalse(captured.scope().isAlive());
     }
 
+    @Test
+    public void testOtherPointer() throws Throwable {
+        AtomicReference<MemorySegment> capturedSegment = new AtomicReference<>();
+        MethodHandle target = methodHandle((_, addr) -> capturedSegment.set(addr));
+        FunctionDescriptor upcallDesc = FunctionDescriptor.ofVoid(S_PDI_LAYOUT, C_POINTER);
+        MemorySegment argAddr = MemorySegment.ofAddress(42);
+        try (Arena arena = Arena.ofConfined()) {
+            MemorySegment upcallStub = LINKER.upcallStub(target, upcallDesc, arena);
+            MemorySegment argSegment = arena.allocate(S_PDI_LAYOUT);
+            MH_do_upcall_ptr.invoke(upcallStub, argSegment, argAddr);
+        }
+
+        // We've captured the address '42' from the upcall. This should have
+        // the global scope, so it should still be alive here.
+        MemorySegment captured = capturedSegment.get();
+        assertEquals(argAddr, captured);
+        assertTrue(captured.scope().isAlive());
+    }
 }

test/jdk/java/foreign/libTestUpcallStructScope.c

@@ -28,3 +28,7 @@ struct S_PDI { void* p0; double p1; int p2; };
 EXPORT void do_upcall(void (*cb)(struct S_PDI), struct S_PDI a0) {
     cb(a0);
 }
+
+EXPORT void do_upcall_ptr(void (*cb)(struct S_PDI, void*), struct S_PDI a0, void* ptr) {
+    cb(a0, ptr);
+}