diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md index 761207de4..c5fdab8f7 100644 --- a/.github/CONTRIBUTING.md +++ b/.github/CONTRIBUTING.md @@ -101,7 +101,7 @@ Enhancement suggestions are tracked as [GitHub issues](https://github.com/Aikido ### Your First Code Contribution - clone the repository to your local machine -- run `$ make build` to build a `dist/` folder which includes binaries, agent.jar file and the agent_api.jar file +- run `$ make build` to build a `dist/` folder which includes the agent.jar and agent_api.jar files - run `$ make mock_init` to build and start the mock aikido server (uses docker) - run `$ make test` to test the library with JUnit 5 - run `$ make clean` to clean up the repository diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index 23f15a3df..f06973db4 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -23,7 +23,7 @@ jobs: working-directory: ./ run: | chmod +x gradlew - make binaries + make wasm make build - name: Upload build artifacts diff --git a/.github/workflows/end2end.yml b/.github/workflows/end2end.yml index 66c9dac2d..05320cf23 100644 --- a/.github/workflows/end2end.yml +++ b/.github/workflows/end2end.yml @@ -23,7 +23,7 @@ jobs: working-directory: ./ run: | chmod +x gradlew - make binaries + make wasm make build - name: Upload build artifacts diff --git a/.github/workflows/gradle-tests.yml b/.github/workflows/gradle-tests.yml index 15e7d71b3..f2afccba6 100644 --- a/.github/workflows/gradle-tests.yml +++ b/.github/workflows/gradle-tests.yml @@ -26,7 +26,7 @@ jobs: working-directory: ./ run: | chmod +x gradlew - make binaries + make wasm make build - name: Upload build artifacts diff --git a/.github/workflows/opentel.yml b/.github/workflows/opentel.yml index e9e5aae41..78bd9d744 100644 --- a/.github/workflows/opentel.yml +++ b/.github/workflows/opentel.yml @@ -23,7 +23,7 @@ jobs: working-directory: ./ run: | chmod +x gradlew - make binaries + make wasm make build - name: Upload build artifacts diff --git a/.github/workflows/qa-tests.yml b/.github/workflows/qa-tests.yml index ce01cadd6..6e6742743 100644 --- a/.github/workflows/qa-tests.yml +++ b/.github/workflows/qa-tests.yml @@ -33,7 +33,7 @@ jobs: working-directory: ./firewall-java run: | chmod +x gradlew - make binaries + make wasm make build # Move the build jars to demo app diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 662236c90..f35d93d08 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -31,8 +31,11 @@ jobs: with: java-version: '21' distribution: 'adopt' - - name: Download binaries & Build with Gradle - run: chmod +x gradlew && make binaries && make build + - name: Build with Gradle + run: | + chmod +x gradlew + make wasm + make build - name: Create zip and tar.gz files of the build run: | mv dist/ zen/ diff --git a/.github/workflows/smoke-tests.yml b/.github/workflows/smoke-tests.yml index 5696201c8..be12bb159 100644 --- a/.github/workflows/smoke-tests.yml +++ b/.github/workflows/smoke-tests.yml @@ -28,13 +28,13 @@ jobs: working-directory: ./ run: | chmod +x gradlew - make binaries + make wasm make build - - name: Run RustSQLInterfaceTest + - name: Run WasmSQLInterfaceTest working-directory: ./ run: | - ./gradlew test --tests "vulnerabilities.RustSQLInterfaceTest" --info + ./gradlew test --tests "vulnerabilities.WasmSQLInterfaceTest" --info smoke-test-musl: name: Smoke Test (${{matrix.image}}, Java 21) @@ -56,12 +56,12 @@ jobs: working-directory: ./ run: | chmod +x gradlew - make binaries + make wasm make build - - name: Run RustSQLInterfaceTest on ${{matrix.image}} + - name: Run WasmSQLInterfaceTest on ${{matrix.image}} run: | docker run --rm -v "$(pwd):/app" -w /app ${{matrix.image}} sh -c " chmod +x gradlew && \ - AIKIDO_DEBUG=true ./gradlew test --tests 'vulnerabilities.RustSQLInterfaceTest' --info + AIKIDO_DEBUG=true ./gradlew test --tests 'vulnerabilities.WasmSQLInterfaceTest' --info " diff --git a/.github/workflows/test-ddtrace.yml b/.github/workflows/test-ddtrace.yml index 87060d047..2db15e344 100644 --- a/.github/workflows/test-ddtrace.yml +++ b/.github/workflows/test-ddtrace.yml @@ -23,7 +23,7 @@ jobs: working-directory: ./ run: | chmod +x gradlew - make binaries + make wasm make build - name: Upload build artifacts uses: actions/upload-artifact@v4 diff --git a/Makefile b/Makefile index 9822afeb3..0ff6ce559 100644 --- a/Makefile +++ b/Makefile @@ -2,11 +2,34 @@ clean: rm -rf dist/ ./gradlew clean -build: clean check_binaries - mkdir -p dist/ +ZEN_INTERNALS_VERSION = v0.1.60 +WASM_BASE_URL = https://github.com/AikidoSec/zen-internals/releases/download/$(ZEN_INTERNALS_VERSION) +WASM_RESOURCE_DIR = agent_api/src/main/resources + +.PHONY: wasm download-wasm check-wasm +wasm: download-wasm check-wasm + +download-wasm: + mkdir -p $(WASM_RESOURCE_DIR) + @set -e; \ + tmp_dir=$$(mktemp -d); \ + trap 'rm -rf "$$tmp_dir"' 0; \ + curl -fL -o "$$tmp_dir/zen_internals.wasm" $(WASM_BASE_URL)/libzen_internals.wasm; \ + curl -fL -o "$$tmp_dir/checksum" $(WASM_BASE_URL)/libzen_internals.wasm.sha256sum; \ + sed 's/libzen_internals\.wasm/zen_internals.wasm/' "$$tmp_dir/checksum" > "$$tmp_dir/zen_internals.wasm.sha256sum"; \ + mv "$$tmp_dir/zen_internals.wasm" $(WASM_RESOURCE_DIR)/zen_internals.wasm; \ + mv "$$tmp_dir/zen_internals.wasm.sha256sum" $(WASM_RESOURCE_DIR)/zen_internals.wasm.sha256sum + +check-wasm: + @expected=$$(awk '{print $$1}' $(WASM_RESOURCE_DIR)/zen_internals.wasm.sha256sum); \ + actual=$$(shasum -a 256 $(WASM_RESOURCE_DIR)/zen_internals.wasm | awk '{print $$1}'); \ + if [ "$$expected" != "$$actual" ]; then \ + echo "WASM checksum mismatch: expected $$expected, got $$actual"; \ + exit 1; \ + fi - @echo "Copying binaries from .cache folder" - cp -r .cache/binaries dist/binaries +build: clean + mkdir -p dist/ ./gradlew agent:shadowJar cp agent/build/libs/agent*-all.jar dist/agent.jar @@ -23,49 +46,13 @@ mock_restart: mock_stop: docker kill mock_core && docker rm mock_core -test: check_binaries +test: AIKIDO_LOG_LEVEL="error" AIKIDO_TOKEN="token" ./gradlew test -cov: check_binaries +cov: AIKIDO_LOG_LEVEL="error" AIKIDO_TOKEN="token" ./gradlew test --rerun-tasks -PcoverageRun jacocoTestReport jacocoTestCoverageVerification -# Binaries : - -BASE_URL = https://github.com/AikidoSec/zen-internals/releases/download/v0.1.60 -FILES = \ - libzen_internals_aarch64-apple-darwin.dylib \ - libzen_internals_aarch64-apple-darwin.dylib.sha256sum \ - libzen_internals_aarch64-unknown-linux-gnu.so \ - libzen_internals_aarch64-unknown-linux-gnu.so.sha256sum \ - libzen_internals_aarch64-unknown-linux-musl.so \ - libzen_internals_aarch64-unknown-linux-musl.so.sha256sum \ - libzen_internals_x86_64-apple-darwin.dylib \ - libzen_internals_x86_64-apple-darwin.dylib.sha256sum \ - libzen_internals_x86_64-pc-windows-gnu.dll \ - libzen_internals_x86_64-pc-windows-gnu.dll.sha256sum \ - libzen_internals_x86_64-unknown-linux-gnu.so \ - libzen_internals_x86_64-unknown-linux-gnu.so.sha256sum \ - libzen_internals_x86_64-unknown-linux-musl.so \ - libzen_internals_x86_64-unknown-linux-musl.so.sha256sum \ - -binaries: binaries_make_dir $(addprefix .cache/binaries/, $(FILES)) -binaries_make_dir: - rm -rf .cache/binaries - mkdir -p .cache/binaries/ -.cache/binaries/%: - @echo "Downloading $*..." - curl -L -o $@ $(BASE_URL)/$* -.PHONY: check_binaries -check_binaries: - @if [ -d ".cache/binaries" ]; then \ - echo "Cache directory exists."; \ - else \ - echo "Cache directory is empty. Running 'make binaries'..."; \ - $(MAKE) binaries; \ - fi - - # Automatic versioning for releases : VERSION_FILES = ./build.gradle ./agent_api/src/main/java/dev/aikido/agent_api/Config.java @@ -80,4 +67,3 @@ replace_version: sed -i.bak "s/1.0-REPLACE-VERSION/$$version/g" $$file; \ rm $$file.bak; \ done; - diff --git a/README.md b/README.md index fa8194501..ff89b0252 100644 --- a/README.md +++ b/README.md @@ -73,26 +73,6 @@ To activate Zen you then just have to add the following `-javaagent` to your Jav ``` java -javaagent:/opt/zen/agent.jar -jar build/myapp.jar ``` -Replace `/opt/zen` with your directory of choice. Keep `agent.jar` together with the `binaries` folder from the -release - if you copy `agent.jar` into a different location or a different Docker build stage, copy the whole -directory, not just the jar. Without `binaries`, SQL injection detection is disabled. - -
-Seeing System::load has been called ... in an unnamed module on Java 22+? - -This warning is safe to ignore for now - SQL injection detection still works, since `--illegal-native-access` -still defaults to `warn`, not `deny`, on Java 25. - -To get rid of it, add `--enable-native-access=ALL-UNNAMED` to your Java command, or set it through the -`JDK_JAVA_OPTIONS` environment variable so you don't have to touch your existing startup command: -``` -java --enable-native-access=ALL-UNNAMED -javaagent:/opt/zen/agent.jar -jar build/myapp.jar -``` -``` -JDK_JAVA_OPTIONS=--enable-native-access=ALL-UNNAMED -``` -
- To use user-blocking and/or rate-limiting features, you will have to include the following Jarfile into your repository ### Gradle Add the following code to your `build.gradle` file. diff --git a/agent/src/main/java/dev/aikido/agent/Agent.java b/agent/src/main/java/dev/aikido/agent/Agent.java index d88d1df54..fbe0514fb 100644 --- a/agent/src/main/java/dev/aikido/agent/Agent.java +++ b/agent/src/main/java/dev/aikido/agent/Agent.java @@ -17,7 +17,7 @@ import static dev.aikido.agent.ByteBuddyInitializer.createAgentBuilder; import static dev.aikido.agent.DaemonStarter.startDaemon; import static dev.aikido.agent.Wrappers.WRAPPERS; -import static dev.aikido.agent_api.vulnerabilities.sql_injection.RustSQLInterface.loadLibrary; +import static dev.aikido.agent_api.vulnerabilities.sql_injection.WasmSQLInterface.initialize; public class Agent { private static final Logger logger = LogManager.getLogger(Agent.class); @@ -38,8 +38,8 @@ public static void premain(String agentArgs, Instrumentation inst) { logger.info("Zen by Aikido v%s starting.", Config.pkgVersion); setAikidoSysProperties(); - // Test loading of zen binaries : - loadLibrary(); + // Load the zen-internals WASM module. + initialize(); ElementMatcher.Junction wrapperTypeDescriptors = ElementMatchers.none(); for(Wrapper wrapper: WRAPPERS) { diff --git a/agent_api/build.gradle b/agent_api/build.gradle index 62a87ebd4..39f07b1aa 100644 --- a/agent_api/build.gradle +++ b/agent_api/build.gradle @@ -12,7 +12,8 @@ jacoco { dependencies { implementation 'com.github.seancfoley:ipaddress:5.5.1' implementation 'com.google.code.gson:gson:2.11.0' - implementation 'com.github.jnr:jnr-ffi:2.2.17' + implementation 'com.dylibso.chicory:runtime:1.7.5' + implementation 'com.dylibso.chicory:compiler:1.7.5' // Junixsocket imports : implementation 'com.kohlschutter.junixsocket:junixsocket-core:2.10.1' implementation 'com.kohlschutter.junixsocket:junixsocket-server:2.10.1' @@ -48,7 +49,6 @@ test { jvmArgs '--add-opens', 'java.base/java.lang=ALL-UNNAMED' jvmArgs '-Dnet.bytebuddy.experimental=true' // Mockito support. - systemProperty 'AIK_agent_dir', "${project.rootDir}/dist" if (project.hasProperty('coverageRun')) { systemProperty 'AIK_INTERNAL_coverage_run', '1' } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/GetBinaryPath.java b/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/GetBinaryPath.java deleted file mode 100644 index e1810c527..000000000 --- a/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/GetBinaryPath.java +++ /dev/null @@ -1,67 +0,0 @@ -package dev.aikido.agent_api.vulnerabilities.sql_injection; - -import dev.aikido.agent_api.helpers.logging.LogManager; -import dev.aikido.agent_api.helpers.logging.Logger; -import jnr.a64asm.INST_CODE; -import jnr.ffi.Library; -import jnr.ffi.LibraryLoader; - -import java.io.BufferedReader; -import java.io.IOException; -import java.io.InputStreamReader; -import java.io.InterruptedIOException; -import java.lang.annotation.Native; - -public final class GetBinaryPath { - private GetBinaryPath() {} - private static final Logger logger = LogManager.getLogger(GetBinaryPath.class); - - public static String getPathForBinary() { - String fileName = getFileName(); - String aikidoDirectory = System.getProperty("AIK_agent_dir"); - if (aikidoDirectory == null) { - return null; - } - return aikidoDirectory + "/binaries/" + fileName; - } - private static String getFileName() { - String os = System.getProperty("os.name").toLowerCase(); - String architecture = System.getProperty("os.arch").toLowerCase(); - StringBuilder fileName = new StringBuilder(); - - fileName.append("libzen_internals_"); // Start of file - - if (architecture.contains("aarch64")) { - fileName.append("aarch64-"); // Add architecture to file name - } else if (architecture.contains("64")) { - fileName.append("x86_64-"); // Add architecture to file name - } else { - fileName.append("x86_64-"); // Default to x86-64 - } - - if (os.contains("win")) { - fileName.append("pc-windows-gnu.dll"); // Windows - } else if (os.contains("mac")) { - fileName.append("apple-darwin.dylib"); // macOS - } else { - // Default to linux - fileName.append(String.format("unknown-linux-%s.so", getLibCVariant())); - } - return fileName.toString(); - } - - public interface Libc { - Libc INSTANCE = LibraryLoader.create(Libc.class).load("c"); - String gnu_get_libc_version(); - } - - private static String getLibCVariant() { - // gnu_get_libc_version only works for systems with gnu installed. - try { - Libc.INSTANCE.gnu_get_libc_version(); - } catch (UnsatisfiedLinkError e) { - return "musl"; - } - return "gnu"; - } -} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/RustSQLInterface.java b/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/RustSQLInterface.java deleted file mode 100644 index 1e14a3816..000000000 --- a/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/RustSQLInterface.java +++ /dev/null @@ -1,77 +0,0 @@ -package dev.aikido.agent_api.vulnerabilities.sql_injection; - -import jnr.ffi.LibraryLoader; -import jnr.ffi.LibraryOption; -import jnr.ffi.annotations.Encoding; -import dev.aikido.agent_api.helpers.logging.LogManager; -import dev.aikido.agent_api.helpers.logging.Logger; - -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.HashMap; -import java.util.Map; -import java.nio.charset.StandardCharsets; - -import static dev.aikido.agent_api.vulnerabilities.sql_injection.GetBinaryPath.getPathForBinary; - -public final class RustSQLInterface { - private RustSQLInterface() {} - - private static final Logger logger = LogManager.getLogger(RustSQLInterface.class); - - @Encoding("UTF-8") - public interface SqlLib { - int detect_sql_injection(String query, long queryLen, String userinput, long userinputLen, int dialect); - } - - public static boolean detectSqlInjection(String query, String userInput, Dialect dialect) { - int dialectInteger = dialect.getDialectInteger(); - try { - SqlLib lib = loadLibrary(); - if (lib != null) { - long queryLen = query != null ? query.getBytes(StandardCharsets.UTF_8).length : 0; - long userInputLen = userInput != null ? userInput.getBytes(StandardCharsets.UTF_8).length : 0; - int result = lib.detect_sql_injection(query, queryLen, userInput, userInputLen, dialectInteger); - return result == 1; - } - } catch (Throwable e) { - logger.trace(e); - } - return false; - } - public static SqlLib loadLibrary() { - String path = getPathForBinary(); - if (path == null) { - logger.error("Could not load zen binaries: AIK_agent_dir is not set. SQL injection detection is disabled."); - return null; - } - Path binariesDir = Path.of(path).getParent(); - if (binariesDir == null || !Files.isDirectory(binariesDir)) { - logger.error("Could not load zen binaries: the 'binaries' directory is missing (%s). Copy it next to " + - "agent.jar - check if a build step (e.g. Docker COPY) left it out. SQL injection detection is disabled.", binariesDir); - return null; - } - if (!Files.exists(Path.of(path))) { - logger.error("Could not load zen binaries: file not found: %s. SQL injection detection is disabled.", path); - return null; - } - Map libraryOptions = new HashMap<>(); - libraryOptions.put(LibraryOption.LoadNow, true); // load immediately instead of lazily (ie on first use) - libraryOptions.put(LibraryOption.IgnoreError, true); // calls shouldn't save last errno after call - - SqlLib library = null; - try { - library = LibraryLoader.loadLibrary(SqlLib.class, libraryOptions, path); - } catch (Throwable e) { - String os = System.getProperty("os.name").toLowerCase(); - String architecture = System.getProperty("os.arch").toLowerCase(); - logger.error("Failed to load Zen Internals (%s, %s)", os, architecture); - throw e; - } - - if (library == null) { - logger.error("Failed to load zen binaries."); - } - return library; - } -} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/SqlDetector.java b/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/SqlDetector.java index 156ec2727..a882a0bff 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/SqlDetector.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/SqlDetector.java @@ -39,7 +39,7 @@ public static boolean detectSqlInjection(String query, String userInput, Dialect if (shouldReturnEarly(queryLower, userInputNormalized)) { return false; } - return RustSQLInterface.detectSqlInjection(queryLower, userInputNormalized, dialect); + return WasmSQLInterface.detectSqlInjection(queryLower, userInputNormalized, dialect); } /** * Input : Lowercased query and user_input. @@ -71,4 +71,3 @@ public static boolean shouldReturnEarly(String query, String userInput) { return pattern.matcher(cleanedInputForList).matches(); } } - diff --git a/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/WasmInstance.java b/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/WasmInstance.java new file mode 100644 index 000000000..8ca9b500b --- /dev/null +++ b/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/WasmInstance.java @@ -0,0 +1,69 @@ +package dev.aikido.agent_api.vulnerabilities.sql_injection; + +import com.dylibso.chicory.runtime.ExportFunction; +import com.dylibso.chicory.runtime.Instance; +import com.dylibso.chicory.runtime.Memory; + +import java.nio.charset.StandardCharsets; + +final class WasmInstance { + private final ExportFunction alloc; + private final ExportFunction free; + private final ExportFunction detectSql; + private final Memory memory; + + WasmInstance(Instance instance) { + alloc = instance.export("wasm_alloc"); + free = instance.export("wasm_free"); + detectSql = instance.export("detect_sql_injection"); + memory = instance.memory(); + } + + int detect(String query, String userInput, int dialect) { + byte[] queryBytes = query == null + ? new byte[0] + : query.getBytes(StandardCharsets.UTF_8); + byte[] userInputBytes = userInput == null + ? new byte[0] + : userInput.getBytes(StandardCharsets.UTF_8); + try (Allocation queryAllocation = allocate(queryBytes.length); + Allocation userInputAllocation = allocate(userInputBytes.length)) { + memory.write(queryAllocation.pointer, queryBytes); + memory.write(userInputAllocation.pointer, userInputBytes); + long[] result = detectSql.apply( + Integer.toUnsignedLong(queryAllocation.pointer), + queryBytes.length, + Integer.toUnsignedLong(userInputAllocation.pointer), + userInputBytes.length, + dialect + ); + if (result.length == 0) { + throw new IllegalStateException("detect_sql_injection returned no result"); + } + return Math.toIntExact(result[0]); + } + } + + private Allocation allocate(int length) { + long[] result = alloc.apply(length); + if (result.length == 0 || result[0] > 0xffff_ffffL) { + throw new IllegalStateException("invalid WASM allocation result"); + } + return new Allocation((int) result[0], length); + } + + private final class Allocation implements AutoCloseable { + private final int pointer; + private final int length; + + private Allocation(int pointer, int length) { + this.pointer = pointer; + this.length = length; + } + + @Override + public void close() { + free.apply(Integer.toUnsignedLong(pointer), length); + } + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/WasmSQLInterface.java b/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/WasmSQLInterface.java new file mode 100644 index 000000000..a8626724b --- /dev/null +++ b/agent_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/WasmSQLInterface.java @@ -0,0 +1,186 @@ +package dev.aikido.agent_api.vulnerabilities.sql_injection; + +import com.dylibso.chicory.compiler.InterpreterFallback; +import com.dylibso.chicory.compiler.MachineFactoryCompiler; +import com.dylibso.chicory.runtime.Instance; +import com.dylibso.chicory.runtime.Machine; +import com.dylibso.chicory.wasm.Parser; +import com.dylibso.chicory.wasm.WasmModule; +import dev.aikido.agent_api.helpers.logging.LogManager; +import dev.aikido.agent_api.helpers.logging.Logger; + +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.ArrayDeque; +import java.util.HexFormat; +import java.util.Set; +import java.util.function.Function; + +public final class WasmSQLInterface { + private static final Logger logger = LogManager.getLogger(WasmSQLInterface.class); + private static final String WASM_RESOURCE = "zen_internals.wasm"; + private static final String CHECKSUM_RESOURCE = "zen_internals.wasm.sha256sum"; + private static final int MAX_IDLE_INSTANCES = 10; + // These functions exceed Java's 65,535-byte method limit, so run them in the interpreter and fail if any others cannot compile. + private static final Set INTERPRETED_FUNCTIONS = Set.of(865, 3165); + + private static volatile RuntimeState runtime; + private static volatile Throwable initializationFailure; + + private WasmSQLInterface() {} + + public static boolean initialize() { + try { + getRuntime(); + return true; + } catch (Throwable e) { + logger.error("Failed to initialize zen-internals WASM: %s", e.getMessage()); + logger.trace(e); + return false; + } + } + + public static boolean detectSqlInjection(String query, String userInput, Dialect dialect) { + RuntimeState current; + WasmInstance instance; + try { + current = getRuntime(); + instance = current.pool.acquire(); + } catch (Throwable e) { + logger.trace(e); + // If it cannot initialize or create an instance, let the query proceed. + return false; + } + + boolean reusable = false; + try { + int result = instance.detect(query, userInput, dialect.getDialectInteger()); + reusable = true; + return result == 1; + } catch (Throwable e) { + logger.trace(e); + // If zen-internals fails while checking the query, let the query proceed. + return false; + } finally { + if (reusable) { + current.pool.release(instance); + } + } + } + + private static RuntimeState getRuntime() { + RuntimeState current = runtime; + if (current != null) { + return current; + } + Throwable failure = initializationFailure; + if (failure != null) { + throw new IllegalStateException("zen-internals WASM initialization failed", failure); + } + + synchronized (WasmSQLInterface.class) { + current = runtime; + if (current != null) { + return current; + } + failure = initializationFailure; + if (failure != null) { + throw new IllegalStateException("zen-internals WASM initialization failed", failure); + } + try { + current = createRuntime(); + runtime = current; + return current; + } catch (Throwable e) { + initializationFailure = e; + throw new IllegalStateException("zen-internals WASM initialization failed", e); + } + } + } + + private static RuntimeState createRuntime() throws IOException { + byte[] wasm = readResource(WASM_RESOURCE); + String[] checksumFields = new String(readResource(CHECKSUM_RESOURCE), StandardCharsets.UTF_8) + .trim() + .split("\\s+", 2); + if (checksumFields[0].isEmpty()) { + throw new IllegalStateException("invalid zen-internals WASM checksum"); + } + String expectedChecksum = checksumFields[0]; + String actualChecksum = sha256(wasm); + if (!actualChecksum.equals(expectedChecksum)) { + throw new IllegalStateException( + "zen-internals WASM checksum mismatch: expected " + expectedChecksum + + ", got " + actualChecksum); + } + + WasmModule module = Parser.parse(wasm); + Function machineFactory = + MachineFactoryCompiler.builder(module) + .withInterpreterFallback(InterpreterFallback.FAIL) + .withInterpretedFunctions(INTERPRETED_FUNCTIONS) + .compile(); + return new RuntimeState(module, machineFactory); + } + + private static byte[] readResource(String name) throws IOException { + ClassLoader classLoader = WasmSQLInterface.class.getClassLoader(); + try (InputStream stream = classLoader.getResourceAsStream(name)) { + if (stream == null) { + throw new IOException("resource not found: " + name); + } + return stream.readAllBytes(); + } + } + + private static String sha256(byte[] bytes) { + try { + return HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(bytes)); + } catch (NoSuchAlgorithmException e) { + throw new AssertionError("SHA-256 is required by every supported JDK", e); + } + } + + private static final class RuntimeState { + private final WasmModule module; + private final Function machineFactory; + private final InstancePool pool; + + private RuntimeState(WasmModule module, Function machineFactory) { + this.module = module; + this.machineFactory = machineFactory; + this.pool = new InstancePool(this); + } + + private WasmInstance newInstance() { + Instance instance = Instance.builder(module) + .withMachineFactory(machineFactory) + .build(); + return new WasmInstance(instance); + } + } + + private static final class InstancePool { + private final RuntimeState runtimeState; + private final ArrayDeque idleInstances = new ArrayDeque<>(); + + private InstancePool(RuntimeState runtimeState) { + this.runtimeState = runtimeState; + } + + private synchronized WasmInstance acquire() { + WasmInstance instance = idleInstances.pollLast(); + return instance != null ? instance : runtimeState.newInstance(); + } + + private synchronized void release(WasmInstance instance) { + if (idleInstances.size() < MAX_IDLE_INSTANCES) { + idleInstances.addLast(instance); + } + } + } + +} diff --git a/agent_api/src/main/resources/zen_internals.wasm b/agent_api/src/main/resources/zen_internals.wasm new file mode 100644 index 000000000..e08883503 Binary files /dev/null and b/agent_api/src/main/resources/zen_internals.wasm differ diff --git a/agent_api/src/main/resources/zen_internals.wasm.sha256sum b/agent_api/src/main/resources/zen_internals.wasm.sha256sum new file mode 100644 index 000000000..b044c7238 --- /dev/null +++ b/agent_api/src/main/resources/zen_internals.wasm.sha256sum @@ -0,0 +1 @@ +61a3913c7b5c0affaf21372158bdf1808e385857212c96d81f241504df51205d zen_internals.wasm diff --git a/agent_api/src/test/java/vulnerabilities/RustSQLInterfaceTest.java b/agent_api/src/test/java/vulnerabilities/RustSQLInterfaceTest.java deleted file mode 100644 index 1b98e8315..000000000 --- a/agent_api/src/test/java/vulnerabilities/RustSQLInterfaceTest.java +++ /dev/null @@ -1,18 +0,0 @@ -package vulnerabilities; - -import dev.aikido.agent_api.vulnerabilities.sql_injection.Dialect; -import dev.aikido.agent_api.vulnerabilities.sql_injection.RustSQLInterface; -import org.junit.jupiter.api.Test; - -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; - -public class RustSQLInterfaceTest { - @Test - public void testItWorks() { - boolean injectionResult = RustSQLInterface.detectSqlInjection("SELECT * FROM table;", "table;", new Dialect("postgresql")); - assertTrue(injectionResult); - injectionResult = RustSQLInterface.detectSqlInjection("SELECT * FROM table;", "table", new Dialect("postgresql")); - assertFalse(injectionResult); - } -} diff --git a/agent_api/src/test/java/vulnerabilities/WasmSQLInterfaceTest.java b/agent_api/src/test/java/vulnerabilities/WasmSQLInterfaceTest.java new file mode 100644 index 000000000..3edefc7b4 --- /dev/null +++ b/agent_api/src/test/java/vulnerabilities/WasmSQLInterfaceTest.java @@ -0,0 +1,63 @@ +package vulnerabilities; + +import dev.aikido.agent_api.vulnerabilities.sql_injection.Dialect; +import dev.aikido.agent_api.vulnerabilities.sql_injection.WasmSQLInterface; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class WasmSQLInterfaceTest { + @Test + public void detectsSqlInjection() { + Dialect postgresql = new Dialect("postgresql"); + assertTrue(WasmSQLInterface.initialize()); + + assertTrue( + WasmSQLInterface.detectSqlInjection( + "SELECT * FROM users WHERE id = '1' OR 1=1", + "1' OR 1=1", + postgresql + ) + ); + assertFalse( + WasmSQLInterface.detectSqlInjection( + "SELECT * FROM users WHERE id = '1'", + "1", + postgresql + ) + ); + } + + @Test + public void supportsConcurrentCallsWithIndependentInstances() throws Exception { + assertTrue(WasmSQLInterface.initialize()); + ExecutorService executor = Executors.newFixedThreadPool(8); + try { + List> calls = new ArrayList<>(); + for (int i = 0; i < 64; i++) { + calls.add(() -> WasmSQLInterface.detectSqlInjection( + "SELECT * FROM users WHERE id = '1' OR 1=1", + "1' OR 1=1", + new Dialect("mysql"))); + } + for (Future call : executor.invokeAll(calls)) { + assertTrue(get(call)); + } + } finally { + executor.shutdownNow(); + } + } + + private static boolean get(Future call) throws ExecutionException, InterruptedException { + return call.get(); + } +} diff --git a/agent_api/src/test/java/vulnerabilities/sql_injection/GetBinaryPathTest.java b/agent_api/src/test/java/vulnerabilities/sql_injection/GetBinaryPathTest.java deleted file mode 100644 index 8132799b4..000000000 --- a/agent_api/src/test/java/vulnerabilities/sql_injection/GetBinaryPathTest.java +++ /dev/null @@ -1,124 +0,0 @@ -package vulnerabilities.sql_injection; - -import dev.aikido.agent_api.vulnerabilities.sql_injection.GetBinaryPath; -import org.junit.jupiter.api.*; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNull; - -public class GetBinaryPathTest { - @BeforeAll - public static void copyStart() { - System.setProperty("copy.AIK_agent_dir", System.getProperty("AIK_agent_dir")); - System.setProperty("copy.os.name", System.getProperty("os.name")); - System.setProperty("copy.os.arch", System.getProperty("os.arch")); - } - @AfterAll - public static void cleanup() { - System.setProperty("AIK_agent_dir", System.getProperty("copy.AIK_agent_dir")); - System.setProperty("os.name", System.getProperty("copy.os.name")); - System.setProperty("os.arch", System.getProperty("copy.os.arch")); - } - - @BeforeEach - public void setUp() { - // Clear system properties before each test - System.clearProperty("AIK_agent_dir"); - System.clearProperty("os.name"); - System.clearProperty("os.arch"); - } - - @Test - public void testGetPathForBinary_WithWindows64() { - System.setProperty("AIK_agent_dir", "/path/to/agent"); - System.setProperty("os.name", "Windows 10"); - System.setProperty("os.arch", "amd64"); - - String expectedPath = "/path/to/agent/binaries/libzen_internals_x86_64-pc-windows-gnu.dll"; - String actualPath = GetBinaryPath.getPathForBinary(); - assertEquals(expectedPath, actualPath); - } - - @Test - public void testGetPathForBinary_WithWindowsARM64() { - System.setProperty("AIK_agent_dir", "/path/to/agent"); - System.setProperty("os.name", "Windows 10"); - System.setProperty("os.arch", "aarch64"); - - String expectedPath = "/path/to/agent/binaries/libzen_internals_aarch64-pc-windows-gnu.dll"; - String actualPath = GetBinaryPath.getPathForBinary(); - assertEquals(expectedPath, actualPath); - } - - @Test - public void testGetPathForBinary_WithMac() { - System.setProperty("AIK_agent_dir", "/path/to/agent"); - System.setProperty("os.name", "Mac OS X"); - System.setProperty("os.arch", "x86_64"); - - String expectedPath = "/path/to/agent/binaries/libzen_internals_x86_64-apple-darwin.dylib"; - String actualPath = GetBinaryPath.getPathForBinary(); - assertEquals(expectedPath, actualPath); - } - - @Test - public void testGetPathForBinary_WithLinux() { - System.setProperty("AIK_agent_dir", "/path/to/agent"); - System.setProperty("os.name", "Linux"); - System.setProperty("os.arch", "x86_64"); - - String expectedPath = "/path/to/agent/binaries/libzen_internals_x86_64-unknown-linux-gnu.so"; - String actualPath = GetBinaryPath.getPathForBinary(); - assertEquals(expectedPath, actualPath); - } - - @Test - public void testGetPathForBinary_WithLinuxARM64() { - System.setProperty("AIK_agent_dir", "/path/to/agent"); - System.setProperty("os.name", "Linux"); - System.setProperty("os.arch", "aarch64"); - - String expectedPath = "/path/to/agent/binaries/libzen_internals_aarch64-unknown-linux-gnu.so"; - String actualPath = GetBinaryPath.getPathForBinary(); - assertEquals(expectedPath, actualPath); - } - - @Test - public void testGetPathForBinary_WithNoAgentDir() { - System.setProperty("os.name", "Linux"); - System.setProperty("os.arch", "x86_64"); - - String actualPath = GetBinaryPath.getPathForBinary(); - assertNull(actualPath); - } - - @Test - public void testGetPathForBinary_WithEmptyAgentDir() { - System.setProperty("AIK_agent_dir", ""); - System.setProperty("os.name", "Linux"); - System.setProperty("os.arch", "x86_64"); - - String actualPath = GetBinaryPath.getPathForBinary(); - assertEquals("/binaries/libzen_internals_x86_64-unknown-linux-gnu.so", actualPath); - } - @Test - public void testGetPathForBinary_WithUnknownOS() { - System.setProperty("AIK_agent_dir", "/path/to/agent"); - System.setProperty("os.name", "Unknown OS"); - System.setProperty("os.arch", "x86_64"); - - String expectedPath = "/path/to/agent/binaries/libzen_internals_x86_64-unknown-linux-gnu.so"; - String actualPath = GetBinaryPath.getPathForBinary(); - assertEquals(expectedPath, actualPath); - } - @Test - public void testGetPathForBinary_WithUnknownArchitecture() { - System.setProperty("AIK_agent_dir", "/path/to/agent"); - System.setProperty("os.name", "Linux"); - System.setProperty("os.arch", "unknown-arch"); - - String expectedPath = "/path/to/agent/binaries/libzen_internals_x86_64-unknown-linux-gnu.so"; - String actualPath = GetBinaryPath.getPathForBinary(); - assertEquals(expectedPath, actualPath); - } -} \ No newline at end of file