Move Windows native functions into NativeAccess (#108873)

Elasticsearch uses a couple windows specific functions, specifically
gettting a short path, and registering a console control handler for
shutdown notification. This commit moves this functionality from the
existing jna natives into NativeAccess.

relates #104876

libs/native/jna/src/main/java/org/elasticsearch/nativeaccess/jna/JnaKernel32Library.java

@@ -13,8 +13,10 @@
 import com.sun.jna.NativeLong;
 import com.sun.jna.Pointer;
 import com.sun.jna.Structure;
+import com.sun.jna.WString;
 import com.sun.jna.win32.StdCallLibrary;
 
+import org.elasticsearch.nativeaccess.WindowsFunctions.ConsoleCtrlHandler;
 import org.elasticsearch.nativeaccess.lib.Kernel32Library;
 
 import java.util.List;
@@ -96,6 +98,22 @@ public long Type() {
         }
     }
 
+    /**
+     * JNA adaptation of {@link ConsoleCtrlHandler}
+     */
+    public static class NativeHandlerCallback implements StdCallLibrary.StdCallCallback {
+
+        private final ConsoleCtrlHandler handler;
+
+        public NativeHandlerCallback(ConsoleCtrlHandler handler) {
+            this.handler = handler;
+        }
+
+        public boolean callback(long dwCtrlType) {
+            return handler.handle((int) dwCtrlType);
+        }
+    }
+
     private interface NativeFunctions extends StdCallLibrary {
         Pointer GetCurrentProcess();
 
@@ -106,9 +124,14 @@ private interface NativeFunctions extends StdCallLibrary {
         int VirtualQueryEx(Pointer handle, Pointer address, JnaMemoryBasicInformation memoryInfo, int length);
 
         boolean SetProcessWorkingSetSize(Pointer handle, SizeT minSize, SizeT maxSize);
+
+        int GetShortPathNameW(WString lpszLongPath, char[] lpszShortPath, int cchBuffer);
+
+        boolean SetConsoleCtrlHandler(StdCallLibrary.StdCallCallback handler, boolean add);
     }
 
     private final NativeFunctions functions;
+    private NativeHandlerCallback consoleCtrlHandlerCallback = null;
 
     JnaKernel32Library() {
         this.functions = Native.load("kernel32", NativeFunctions.class);
@@ -161,4 +184,17 @@ public boolean SetProcessWorkingSetSize(Handle handle, long minSize, long maxSiz
         var jnaHandle = (JnaHandle) handle;
         return functions.SetProcessWorkingSetSize(jnaHandle.pointer, new SizeT(minSize), new SizeT(maxSize));
     }
+
+    @Override
+    public int GetShortPathNameW(String lpszLongPath, char[] lpszShortPath, int cchBuffer) {
+        var wideFileName = new WString(lpszLongPath);
+        return functions.GetShortPathNameW(wideFileName, lpszShortPath, cchBuffer);
+    }
+
+    @Override
+    public boolean SetConsoleCtrlHandler(ConsoleCtrlHandler handler, boolean add) {
+        assert consoleCtrlHandlerCallback == null;
+        consoleCtrlHandlerCallback = new NativeHandlerCallback(handler);
+        return functions.SetConsoleCtrlHandler(consoleCtrlHandlerCallback, true);
+    }
 }

libs/native/src/main/java/org/elasticsearch/nativeaccess/NativeAccess.java

@@ -52,6 +52,13 @@ static NativeAccess instance() {
      */
     Zstd getZstd();
 
+    /**
+     * Returns an accessor for native functions only available on Windows, or {@code null} if not on Windows.
+     */
+    default WindowsFunctions getWindowsFunctions() {
+        return null;
+    }
+
     /*
      * Returns the vector similarity functions, or an empty optional.
      */

libs/native/src/main/java/org/elasticsearch/nativeaccess/WindowsFunctions.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess;
+
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
+import org.elasticsearch.nativeaccess.lib.Kernel32Library;
+
+/**
+ * Native functions specific to the Windows operating system.
+ */
+public class WindowsFunctions {
+    private static final Logger logger = LogManager.getLogger(Systemd.class);
+
+    private final Kernel32Library kernel;
+
+    WindowsFunctions(Kernel32Library kernel) {
+        this.kernel = kernel;
+    }
+
+    /**
+     * Retrieves the short path form of the specified path.
+     *
+     * @param path the path
+     * @return the short path name, or the original path name if unsupported or unavailable
+     */
+    public String getShortPathName(String path) {
+        String longPath = "\\\\?\\" + path;
+        // first we get the length of the buffer needed
+        final int length = kernel.GetShortPathNameW(longPath, null, 0);
+        if (length == 0) {
+            logger.warn("failed to get short path name: {}", kernel.GetLastError());
+            return path;
+        }
+        final char[] shortPath = new char[length];
+        // knowing the length of the buffer, now we get the short name
+        if (kernel.GetShortPathNameW(longPath, shortPath, length) > 0) {
+            assert shortPath[length - 1] == '\0';
+            return new String(shortPath, 0, length - 1);
+        } else {
+            logger.warn("failed to get short path name: {}", kernel.GetLastError());
+            return path;
+        }
+    }
+
+    /**
+     * Adds a Console Ctrl Handler for Windows. On non-windows this is a noop.
+     *
+     * @return true if the handler is correctly set
+     */
+    public boolean addConsoleCtrlHandler(ConsoleCtrlHandler handler) {
+        return kernel.SetConsoleCtrlHandler(dwCtrlType -> {
+            if (logger.isDebugEnabled()) {
+                logger.debug("console control handler received event [{}]", dwCtrlType);
+            }
+            return handler.handle(dwCtrlType);
+        }, true);
+    }
+
+    /**
+     * Windows callback for console events
+     *
+     * @see <a href="http://msdn.microsoft.com/en-us/library/windows/desktop/ms683242%28v=vs.85%29.aspx">HandlerRoutine docs</a>
+     */
+    public interface ConsoleCtrlHandler {
+
+        int CTRL_CLOSE_EVENT = 2;
+
+        /**
+         * Handles the Ctrl event.
+         *
+         * @param code the code corresponding to the Ctrl sent.
+         * @return true if the handler processed the event, false otherwise. If false, the next handler will be called.
+         */
+        boolean handle(int code);
+    }
+}

libs/native/src/main/java/org/elasticsearch/nativeaccess/WindowsNativeAccess.java

@@ -28,10 +28,12 @@ class WindowsNativeAccess extends AbstractNativeAccess {
     public static final int MEM_COMMIT = 0x1000;
 
     private final Kernel32Library kernel;
+    private final WindowsFunctions windowsFunctions;
 
     WindowsNativeAccess(NativeLibraryProvider libraryProvider) {
         super("Windows", libraryProvider);
         this.kernel = libraryProvider.getLibrary(Kernel32Library.class);
+        this.windowsFunctions = new WindowsFunctions(kernel);
     }
 
     @Override
@@ -71,6 +73,11 @@ public ProcessLimits getProcessLimits() {
         return new ProcessLimits(ProcessLimits.UNKNOWN, ProcessLimits.UNKNOWN, ProcessLimits.UNKNOWN);
     }
 
+    @Override
+    public WindowsFunctions getWindowsFunctions() {
+        return windowsFunctions;
+    }
+
     @Override
     public Optional<VectorSimilarityFunctions> getVectorSimilarityFunctions() {
         return Optional.empty(); // not supported yet

libs/native/src/main/java/org/elasticsearch/nativeaccess/lib/Kernel32Library.java

@@ -8,6 +8,8 @@
 
 package org.elasticsearch.nativeaccess.lib;
 
+import org.elasticsearch.nativeaccess.WindowsFunctions.ConsoleCtrlHandler;
+
 public non-sealed interface Kernel32Library extends NativeLibrary {
     interface Handle {}
 
@@ -78,4 +80,25 @@ interface MemoryBasicInformation {
      * @see <a href="https://msdn.microsoft.com/en-us/library/windows/desktop/ms686234%28v=vs.85%29.aspx">SetProcessWorkingSetSize docs</a>
      */
     boolean SetProcessWorkingSetSize(Handle handle, long minSize, long maxSize);
+
+    /**
+     * Retrieves the short path form of the specified path.
+     *
+     * @param lpszLongPath  the path string
+     * @param lpszShortPath a buffer to receive the short name
+     * @param cchBuffer     the size of the buffer
+     * @return the length of the string copied into {@code lpszShortPath}, otherwise zero for failure
+     * @see <a href="https://msdn.microsoft.com/en-us/library/windows/desktop/aa364989.aspx">GetShortPathName docs</a>
+     */
+    int GetShortPathNameW(String lpszLongPath, char[] lpszShortPath, int cchBuffer);
+
+    /**
+     * Native call to the Kernel32 API to set a new Console Ctrl Handler.
+     *
+     * @param handler A callback to handle control events
+     * @param add     True if the handler should be added, false if it should replace existing handlers
+     * @return true if the handler is correctly set
+     * @see <a href="https://learn.microsoft.com/en-us/windows/console/setconsolectrlhandler">SetConsoleCtrlHandler docs</a>
+     */
+    boolean SetConsoleCtrlHandler(ConsoleCtrlHandler handler, boolean add);
 }

libs/native/src/main21/java/org/elasticsearch/nativeaccess/jdk/ArenaUtil.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess.jdk;
+
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemoryLayout;
+import java.lang.foreign.MemorySegment;
+import java.nio.charset.Charset;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+
+/**
+ * Utility methods to act on Arena apis which have changed in subsequent JDK releases.
+ */
+class ArenaUtil {
+
+    /**
+     * Allocate an array of the given memory layout.
+     */
+    static MemorySegment allocate(Arena arena, MemoryLayout layout, int count) {
+        return arena.allocateArray(layout, count);
+    }
+
+    /**
+     * Allocate and copy the given string into native memory.
+     */
+    static MemorySegment allocateFrom(Arena arena, String str, Charset charset) {
+        return arena.allocateArray(JAVA_BYTE, str.getBytes(charset));
+    }
+
+    private ArenaUtil() {}
+}

libs/native/src/main21/java/org/elasticsearch/nativeaccess/jdk/JdkKernel32Library.java

@@ -8,6 +8,7 @@
 
 package org.elasticsearch.nativeaccess.jdk;
 
+import org.elasticsearch.nativeaccess.WindowsFunctions.ConsoleCtrlHandler;
 import org.elasticsearch.nativeaccess.lib.Kernel32Library;
 
 import java.lang.foreign.Arena;
@@ -18,14 +19,18 @@
 import java.lang.foreign.StructLayout;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.VarHandle;
+import java.nio.charset.StandardCharsets;
 
 import static java.lang.foreign.MemoryLayout.PathElement.groupElement;
 import static java.lang.foreign.MemoryLayout.paddingLayout;
 import static java.lang.foreign.ValueLayout.ADDRESS;
 import static java.lang.foreign.ValueLayout.JAVA_BOOLEAN;
+import static java.lang.foreign.ValueLayout.JAVA_CHAR;
 import static java.lang.foreign.ValueLayout.JAVA_INT;
 import static java.lang.foreign.ValueLayout.JAVA_LONG;
 import static org.elasticsearch.nativeaccess.jdk.LinkerHelper.downcallHandle;
+import static org.elasticsearch.nativeaccess.jdk.LinkerHelper.upcallHandle;
+import static org.elasticsearch.nativeaccess.jdk.LinkerHelper.upcallStub;
 import static org.elasticsearch.nativeaccess.jdk.MemorySegmentUtil.varHandleWithoutOffset;
 
 class JdkKernel32Library implements Kernel32Library {
@@ -52,6 +57,21 @@ class JdkKernel32Library implements Kernel32Library {
         "SetProcessWorkingSetSize",
         FunctionDescriptor.of(ADDRESS, JAVA_LONG, JAVA_LONG)
     );
+    private static final MethodHandle GetShortPathNameW$mh = downcallHandleWithError(
+        "GetShortPathNameW",
+        FunctionDescriptor.of(JAVA_INT, ADDRESS, ADDRESS, JAVA_INT)
+    );
+    private static final MethodHandle SetConsoleCtrlHandler$mh = downcallHandleWithError(
+        "SetConsoleCtrlHandler",
+        FunctionDescriptor.of(JAVA_BOOLEAN, ADDRESS, JAVA_BOOLEAN)
+    );
+
+    private static final FunctionDescriptor ConsoleCtrlHandler_handle$fd = FunctionDescriptor.of(JAVA_BOOLEAN, JAVA_INT);
+    private static final MethodHandle ConsoleCtrlHandler_handle$mh = upcallHandle(
+        ConsoleCtrlHandler.class,
+        "handle",
+        ConsoleCtrlHandler_handle$fd
+    );
 
     private static MethodHandle downcallHandleWithError(String function, FunctionDescriptor functionDescriptor) {
         return downcallHandle(function, functionDescriptor, CAPTURE_GETLASTERROR_OPTION);
@@ -208,4 +228,38 @@ public boolean SetProcessWorkingSetSize(Handle process, long minSize, long maxSi
             throw new AssertionError(t);
         }
     }
+
+    @Override
+    public int GetShortPathNameW(String lpszLongPath, char[] lpszShortPath, int cchBuffer) {
+        try (Arena arena = Arena.ofConfined()) {
+            MemorySegment wideFileName = ArenaUtil.allocateFrom(arena, lpszLongPath + "\0", StandardCharsets.UTF_16LE);
+            MemorySegment shortPath;
+            if (lpszShortPath != null) {
+                shortPath = ArenaUtil.allocate(arena, JAVA_CHAR, cchBuffer);
+            } else {
+                shortPath = MemorySegment.NULL;
+            }
+
+            int ret = (int) GetShortPathNameW$mh.invokeExact(lastErrorState, wideFileName, shortPath, cchBuffer);
+            if (shortPath != MemorySegment.NULL) {
+                for (int i = 0; i < cchBuffer; ++i) {
+                    lpszShortPath[i] = shortPath.getAtIndex(JAVA_CHAR, i);
+                }
+            }
+            return ret;
+        } catch (Throwable t) {
+            throw new AssertionError(t);
+        }
+    }
+
+    @Override
+    public boolean SetConsoleCtrlHandler(ConsoleCtrlHandler handler, boolean add) {
+        // use the global arena so the handler will have the lifetime of the jvm
+        MemorySegment nativeHandler = upcallStub(ConsoleCtrlHandler_handle$mh, handler, ConsoleCtrlHandler_handle$fd, Arena.global());
+        try {
+            return (boolean) SetConsoleCtrlHandler$mh.invokeExact(lastErrorState, nativeHandler, add);
+        } catch (Throwable t) {
+            throw new AssertionError(t);
+        }
+    }
 }

libs/native/src/main21/java/org/elasticsearch/nativeaccess/jdk/LinkerHelper.java

@@ -22,7 +22,7 @@
 class LinkerHelper {
     private static final Linker LINKER = Linker.nativeLinker();
     private static final SymbolLookup SYMBOL_LOOKUP;
-    private static final MethodHandles.Lookup MH_LOOKUP = MethodHandles.publicLookup();
+    private static final MethodHandles.Lookup MH_LOOKUP = MethodHandles.lookup();
 
     static {
         // We first check the loader lookup, which contains libs loaded by System.load and System.loadLibrary.

libs/native/src/main22/java/org/elasticsearch/nativeaccess/jdk/ArenaUtil.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess.jdk;
+
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemoryLayout;
+import java.lang.foreign.MemorySegment;
+import java.nio.charset.Charset;
+
+public class ArenaUtil {
+
+    /**
+     * Allocate an array of the given memory layout.
+     */
+    static MemorySegment allocate(Arena arena, MemoryLayout layout, int count) {
+        return arena.allocate(layout, count);
+    }
+
+    /**
+     * Allocate and copy the given string into native memory.
+     */
+    static MemorySegment allocateFrom(Arena arena, String str, Charset charset) {
+        return arena.allocateFrom(str, charset);
+    }
+
+    private ArenaUtil() {}
+}