8308445: Linker should check that capture state segment is big enough

Reviewed-by: mcimadamore

Author: JornVernee

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

@@ -96,6 +96,7 @@ private final MethodHandle downcallHandle0(FunctionDescriptor function, Option..
             FunctionDescriptor fd = linkRequest.descriptor();
             MethodType type = fd.toMethodType();
             MethodHandle handle = arrangeDowncall(type, fd, linkRequest.options());
+            handle = SharedUtils.maybeCheckCaptureSegment(handle, linkRequest.options());
             handle = SharedUtils.maybeInsertAllocator(fd, handle);
             return handle;
         });

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

@@ -78,6 +78,7 @@ private SharedUtils() {
     private static final MethodHandle MH_BUFFER_COPY;
     private static final MethodHandle MH_REACHABILITY_FENCE;
     public static final MethodHandle MH_CHECK_SYMBOL;
+    private static final MethodHandle MH_CHECK_CAPTURE_SEGMENT;
 
     public static final AddressLayout C_POINTER = ADDRESS
             .withTargetLayout(MemoryLayout.sequenceLayout(JAVA_BYTE));
@@ -110,6 +111,8 @@ public void close() {
                     methodType(void.class, Object.class));
             MH_CHECK_SYMBOL = lookup.findStatic(SharedUtils.class, "checkSymbol",
                     methodType(void.class, MemorySegment.class));
+            MH_CHECK_CAPTURE_SEGMENT = lookup.findStatic(SharedUtils.class, "checkCaptureSegment",
+                    methodType(MemorySegment.class, MemorySegment.class));
         } catch (ReflectiveOperationException e) {
             throw new BootstrapMethodError(e);
         }
@@ -343,6 +346,23 @@ public static MethodHandle maybeInsertAllocator(FunctionDescriptor descriptor, M
         return handle;
     }
 
+    public static MethodHandle maybeCheckCaptureSegment(MethodHandle handle, LinkerOptions options) {
+        if (options.hasCapturedCallState()) {
+            // (<target address>, SegmentAllocator, <capture segment>, ...) -> ...
+            handle = MethodHandles.filterArguments(handle, 2, MH_CHECK_CAPTURE_SEGMENT);
+        }
+        return handle;
+    }
+
+    @ForceInline
+    public static MemorySegment checkCaptureSegment(MemorySegment captureSegment) {
+        Objects.requireNonNull(captureSegment);
+        if (captureSegment.equals(MemorySegment.NULL)) {
+            throw new IllegalArgumentException("Capture segment is NULL: " + captureSegment);
+        }
+        return captureSegment.asSlice(0, CapturableState.LAYOUT);
+    }
+
     @ForceInline
     public static void checkSymbol(MemorySegment symbol) {
         Objects.requireNonNull(symbol);

src/java.base/share/classes/jdk/internal/foreign/abi/fallback/FallbackLinker.java

@@ -49,7 +49,6 @@
 import java.util.function.Consumer;
 
 import static java.lang.foreign.ValueLayout.ADDRESS;
-import static java.lang.foreign.ValueLayout.JAVA_LONG;
 import static java.lang.invoke.MethodHandles.foldArguments;
 
 public final class FallbackLinker extends AbstractLinker {
@@ -161,7 +160,7 @@ private static Object doDowncall(SegmentAllocator returnAllocator, Object[] args
 
             MemorySegment capturedState = null;
             if (invData.capturedStateMask() != 0) {
-                capturedState = (MemorySegment) args[argStart++];
+                capturedState = SharedUtils.checkCaptureSegment((MemorySegment) args[argStart++]);
                 MemorySessionImpl capturedStateImpl = ((AbstractMemorySegmentImpl) capturedState).sessionImpl();
                 capturedStateImpl.acquire0();
                 acquiredSessions.add(capturedStateImpl);

test/jdk/java/foreign/capturecallstate/TestCaptureCallState.java

@@ -50,6 +50,7 @@
 import static java.lang.foreign.ValueLayout.JAVA_INT;
 import static java.lang.foreign.ValueLayout.JAVA_LONG;
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
 
 public class TestCaptureCallState extends NativeTestHelper {
 
@@ -85,6 +86,21 @@ public void testSavedThreadLocal(SaveValuesCase testCase) throws Throwable {
         }
     }
 
+    @Test(dataProvider = "invalidCaptureSegmentCases")
+    public void testInvalidCaptureSegment(MemorySegment captureSegment,
+                                          Class<?> expectedExceptionType, String expectedExceptionMessage) {
+        Linker.Option stl = Linker.Option.captureCallState("errno");
+        MethodHandle handle = downcallHandle("set_errno_V", FunctionDescriptor.ofVoid(C_INT), stl);
+
+        try {
+            int testValue = 42;
+            handle.invoke(captureSegment, testValue); // should throw
+        } catch (Throwable t) {
+            assertTrue(expectedExceptionType.isInstance(t));
+            assertTrue(t.getMessage().matches(expectedExceptionMessage));
+        }
+    }
+
     @DataProvider
     public static Object[][] cases() {
         List<SaveValuesCase> cases = new ArrayList<>();
@@ -128,4 +144,13 @@ static SaveValuesCase structCase(String name, Map<MemoryLayout, Object> fields)
         return new SaveValuesCase("set_errno_" + name, FunctionDescriptor.of(layout, JAVA_INT), "errno", check);
     }
 
+    @DataProvider
+    public static Object[][] invalidCaptureSegmentCases() {
+        return new Object[][]{
+            {Arena.ofAuto().allocate(1), IndexOutOfBoundsException.class, ".*Out of bound access on segment.*"},
+            {MemorySegment.NULL, IllegalArgumentException.class, ".*Capture segment is NULL.*"},
+            {Arena.ofAuto().allocate(Linker.Option.captureStateLayout().byteSize() + 3).asSlice(3), // misaligned
+                    IllegalArgumentException.class, ".*Target offset incompatible with alignment constraints.*"},
+        };
+    }
 }