Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history 
Signed-off-by: Edoardo Vacchi <[email protected]>
  • Loading branch information
@evacchi
evacchi committed Dec 24, 2024
1 parent 0a30ff2 commit d5e23db
Show file tree
Hide file tree
Showing 6 changed files with 248 additions and 13 deletions.
15 changes: 14 additions & 1 deletion pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,8 @@
<maven-source-plugin.version>3.2.1</maven-source-plugin.version>
<maven-gpg-plugin.version>3.2.7</maven-gpg-plugin.version>
<jimfs.version>1.3.0</jimfs.version>
<jakarta.json-api.version>2.1.3</jakarta.json-api.version>
<jakarta.json.version>1.1.7</jakarta.json.version>
</properties>

<dependencies>
Expand Down Expand Up @@ -102,7 +104,18 @@
<version>${jimfs.version}</version>
<scope>test</scope>
</dependency>

<dependency>
<groupId>jakarta.json</groupId>
<artifactId>jakarta.json-api</artifactId>
<version>${jakarta.json-api.version}</version>
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.eclipse.parsson</groupId>
<artifactId>jakarta.json</artifactId>
<version>${jakarta.json.version}</version>
<scope>test</scope>
</dependency>
</dependencies>

<profiles>
Expand Down
1 change: 0 additions & 1 deletion src/main/java/org/extism/chicory/sdk/ChicoryModule.java
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
package org.extism.chicory.sdk;

import com.dylibso.chicory.experimental.aot.AotMachine;
import com.dylibso.chicory.runtime.Instance;
import com.dylibso.chicory.wasm.Parser;
import com.dylibso.chicory.wasm.WasmModule;
Expand Down
198 changes: 189 additions & 9 deletions src/main/java/org/extism/chicory/sdk/HostEnv.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,23 @@
import com.dylibso.chicory.log.Logger;
import com.dylibso.chicory.runtime.HostFunction;
import com.dylibso.chicory.runtime.Instance;
import com.dylibso.chicory.wasm.types.Value;
import com.dylibso.chicory.wasm.types.ValueType;

import jakarta.json.Json;
import jakarta.json.JsonObject;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import static com.dylibso.chicory.wasm.types.Value.i64;

public class HostEnv {

private final Kernel kernel;
Expand All @@ -22,14 +28,16 @@ public class HostEnv {
private final Log log;
private final Var var;
private final Config config;
private final Http http;

public HostEnv(Kernel kernel, Map<String, String> config, Logger logger) {
public HostEnv(Kernel kernel, Map<String, String> config, String[] allowedHosts, Logger logger) {
this.kernel = kernel;
this.memory = new Memory();
this.logger = logger;
this.config = new Config(config);
this.var = new Var();
this.log = new Log();
this.http = new Http(allowedHosts);
}

public Log log() {
Expand All @@ -44,12 +52,17 @@ public Config config() {
return config;
}

public Http http() {
return http;
}

public HostFunction[] toHostFunctions() {
return concat(
kernel.toHostFunctions(),
log.toHostFunctions(),
var.toHostFunctions(),
config.toHostFunctions());
config.toHostFunctions(),
http.toHostFunctions());
}

private HostFunction[] concat(HostFunction[]... hfs) {
Expand Down Expand Up @@ -108,7 +121,8 @@ long writeString(String s) {
}

public class Log {
private Log(){}
private Log() {
}

public void log(LogLevel level, String message) {
logger.log(level.toChicoryLogLevel(), message, null);
Expand Down Expand Up @@ -156,9 +170,10 @@ HostFunction[] toHostFunctions() {
}

public class Var {
private final Map<String, byte[]> vars = new ConcurrentHashMap<>();
private final Map<String, byte[]> vars = new ConcurrentHashMap<>();

private Var() {}
private Var() {
}

public byte[] get(String key) {
return vars.get(key);
Expand Down Expand Up @@ -246,7 +261,172 @@ HostFunction[] toHostFunctions() {

}

public class Http {
private final HostPattern[] hostPatterns;
HttpClient httpClient;
HttpResponse<byte[]> lastResponse;

public Http(String[] allowedHosts) {
this.hostPatterns = new HostPattern[allowedHosts.length];
for (int i = 0; i < allowedHosts.length; i++) {
this.hostPatterns[i] = new HostPattern(allowedHosts[i]);
}
}

public HttpClient httpClient() {
if (httpClient == null) {
httpClient = HttpClient.newHttpClient();
}
return httpClient;
}

long[] request(Instance instance, long... args) {
var result = new long[1];

var requestOffset = args[0];
var bodyOffset = args[1];

var requestJson = memory().readBytes(requestOffset);
kernel.free.apply(requestOffset);

byte[] requestBody;
if (bodyOffset == 0) {
requestBody = new byte[0];
} else {
requestBody = memory().readBytes(bodyOffset);
kernel.free.apply(bodyOffset);
}

var request = Json.createReader(new ByteArrayInputStream(requestJson))
.readObject();

var method = request.getJsonString("method").getString();
var uri = URI.create(request.getJsonString("url").getString());
var headers = request.getJsonObject("headers");

Map<String, String> headersMap = new HashMap<>();
for (var key : headers.keySet()) {
headersMap.put(key, headers.getString(key));
}

byte[] body = request(method, uri, headersMap, requestBody);
if (body.length == 0) {
result[0] = 0;
} else {
result[0] = memory().writeBytes(body);
}

return result;
}

byte[] request(String method, URI uri, Map<String, String> headers, byte[] requestBody) {
HttpRequest.BodyPublisher bodyPublisher;
if (requestBody.length == 0) {
bodyPublisher = HttpRequest.BodyPublishers.noBody();
} else {
bodyPublisher = HttpRequest.BodyPublishers.ofByteArray(requestBody);
}

var host = uri.getHost();
if (Arrays.stream(hostPatterns).anyMatch(p -> !p.matches(host))) {
throw new ExtismException(String.format("HTTP request to '%s' is not allowed", host));
}

var reqBuilder = HttpRequest.newBuilder().uri(uri);
for (var key : headers.keySet()) {
reqBuilder.header(key, headers.get(key));
}

var req = reqBuilder.method(method, bodyPublisher).build();

try {
this.lastResponse =
httpClient().send(req, HttpResponse.BodyHandlers.ofByteArray());
return lastResponse.body();
} catch (IOException | InterruptedException e) {
// FIXME gracefully handle the interruption
throw new ExtismException(e);
}
}

long[] statusCode(Instance instance, long... args) {
return new long[]{lastResponse == null ? 0 : lastResponse.statusCode()};
}


long[] headers(Instance instance, long[] longs) {
var result = new long[1];
if (lastResponse == null) {
return result;
}

// FIXME duplicated headers are effectively overwriting duplicate values!
var objBuilder = Json.createObjectBuilder();
for (var entry : lastResponse.headers().map().entrySet()) {
for (var v : entry.getValue()) {
objBuilder.add(entry.getKey(), v);
}
}

var bytes = objBuilder.build().toString().getBytes(StandardCharsets.UTF_8);
result[0] = memory().writeBytes(bytes);
return result;
}


public HostFunction[] toHostFunctions() {
return new HostFunction[]{
new HostFunction(
Kernel.IMPORT_MODULE_NAME,
"http_request",
List.of(ValueType.I64, ValueType.I64),
List.of(ValueType.I64),
this::request),
new HostFunction(
Kernel.IMPORT_MODULE_NAME,
"http_status_code",
List.of(),
List.of(ValueType.I32),
this::statusCode),
new HostFunction(
Kernel.IMPORT_MODULE_NAME,
"http_headers",
List.of(),
List.of(ValueType.I64),
this::headers),

};
}
}

private static class HostPattern {
private final String pattern;
private final boolean exact;

public HostPattern(String pattern) {
if (pattern.indexOf('*', 1) != -1) {
throw new ExtismException("Illegal pattern " + pattern);
}
int wildcard = pattern.indexOf('*');
if (wildcard < 0) {
this.exact = true;
this.pattern = pattern;
} else if (wildcard == 0) {
this.exact = false;
this.pattern = pattern.substring(1);
} else {
throw new ExtismException("Illegal pattern " + pattern);
}
}

public boolean matches(String host) {
if (exact) {
return host.equals(pattern);
} else {
return host.endsWith(pattern);
}
}
}


}
5 changes: 4 additions & 1 deletion src/main/java/org/extism/chicory/sdk/Linker.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,21 +36,24 @@ Plugin link() {
var dg = new DependencyGraph(logger);

Map<String, String> config;
String[] allowedHosts;
WasiOptions wasiOptions;
CachedAotMachineFactory aotMachineFactory;
if (manifest.options == null) {
config = Map.of();
allowedHosts = new String[0];
wasiOptions = null;
aotMachineFactory = null;
} else {
dg.setOptions(manifest.options);
config = manifest.options.config;
allowedHosts = manifest.options.allowedHosts;
wasiOptions = manifest.options.wasiOptions;
aotMachineFactory = manifest.options.aot? new CachedAotMachineFactory() : null;
}

// Register the HostEnv exports.
var hostEnv = new HostEnv(new Kernel(aotMachineFactory), config, logger);
var hostEnv = new HostEnv(new Kernel(aotMachineFactory), config, allowedHosts, logger);
dg.registerFunctions(hostEnv.toHostFunctions());

// Register the WASI host functions.
Expand Down
13 changes: 13 additions & 0 deletions src/main/java/org/extism/chicory/sdk/Manifest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import com.dylibso.chicory.wasi.WasiOptions;

import java.nio.file.Path;
import java.util.EnumSet;
import java.util.List;
import java.util.Map;
Expand All @@ -17,6 +18,7 @@ public static class Options {
EnumSet<Validation> validationFlags = EnumSet.noneOf(Validation.class);
Map<String, String> config;
WasiOptions wasiOptions;
String[] allowedHosts;

public Options withAoT() {
this.aot = true;
Expand All @@ -33,6 +35,17 @@ public Options withValidation(Validation... vs) {
return this;
}

public Options withAllowedHosts(String... allowedHosts) {
for (String allowedHost : allowedHosts) {
// Wildcards are only allowed at starting position and may occur only once.
if (allowedHost.indexOf('*') > 0 || allowedHost.indexOf('*', 1) != -1) {
throw new ExtismException("Illegal pattern " + allowedHost);
}
}
this.allowedHosts = allowedHosts;
return this;
}

public Options withWasi(WasiOptions wasiOptions) {
this.wasiOptions = wasiOptions;
return this;
Expand Down
Loading

0 comments on commit d5e23db

Please sign in to comment.