DemoQueryAllocatedRanges.java
Posted by Anonymous on Thu 20th Feb 2025 11:48
raw | new post


/*
 * run with
 * $ java --enable-native-access DemoQueryAllocatedRanges -classpath . DemoQueryAllocatedRanges
 * compile with:
 * javac --enable-preview --add-modules java.base --release 23 DemoQueryAllocatedRanges.java
 *
 * Currently fails like this:
 * ---- snip ----
 * $ java --enable-native-access DemoQueryAllocatedRanges -classpath . DemoQueryAllocatedRanges
 * WARNING: Unknown module: DemoQueryAllocatedRanges specified to --enable-native-access
 * Running QueryAllocatedRanges for file: foo.bin
 * Error running QueryAllocatedRanges:
 * java.lang.ExceptionInInitializerError
 *         at DemoQueryAllocatedRanges.main(DemoQueryAllocatedRanges.java:299)
 * Caused by: java.lang.RuntimeException: java.lang.UnsatisfiedLinkError: Function CreateFileW not found in library kernel32
 *         at WindowsLinker.downcallHandle(DemoQueryAllocatedRanges.java:270)
 *         at QueryAllocatedRanges.<clinit>(DemoQueryAllocatedRanges.java:77)
 *         ... 1 more
 * Caused by: java.lang.UnsatisfiedLinkError: Function CreateFileW not found in library kernel32
 *         at WindowsLinker.lambda$downcallHandle$0(DemoQueryAllocatedRanges.java:263)
 *         at java.base/java.util.Optional.orElseThrow(Optional.java:403)
 *         at WindowsLinker.downcallHandle(DemoQueryAllocatedRanges.java:262)
 *         ... 2 more
 * Demo finished.
 * ---- snip ----
 */

import java.io.IOException;
import java.lang.foreign.Arena;
import java.lang.foreign.MemoryLayout;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.ValueLayout;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.CharBuffer;
import java.nio.channels.FileChannel;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

import static java.lang.foreign.ValueLayout.ADDRESS;
import static java.lang.foreign.ValueLayout.JAVA_LONG;
import static java.lang.foreign.ValueLayout.JAVA_INT;
import static java.lang.foreign.ValueLayout.JAVA_BYTE;
import static java.lang.foreign.MemoryLayout.PathElement; // Import PathElement

/* needed for demo code */
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;

/* public */ class QueryAllocatedRanges {

    // Define FSCTL_QUERY_ALLOCATED_RANGES control code
    private static final int FILE_DEVICE_FILE_SYSTEM = 0x00000009;
    private static final int METHOD_BUFFERED = 0;
    private static final int FILE_ANY_ACCESS = 0;
    private static final int FSCTL_QUERY_ALLOCATED_RANGES = (FILE_DEVICE_FILE_SYSTEM << 16) | (FILE_ANY_ACCESS << 14) | (51 << 2) | METHOD_BUFFERED;

    // Define FILE_ATTRIBUTE_NORMAL for CreateFile
    private static final int FILE_ATTRIBUTE_NORMAL = 0x00000080;
    // Define GENERIC_READ for CreateFile
    private static final int GENERIC_READ = 0x80000000;
    // Define GENERIC_WRITE for CreateFile - not needed for this FSCTL but might be needed for others
    private static final int GENERIC_WRITE = 0x40000000;
    // Define OPEN_EXISTING for CreateFile
    private static final int OPEN_EXISTING = 3;
    // Define INVALID_HANDLE_VALUE
    private static final long INVALID_HANDLE_VALUE = -1;

    // Define constants for Windows API functions and structs using Foreign Function and Memory API

    // HANDLE CreateFileW(
    //   LPCWSTR               lpFileName,
    //   DWORD                 dwDesiredAccess,
    //   DWORD                 dwShareMode,
    //   LPSECURITY_ATTRIBUTES lpSecurityAttributes,
    //   DWORD                 dwCreationDisposition,
    //   DWORD                 dwFlagsAndAttributes,
    //   HANDLE                hTemplateFile
    // );
    private static final java.lang.foreign.FunctionDescriptor CreateFileWDescriptor = java.lang.foreign.FunctionDescriptor.of(
            ADDRESS, // HANDLE
            ADDRESS, // LPCWSTR lpFileName
            JAVA_INT, // DWORD dwDesiredAccess
            JAVA_INT, // DWORD dwShareMode
            ADDRESS, // LPSECURITY_ATTRIBUTES lpSecurityAttributes
            JAVA_INT, // DWORD dwCreationDisposition
            JAVA_INT, // DWORD dwFlagsAndAttributes
            ADDRESS  // HANDLE hTemplateFile
    );
    private static final java.lang.invoke.MethodHandle CreateFileW = WindowsLinker.downcallHandle(
            "kernel32", "CreateFileW", CreateFileWDescriptor
    );

    // BOOL DeviceIoControl(
    //   HANDLE                 hDevice,
    //   DWORD                  dwIoControlCode,
    //   LPVOID                 lpInBuffer,
    //   DWORD                  nInBufferSize,
    //   LPVOID                 lpOutBuffer,
    //   DWORD                  nOutBufferSize,
    //   LPDWORD                lpBytesReturned,
    //   LPOVERLAPPED             lpOverlapped
    // );
    private static final java.lang.foreign.FunctionDescriptor DeviceIoControlDescriptor = java.lang.foreign.FunctionDescriptor.of(
            JAVA_INT, // BOOL
            ADDRESS, // HANDLE hDevice
            JAVA_INT, // DWORD dwIoControlCode
            ADDRESS, // LPVOID lpInBuffer
            JAVA_INT, // DWORD nInBufferSize
            ADDRESS, // LPVOID lpOutBuffer
            JAVA_INT, // DWORD nOutBufferSize
            ADDRESS, // LPDWORD lpBytesReturned
            ADDRESS  // LPOVERLAPPED lpOverlapped
    );
    private static final java.lang.invoke.MethodHandle DeviceIoControl = WindowsLinker.downcallHandle(
            "kernel32.dll", "DeviceIoControl", DeviceIoControlDescriptor
    );

    // BOOL CloseHandle(
    //   HANDLE hObject
    // );
    private static final java.lang.foreign.FunctionDescriptor CloseHandleDescriptor = java.lang.foreign.FunctionDescriptor.of(
            JAVA_INT, // BOOL
            ADDRESS  // HANDLE hObject
    );
    private static final java.lang.invoke.MethodHandle CloseHandle = WindowsLinker.downcallHandle(
            "kernel32.dll", "CloseHandle", CloseHandleDescriptor
    );


    // typedef struct _QUERY_ALLOCATED_RANGES_BUFFER {
    //   DWORD         FileOffset;
    //   DWORD         Length;
    // } QUERY_ALLOCATED_RANGES_BUFFER, *PQUERY_ALLOCATED_RANGES_BUFFER;
    private static final MemoryLayout QUERY_ALLOCATED_RANGES_BUFFER_LAYOUT = MemoryLayout.structLayout(
            JAVA_LONG.withName("FileOffset"),
            JAVA_LONG.withName("Length")
    );

    // typedef struct _FILE_ALLOCATED_RANGE_BUFFER {
    //   LARGE_INTEGER FileOffset;
    //   LARGE_INTEGER Length;
    // } FILE_ALLOCATED_RANGE_BUFFER, *PFILE_ALLOCATED_RANGE_BUFFER;
    private static final MemoryLayout FILE_ALLOCATED_RANGE_BUFFER_LAYOUT = MemoryLayout.structLayout(
            JAVA_LONG.withName("FileOffset"),
            JAVA_LONG.withName("Length")
    );


    public static void main(String[] args) throws Throwable {
        if (args.length != 1) {
            System.out.println("Usage: java QueryAllocatedRanges <file_path>");
            return;
        }
        String filePath = args[0];

        try (Arena arena = Arena.ofConfined()) {
            byte[] filePathBytes = filePath.getBytes(StandardCharsets.UTF_16LE); // Encode in UTF-16LE
            MemorySegment lpFileName = arena.allocate(filePathBytes.length + 2); // +2 for null terminator
            lpFileName.copyFrom(MemorySegment.ofArray(filePathBytes)); // Use copyFrom instead of writeByteArray
            lpFileName.set(ValueLayout.JAVA_SHORT, filePathBytes.length, (short) 0); // Null-terminate (UTF-16 null is 2 bytes)

            MemorySegment handle = (MemorySegment) CreateFileW.invokeExact(
                    lpFileName.address(), // Pass the address of the MemorySegment
                    GENERIC_READ,
                    0, // dwShareMode = 0 for exclusive access
                    MemorySegment.NULL, // lpSecurityAttributes
                    OPEN_EXISTING,
                    FILE_ATTRIBUTE_NORMAL,
                    MemorySegment.NULL // hTemplateFile
            );

            if (handle.address() == INVALID_HANDLE_VALUE) {
                System.err.println("Error opening file: " + filePath);
                return; // Or throw an exception
            }

            MemorySegment inputBuffer = arena.allocate(QUERY_ALLOCATED_RANGES_BUFFER_LAYOUT);
            inputBuffer.set(JAVA_LONG, QUERY_ALLOCATED_RANGES_BUFFER_LAYOUT.byteOffset(PathElement.groupElement("FileOffset")), 0L); // FileOffset = 0 // Use PathElement
            inputBuffer.set(JAVA_LONG, QUERY_ALLOCATED_RANGES_BUFFER_LAYOUT.byteOffset(PathElement.groupElement("Length")), -1L); // Length = -1 (ALL remaining ranges) // Use PathElement


            int inputBufferSize = (int) QUERY_ALLOCATED_RANGES_BUFFER_LAYOUT.byteSize();

            int initialOutputBufferSize = 4096;
            MemorySegment outputBuffer = arena.allocate(initialOutputBufferSize);
            MemorySegment bytesReturnedSegment = arena.allocate(JAVA_INT);
            int outputBufferSize = initialOutputBufferSize;


            boolean success = (boolean) DeviceIoControl.invokeExact(
                    handle,
                    FSCTL_QUERY_ALLOCATED_RANGES,
                    inputBuffer,
                    inputBufferSize,
                    outputBuffer,
                    outputBufferSize,
                    bytesReturnedSegment,
                    MemorySegment.NULL // lpOverlapped
            );

            if (!success) {
                int lastError = WindowsLastError.GetLastError();
                if (lastError == 234) { // ERROR_MORE_DATA (234) - output buffer was too small
                    int bytesReturned = bytesReturnedSegment.get(JAVA_INT, 0);
                    outputBufferSize = bytesReturned;
                    outputBuffer = arena.allocate(outputBufferSize);

                    success = (boolean) DeviceIoControl.invokeExact(
                            handle,
                            FSCTL_QUERY_ALLOCATED_RANGES,
                            inputBuffer,
                            inputBufferSize,
                            outputBuffer,
                            outputBufferSize,
                            bytesReturnedSegment,
                            MemorySegment.NULL // lpOverlapped
                    );
                    if (!success) {
                        System.err.println("DeviceIoControl failed after buffer reallocation with error code: " + WindowsLastError.GetLastError());
                        CloseHandle.invokeExact(handle);
                        return;
                    }

                } else {
                    System.err.println("DeviceIoControl failed with error code: " + lastError);
                    CloseHandle.invokeExact(handle);
                    return;
                }
            }

            int bytesReturned = bytesReturnedSegment.get(JAVA_INT, 0);
            if (bytesReturned > 0) {
                System.out.println("Allocated Ranges:");
                long offset = 0;
                while (offset < bytesReturned) {
                    MemorySegment rangeSegment = outputBuffer.asSlice(offset, FILE_ALLOCATED_RANGE_BUFFER_LAYOUT.byteSize());
                    long fileOffset = rangeSegment.get(JAVA_LONG, FILE_ALLOCATED_RANGE_BUFFER_LAYOUT.byteOffset(PathElement.groupElement("FileOffset"))); // Use PathElement
                    long length = rangeSegment.get(JAVA_LONG, FILE_ALLOCATED_RANGE_BUFFER_LAYOUT.byteOffset(PathElement.groupElement("Length"))); // Use PathElement
                    System.out.println("  Offset: " + fileOffset + ", Length: " + length);
                    offset += FILE_ALLOCATED_RANGE_BUFFER_LAYOUT.byteSize();
                }
            } else {
                System.out.println("No allocated ranges found.");
            }

            CloseHandle.invokeExact(handle);
        } catch (IOException | InterruptedException |  java.lang.invoke.WrongMethodTypeException e) {
            e.printStackTrace();
        }
    }
}

class WindowsLastError {
    private static final java.lang.foreign.FunctionDescriptor GetLastErrorDescriptor = java.lang.foreign.FunctionDescriptor.of(
            JAVA_INT // DWORD
    );
    private static final java.lang.invoke.MethodHandle GetLastError = WindowsLinker.downcallHandle(
            "kernel32.dll", "GetLastError", GetLastErrorDescriptor
    );

    public static int GetLastError() throws Throwable {
        return (int) GetLastError.invokeExact();
    }
}

class WindowsLinker {
    private static final java.lang.foreign.Linker LINKER = java.lang.foreign.Linker.nativeLinker(); // Or Linker.defaultLinker()

    public static java.lang.invoke.MethodHandle downcallHandle(String libraryName, String functionName, java.lang.foreign.FunctionDescriptor descriptor) {
        try {
            // Use findSymbol and handle Optional<MemorySegment> correctly
            java.lang.foreign.SymbolLookup symbolLookup = LINKER.defaultLookup();
            java.lang.foreign.MemorySegment functionAddress = symbolLookup.find(functionName)
                    .orElseThrow(() -> // orElseThrow on Optional<MemorySegment>
                            new UnsatisfiedLinkError("Function " + functionName + " not found in library " + libraryName));

            return LINKER.downcallHandle(
                    functionAddress, // Pass the MemorySegment (function address)
                    descriptor
            );
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }
    }
}


public class DemoQueryAllocatedRanges {

    public static void main(String[] args) {
        String filePath = "foo.bin";
        Path file = Paths.get(filePath);

        // Create foo.bin if it doesn't exist
        if (!Files.exists(file)) {
            try {
                Files.createFile(file);
                System.out.println("Created file: " + filePath);
            } catch (IOException e) {
                System.err.println("Error creating file: " + filePath);
                e.printStackTrace();
                return; // Exit if file creation fails
            }
        }

        System.out.println("Running QueryAllocatedRanges for file: " + filePath);

        try {
            // Call QueryAllocatedRanges.main() with filePath as argument
            QueryAllocatedRanges.main(new String[]{filePath});
        } catch (Throwable e) {
            System.err.println("Error running QueryAllocatedRanges:");
            e.printStackTrace();
        }

        System.out.println("Demo finished.");
    }
}