diff --git a/libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/JdkVectorLibrary.java b/libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/JdkVectorLibrary.java index 006a2de4a05d2..e9efc267f14b6 100644 --- a/libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/JdkVectorLibrary.java +++ b/libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/JdkVectorLibrary.java @@ -328,6 +328,7 @@ static boolean checkBulkOffsets( Objects.checkFromIndexSize(0L, rowBytes, b.byteSize()); Objects.checkFromIndexSize(0L, (long) count * Integer.BYTES, offsets.byteSize()); Objects.checkFromIndexSize(0L, (long) count * Float.BYTES, result.byteSize()); + assert validateBulkOffsets(a, offsets, count, length, pitch, result, rowBytes); return true; } @@ -350,6 +351,52 @@ static boolean checkBBQBulkOffsets( Objects.checkFromIndexSize(0L, (long) datasetVectorLengthInBytes * (queryBits / dataBits), b.byteSize()); Objects.checkFromIndexSize(0L, (long) count * Integer.BYTES, offsets.byteSize()); Objects.checkFromIndexSize(0L, (long) count * Float.BYTES, result.byteSize()); + assert validateBBQBulkOffsets(a, offsets, count, datasetVectorLengthInBytes, pitch, result); + return true; + } + + static boolean validateBulkOffsets( + MemorySegment a, + MemorySegment offsets, + int count, + int length, + int pitch, + MemorySegment result, + long rowBytes + ) { + if (count < 0) throw new IllegalArgumentException("count must be non-negative: " + count); + if (length <= 0) throw new IllegalArgumentException("length must be positive: " + length); + if (pitch <= 0) throw new IllegalArgumentException("pitch must be positive: " + pitch); + checkSegmentAlignment(offsets, Integer.BYTES, "offsets", "int"); + checkSegmentAlignment(result, Float.BYTES, "result", "float"); + long aSize = a.byteSize(); + for (int i = 0; i < count; i++) { + int offset = offsets.getAtIndex(JAVA_INT, i); + Objects.checkFromIndexSize((long) offset * pitch, rowBytes, aSize); + } + return true; + } + + static boolean validateBBQBulkOffsets( + MemorySegment a, + MemorySegment offsets, + int count, + int datasetVectorLengthInBytes, + int pitch, + MemorySegment result + ) { + if (count < 0) throw new IllegalArgumentException("count must be non-negative: " + count); + if (datasetVectorLengthInBytes <= 0) { + throw new IllegalArgumentException("datasetVectorLengthInBytes must be positive: " + datasetVectorLengthInBytes); + } + if (pitch <= 0) throw new IllegalArgumentException("pitch must be positive: " + pitch); + checkSegmentAlignment(offsets, Integer.BYTES, "offsets", "int"); + checkSegmentAlignment(result, Float.BYTES, "result", "float"); + long aSize = a.byteSize(); + for (int i = 0; i < count; i++) { + int offset = offsets.getAtIndex(JAVA_INT, i); + Objects.checkFromIndexSize((long) offset * pitch, datasetVectorLengthInBytes, aSize); + } return true; } diff --git a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryBBQTests.java b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryBBQTests.java index 0d6b30b9d037c..332cbe3ec9ab8 100644 --- a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryBBQTests.java +++ b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryBBQTests.java @@ -321,6 +321,34 @@ public void testBulkIllegalDims() { assertThat(ex.getMessage(), containsString("out of bounds for length")); } + // Verifies that individual offset values are bounds-checked against the data segment. + public void testBulkOffsetsOutOfRange() { + assumeTrue(notSupportedMsg(), supported()); + final int indexVectorBytes = numBytes(size, type.dataBits()); + final int queryVectorBytes = numBytes(size, type.queryBits()); + final int numVecs = 3; + var indexSegment = arena.allocate((long) indexVectorBytes * numVecs); + var query = arena.allocate(queryVectorBytes); + var scores = arena.allocate((long) numVecs * Float.BYTES); + var offsetsSegment = arena.allocate((long) numVecs * Integer.BYTES); + + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 0, 0); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 1, numVecs); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 2, 0); + Exception ex = expectThrows( + IOOBE, + () -> nativeSimilarityBulkWithOffsets(indexSegment, query, indexVectorBytes, indexVectorBytes, offsetsSegment, numVecs, scores) + ); + assertThat(ex.getMessage(), containsString("out of bounds for length")); + + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 1, -1); + ex = expectThrows( + IOOBE, + () -> nativeSimilarityBulkWithOffsets(indexSegment, query, indexVectorBytes, indexVectorBytes, offsetsSegment, numVecs, scores) + ); + assertThat(ex.getMessage(), containsString("out of bounds for length")); + } + private static void pack(byte[] unpackedVector, byte[] packedVector, byte elementBits) { for (int i = 0; i < unpackedVector.length; i++) { var value = unpackedVector[i]; diff --git a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryFloat32Tests.java b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryFloat32Tests.java index 4e967bd339e5c..da9c1197b4397 100644 --- a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryFloat32Tests.java +++ b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryFloat32Tests.java @@ -226,6 +226,31 @@ public void testFloat32BulkWithOffsetsHeapSegments() { assertArrayEquals(expectedScores, bulkScores, delta); } + // Verifies that individual offset values are bounds-checked against the data segment. + public void testBulkOffsetsOutOfRange() { + assumeTrue(notSupportedMsg(), supported()); + final int dims = size; + final int numVecs = 3; + final int pitch = dims * Float.BYTES; + var vectorsSegment = arena.allocate((long) pitch * numVecs); + var query = arena.allocate((long) dims * Float.BYTES); + var scores = arena.allocate((long) numVecs * Float.BYTES); + var offsetsSegment = arena.allocate((long) numVecs * Integer.BYTES); + + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 0, 0); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 1, numVecs); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 2, 0); + Exception ex = expectThrows( + IOOBE, + () -> similarityBulkWithOffsets(vectorsSegment, query, dims, pitch, offsetsSegment, numVecs, scores) + ); + assertThat(ex.getMessage(), containsString("out of bounds for length")); + + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 1, -1); + ex = expectThrows(IOOBE, () -> similarityBulkWithOffsets(vectorsSegment, query, dims, pitch, offsetsSegment, numVecs, scores)); + assertThat(ex.getMessage(), containsString("out of bounds for length")); + } + public void testBulkIllegalDims() { assumeTrue(notSupportedMsg(), supported()); var segA = arena.allocate((long) size * 3 * Float.BYTES); diff --git a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt4Tests.java b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt4Tests.java index 525c82db46ff0..2123067b38387 100644 --- a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt4Tests.java +++ b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt4Tests.java @@ -272,6 +272,36 @@ public void testBulkIllegalDims() { assertThat(ex.getMessage(), containsString("out of bounds for length")); } + // Verifies that individual offset values are bounds-checked against the data segment. + public void testBulkOffsetsOutOfRange() { + assumeTrue(notSupportedMsg(), supported()); + final int packedLen = size / 2; + // INT4 length is packedLen (bytes) not element count; checkBulkOffsets computes + // rowBytes = packedLen * 4 / 8 which truncates to 0 when packedLen < 2. + assumeTrue("INT4 bounds check requires packedLen >= 2", packedLen >= 2); + final int numVecs = 3; + var packedSegment = arena.allocate((long) packedLen * numVecs); + var query = arena.allocate(size); + var scores = arena.allocate((long) numVecs * Float.BYTES); + var offsetsSegment = arena.allocate((long) numVecs * Integer.BYTES); + + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 0, 0); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 1, numVecs); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 2, 0); + Exception ex = expectThrows( + IOOBE, + () -> similarityBulkWithOffsets(packedSegment, query, packedLen, packedLen, offsetsSegment, numVecs, scores) + ); + assertThat(ex.getMessage(), containsString("out of bounds for length")); + + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 1, -1); + ex = expectThrows( + IOOBE, + () -> similarityBulkWithOffsets(packedSegment, query, packedLen, packedLen, offsetsSegment, numVecs, scores) + ); + assertThat(ex.getMessage(), containsString("out of bounds for length")); + } + int similarity(MemorySegment unpacked, MemorySegment packed, int packedLen) { try { return (int) getVectorDistance().getHandle( diff --git a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt7uTests.java b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt7uTests.java index 94cf52cb47da1..f329e0d8811b7 100644 --- a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt7uTests.java +++ b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt7uTests.java @@ -350,6 +350,32 @@ public void testBulkIllegalDims() { assertThat(ex.getMessage(), containsString("out of bounds for length")); } + // Verifies that individual offset values are bounds-checked against the data segment. + public void testBulkOffsetsOutOfRange() { + assumeTrue(notSupportedMsg(), supported()); + final int dims = size; + final int numVecs = 3; + var vectorsSegment = arena.allocate((long) dims * numVecs); + var query = arena.allocate(dims); + var scores = arena.allocate((long) numVecs * Float.BYTES); + var offsetsSegment = arena.allocate((long) numVecs * Integer.BYTES); + + // One offset beyond the data segment + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 0, 0); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 1, numVecs); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 2, 0); + Exception ex = expectThrows( + IOOBE, + () -> similarityBulkWithOffsets(vectorsSegment, query, dims, dims, offsetsSegment, numVecs, scores) + ); + assertThat(ex.getMessage(), containsString("out of bounds for length")); + + // Negative offset + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 1, -1); + ex = expectThrows(IOOBE, () -> similarityBulkWithOffsets(vectorsSegment, query, dims, dims, offsetsSegment, numVecs, scores)); + assertThat(ex.getMessage(), containsString("out of bounds for length")); + } + // Verifies that bulk sparse similarity rejects invalid arguments (undersized segments, // negative dims/count) with appropriate out-of-bounds exceptions. public void testBulkSparseIllegalArgs() { diff --git a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt8Tests.java b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt8Tests.java index d8ce4b856f6f2..ba3a461e41b79 100644 --- a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt8Tests.java +++ b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt8Tests.java @@ -259,6 +259,30 @@ public void testByteBulkWithOffsetsHeapSegments() { assertArrayEquals(expectedScores, bulkScores, delta); } + // Verifies that individual offset values are bounds-checked against the data segment. + public void testBulkOffsetsOutOfRange() { + assumeTrue(notSupportedMsg(), supported()); + final int dims = size; + final int numVecs = 3; + var vectorsSegment = arena.allocate((long) dims * numVecs); + var query = arena.allocate(dims); + var scores = arena.allocate((long) numVecs * Float.BYTES); + var offsetsSegment = arena.allocate((long) numVecs * Integer.BYTES); + + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 0, 0); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 1, numVecs); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 2, 0); + Exception ex = expectThrows( + IOOBE, + () -> similarityBulkWithOffsets(vectorsSegment, query, dims, dims, offsetsSegment, numVecs, scores) + ); + assertThat(ex.getMessage(), containsString("out of bounds for length")); + + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, 1, -1); + ex = expectThrows(IOOBE, () -> similarityBulkWithOffsets(vectorsSegment, query, dims, dims, offsetsSegment, numVecs, scores)); + assertThat(ex.getMessage(), containsString("out of bounds for length")); + } + public void testBulkIllegalDims() { assumeTrue(notSupportedMsg(), supported()); var segA = arena.allocate((long) size * 3);