keystone/bindings/java/src/main/java/keystone/Keystone.java

/*
 * Copyright (c) 2018 Jämes Ménétrey <james@menetrey.me>
 *
 * This file is part of the Keystone Java bindings which is released under MIT.
 * See file LICENSE in the Java bindings folder for full license details.
 */

package keystone;

import com.sun.jna.Pointer;
import com.sun.jna.ptr.IntByReference;
import com.sun.jna.ptr.PointerByReference;
import keystone.exceptions.AssembleFailedKeystoneException;
import keystone.exceptions.OpenFailedKeystoneException;
import keystone.exceptions.SetOptionFailedKeystoneException;
import keystone.natives.CleanerContainer;
import keystone.natives.DirectMappingKeystoneNative;
import keystone.natives.KeystoneCleanerContainer;
import keystone.utilities.Version;

import java.util.concurrent.atomic.AtomicBoolean;

/**
 * The Keystone engine.
 */
public class Keystone implements AutoCloseable {

    /**
     * The pointer to the Keystone native resource.
     */
    private final Pointer ksEngine;

    /**
     * The cleaner container that frees up the native resource if this object is not properly closed and is
     * candidate for garbage collection.
     */
    private final CleanerContainer ksEngineCleaner;

    /**
     * Indicates whether the current instance of Keystone has been closed.
     */
    private final AtomicBoolean hasBeenClosed;

    /**
     * The memory retention of the symbol resolver callback, in order to prevent the garbage collector
     * to free up the callback, that would result in a crash, as the native library still has a reference to it.
     */
    private SymbolResolverCallback symbolResolverCallback;

    /**
     * Initializes a new instance of the class {@link Keystone}.
     * <p>
     * Some architectures are not supported. Use the static method {@link #isArchitectureSupported} to determine
     * whether the engine support the architecture.
     *
     * @param architecture The architecture of the code generated by Keystone.
     * @param mode         The mode type.
     * @throws OpenFailedKeystoneException if the Keystone library cannot be opened properly.
     */
    public Keystone(KeystoneArchitecture architecture, KeystoneMode mode) {
        ksEngine = initializeKeystoneEngine(architecture, mode);
        ksEngineCleaner = initializeKeystoneCleanerContainer();
        hasBeenClosed = new AtomicBoolean(false);
    }

    /**
     * Determines whether the given architecture is supported by Keystone.
     *
     * @param architecture The architecture type to check.
     * @return The return value is {@code true} if the architecture is supported, otherwise {@code false}.
     */
    public static boolean isArchitectureSupported(KeystoneArchitecture architecture) {
        return DirectMappingKeystoneNative.ks_arch_supported(architecture);
    }

    /**
     * Opens an handle of Keystone.
     *
     * @param architecture The architecture of the code generated by Keystone.
     * @param mode         The mode type.
     * @return The return value is a pointer to the handle of Keystone.
     * @throws OpenFailedKeystoneException if the Keystone library cannot be opened properly.
     */
    private Pointer initializeKeystoneEngine(KeystoneArchitecture architecture, KeystoneMode mode) {
        var pointerToEngine = new PointerByReference();
        var openResult = DirectMappingKeystoneNative.ks_open(architecture, mode, pointerToEngine);

        if (openResult != KeystoneError.Ok) {
            throw new OpenFailedKeystoneException(openResult);
        }

        return pointerToEngine.getValue();
    }

    /**
     * Initializes the cleaner object, that is going to close the native handle of Keystone if
     * the instance is garbage collected.
     *
     * @return The return value is a cleaner container.
     */
    private CleanerContainer initializeKeystoneCleanerContainer() {
        return new KeystoneCleanerContainer(ksEngine);
    }

    /**
     * Assembles a string that contains assembly code.
     *
     * @param assembly The assembly instructions. Use ; or \n to separate statements.
     * @return The return value is the machine code of the assembly instructions.
     * @throws AssembleFailedKeystoneException if the assembly code cannot be assembled properly.
     */
    public KeystoneEncoded assemble(String assembly) {
        return assemble(assembly, 0);
    }

    /**
     * Assembles a string that contains assembly code, located at a given address location.
     *
     * @param assembly The assembly instructions. Use ; or \n to separate statements.
     * @param address  The address of the first assembly instruction.
     * @return The return value is the machine code of the assembly instructions.
     * @throws AssembleFailedKeystoneException if the assembly code cannot be assembled properly.
     */
    public KeystoneEncoded assemble(String assembly, int address) {
        var pointerToMachineCodeBuffer = new PointerByReference();
        var pointerToMachineCodeSize = new IntByReference();
        var pointerToNumberOfStatements = new IntByReference();

        var result = DirectMappingKeystoneNative.ks_asm(ksEngine, assembly, address, pointerToMachineCodeBuffer,
                pointerToMachineCodeSize, pointerToNumberOfStatements);

        if (result != 0) {
            var errorCode = DirectMappingKeystoneNative.ks_errno(ksEngine);
            throw new AssembleFailedKeystoneException(errorCode, assembly);
        }

        var machineCodeBuffer = pointerToMachineCodeBuffer.getValue();
        var machineCode = machineCodeBuffer.getByteArray(0, pointerToMachineCodeSize.getValue());

        DirectMappingKeystoneNative.ks_free(machineCodeBuffer);

        return new KeystoneEncoded(machineCode, address, pointerToNumberOfStatements.getValue());
    }

    /**
     * Assembles a string that contains assembly code.
     *
     * @param assembly A collection of assembly instructions.
     * @return The return value is the machine code of the assembly instructions.
     * @throws AssembleFailedKeystoneException if the assembly code cannot be assembled properly.
     */
    public KeystoneEncoded assemble(Iterable<String> assembly) {
        return assemble(assembly, 0);
    }

    /**
     * Assembles a string that contains assembly code, located at a given address location.
     *
     * @param assembly A collection of assembly instructions.
     * @param address  The address of the first assembly instruction.
     * @return The return value is the machine code of the assembly instructions.
     * @throws AssembleFailedKeystoneException if the assembly code cannot be assembled properly.
     */
    public KeystoneEncoded assemble(Iterable<String> assembly, int address) {
        return assemble(String.join(";", assembly), address);
    }

    /**
     * Gets the major and minor version numbers.
     *
     * @return The returned value is an instance of the class {@link Version}, containing the major and minor version numbers.
     */
    public Version version() {
        var major = new IntByReference();
        var minor = new IntByReference();

        DirectMappingKeystoneNative.ks_version(major, minor);

        return new Version(major.getValue(), minor.getValue());
    }

    /**
     * Sets the syntax of the assembly code used in this instance of Keystone.
     *
     * @param syntax The syntax of the assembly code.
     * @throws SetOptionFailedKeystoneException if the syntax is not supported.
     * @throws SetOptionFailedKeystoneException if the pair of type and value is not valid.
     */
    public void setAssemblySyntax(KeystoneOptionValue.KeystoneOptionSyntax syntax) {
        setOption(KeystoneOptionType.Syntax, syntax.value());
    }

    /**
     * Sets an option for Keystone engine at runtime using a not-strongly typed option value.
     * <p>
     * It is suggested to prefer the methods {@link #setAssemblySyntax(KeystoneOptionValue.KeystoneOptionSyntax)},
     * {@link #setSymbolResolver(SymbolResolverCallback)} or {@link #unsetSymbolResolver()}.
     *
     * @param type  The type of the option.
     * @param value The value of the option.
     * @throws SetOptionFailedKeystoneException if the pair of type and value is not valid.
     */
    public void setOption(KeystoneOptionType type, int value) {
        var result = DirectMappingKeystoneNative.ks_option(ksEngine, type, value);

        if (result != KeystoneError.Ok) {
            throw new SetOptionFailedKeystoneException(result, type, value);
        }
    }

    /**
     * Sets a symbol resolver callback, to resolve unrecognized symbols encountered in the assembly code.
     * <p>
     * If the method is called many times, only the last callback will be triggered.
     *
     * @param callback The symbol resolver callback.
     * @throws SetOptionFailedKeystoneException if the pair of type and value is not valid.
     */
    public void setSymbolResolver(SymbolResolverCallback callback) {
        symbolResolverCallback = callback;
        DirectMappingKeystoneNative.ks_option(ksEngine, KeystoneOptionType.SymbolResolver, callback);
    }

    /**
     * Unsets the current symbol resolver. The symbol resolver instance can be freely collected afterwards.
     *
     * @throws SetOptionFailedKeystoneException if the pair of type and value is not valid.
     */
    public void unsetSymbolResolver() {
        DirectMappingKeystoneNative.ks_option(ksEngine, KeystoneOptionType.SymbolResolver, 0);
        symbolResolverCallback = null;
    }

    /**
     * Closes this resource, relinquishing any underlying resources.
     * This method is invoked automatically on objects managed by the
     * {@code try}-with-resources statement.
     * <p>
     * The call to this method is thread-safe.
     */
    @Override
    public void close() {
        var hasBeenAlreadyClosed = hasBeenClosed.getAndSet(true);

        if (!hasBeenAlreadyClosed) {
            ksEngineCleaner.close();
        }
    }
}