diff --git a/api/src/main/java/ai/djl/repository/AbstractRepository.java b/api/src/main/java/ai/djl/repository/AbstractRepository.java index c28a3b16887..f6fcd9513ca 100644 --- a/api/src/main/java/ai/djl/repository/AbstractRepository.java +++ b/api/src/main/java/ai/djl/repository/AbstractRepository.java @@ -187,7 +187,9 @@ protected void download(Path tmp, URI baseUri, Artifact.Item item, Progress prog } logger.debug("Downloading artifact: {} ...", fileUri); - try (InputStream is = new BufferedInputStream(fileUri.toURL().openStream())) { + // The artifact URI comes from the repository metadata and may be absolute, so route it + // through Utils.openUrl to apply the same URL handling as the rest of the download paths. + try (InputStream is = Utils.openUrl(fileUri.toURL())) { save(is, tmp, item, progress); } } diff --git a/api/src/main/java/ai/djl/repository/SimpleUrlRepository.java b/api/src/main/java/ai/djl/repository/SimpleUrlRepository.java index a2880aedb5c..5e19983c1f3 100644 --- a/api/src/main/java/ai/djl/repository/SimpleUrlRepository.java +++ b/api/src/main/java/ai/djl/repository/SimpleUrlRepository.java @@ -20,7 +20,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.BufferedInputStream; import java.io.IOException; import java.io.InputStream; import java.net.HttpURLConnection; @@ -102,7 +101,10 @@ public List getResources() { protected void download(Path tmp, URI baseUri, Artifact.Item item, Progress progress) throws IOException { logger.debug("Downloading artifact: {} ...", uri); - try (InputStream is = new BufferedInputStream(uri.toURL().openStream())) { + // Route through Utils.openUrl so the download is subject to the same URL handling + // (scheme restriction, non-public address blocking, redirect re-validation) instead of + // calling openStream() directly, which would bypass them. + try (InputStream is = Utils.openUrl(uri.toURL())) { save(is, tmp, item, progress); } } @@ -131,23 +133,31 @@ private synchronized Metadata getMetadata() throws IOException { String hash = Utils.hash(uri.toString()); MRL mrl = model(Application.UNDEFINED, DefaultModelZoo.GROUP_ID, hash); metadata.setRepositoryUri(mrl.toURI()); + resolved = true; return metadata; } private long getContentLength() throws IOException { String scheme = uri.getScheme(); if ("http".equalsIgnoreCase(scheme) || "https".equalsIgnoreCase(scheme)) { + if (Utils.isOfflineMode()) { + // An unknown size is represented as -1, and the artifact may already be cached, so + // reporting the size as unknown keeps a cached model loadable offline. Probing here + // would fail the load before prepare() can find the cached resource directory. + logger.debug("Offline mode is enabled, skipping content length probe: {}", uri); + return -1; + } HttpURLConnection conn = null; try { - resolved = true; - conn = (HttpURLConnection) uri.toURL().openConnection(); - conn.setRequestMethod("HEAD"); + // Route the HEAD probe through the same handling as the download path, so + // redirects are resolved explicitly and every hop is checked. + conn = Utils.openHttpConnection(uri.toURL(), "HEAD", Collections.emptyMap()); int code = conn.getResponseCode(); if (code != 200) { logger.debug("Failed detect content length, error code: {}", code); return -1; } - return conn.getContentLength(); + return conn.getContentLengthLong(); } finally { if (conn != null) { conn.disconnect(); diff --git a/api/src/main/java/ai/djl/util/ClassLoaderUtils.java b/api/src/main/java/ai/djl/util/ClassLoaderUtils.java index dc5bfba8678..53c31687991 100644 --- a/api/src/main/java/ai/djl/util/ClassLoaderUtils.java +++ b/api/src/main/java/ai/djl/util/ClassLoaderUtils.java @@ -242,27 +242,88 @@ public static void nativeLoad(String nativeHelper, String path) { /** * Tries to compile java classes in the directory. * + *

Compiling bundled {@code .java} sources at model-load time is disabled by default. Enable + * it with {@code DJL_COMPILE_JAVA=true} or {@code -Dai.djl.compile_java=true} when loading + * models from a trusted source. Models that ship precompiled {@code .class} or {@code .jar} + * files, or that supply a translator programmatically, are unaffected. + * * @param dir the directory to scan java file. */ public static void compileJavaClass(Path dir) { + String[] files = findJavaSources(dir); + if (files.length == 0) { + // Nothing to compile: the common case of a model without bundled sources stays quiet. + return; + } + if (!isDynamicCompilationEnabled()) { + // Logged at error and not thrown: a model may also ship a precompiled .class or .jar + // that satisfies the lookup, or the sources may be incidental to it, so failing here + // would break loads that do not depend on compiling them. The message names the flag so + // a model that did depend on it is diagnosable. + logger.error( + "Skipping {} bundled java file(s) in {}: compiling bundled sources is disabled" + + " by default. Set DJL_COMPILE_JAVA=true or" + + " -Dai.djl.compile_java=true to enable it for models from a trusted" + + " source, or ship precompiled .class/.jar files instead.", + files.length, + dir); + return; + } + JavaCompiler compiler = ToolProvider.getSystemJavaCompiler(); + if (compiler == null) { + logger.warn( + "Cannot compile {} bundled java file(s) in {}: no system java compiler is" + + " available, a JDK is required.", + files.length, + dir); + return; + } + try { + int rc = compiler.run(null, null, null, files); + if (rc != 0) { + // Report the failure: the caller resolves a translator from the compiled output, so + // a failed compile otherwise surfaces later as a missing class or as a silent + // fallback to a default implementation. Compiler diagnostics go to stderr. + logger.error( + "Failed to compile {} bundled java file(s) in {}, compiler exit code {}." + + " See the compiler output for details.", + files.length, + dir, + rc); + } + } catch (Throwable e) { + logger.warn("Failed to compile bundled java file", e); + } + } + + private static String[] findJavaSources(Path dir) { try { if (!Files.isDirectory(dir)) { - logger.debug("Directory not exists: {}", dir); - return; + logger.debug("Directory does not exist: {}", dir); + return Utils.EMPTY_ARRAY; } - String[] files; try (Stream stream = Files.walk(dir)) { - files = - stream.filter(p -> Files.isRegularFile(p) && p.toString().endsWith(".java")) - .map(p -> p.toAbsolutePath().toString()) - .toArray(String[]::new); - } - JavaCompiler compiler = ToolProvider.getSystemJavaCompiler(); - if (files.length > 0) { - compiler.run(null, null, null, files); + return stream.filter(p -> Files.isRegularFile(p) && p.toString().endsWith(".java")) + .map(p -> p.toAbsolutePath().toString()) + .toArray(String[]::new); } - } catch (Throwable e) { - logger.warn("Failed to compile bundled java file", e); + } catch (IOException | RuntimeException e) { + // Files.walk reports traversal errors from the terminal operation as an unchecked + // UncheckedIOException (an unreadable subdirectory, a symlink loop, a directory removed + // concurrently), so catching IOException alone would turn a scan problem into a failed + // model load. + logger.warn("Failed to scan for bundled java files in {}", dir, e); + return Utils.EMPTY_ARRAY; } } + + /** + * Returns whether dynamic compilation of bundled java sources is explicitly enabled. + * + * @return true if {@code DJL_COMPILE_JAVA} / {@code ai.djl.compile_java} is set true + */ + static boolean isDynamicCompilationEnabled() { + String mode = Utils.getenv("DJL_COMPILE_JAVA", System.getProperty("ai.djl.compile_java")); + return Boolean.parseBoolean(mode); + } } diff --git a/api/src/main/java/ai/djl/util/Utils.java b/api/src/main/java/ai/djl/util/Utils.java index 1fa4fee6e76..ece42065609 100644 --- a/api/src/main/java/ai/djl/util/Utils.java +++ b/api/src/main/java/ai/djl/util/Utils.java @@ -23,7 +23,11 @@ import java.io.IOException; import java.io.InputStream; import java.net.HttpURLConnection; +import java.net.InetAddress; +import java.net.JarURLConnection; import java.net.URL; +import java.net.URLConnection; +import java.net.UnknownHostException; import java.nio.charset.StandardCharsets; import java.nio.file.FileVisitOption; import java.nio.file.Files; @@ -39,6 +43,7 @@ import java.util.Map; import java.util.Objects; import java.util.Scanner; +import java.util.function.Predicate; import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; @@ -49,6 +54,10 @@ public final class Utils { private static final Logger logger = LoggerFactory.getLogger(Utils.class); + private static final int MAX_REDIRECTS = 5; + + private static final int CONNECT_TIMEOUT_MILLIS = 10_000; + public static final String[] EMPTY_ARRAY = new String[0]; private Utils() {} @@ -521,21 +530,398 @@ public static InputStream openUrl(URL url) throws IOException { * @param url the URL to open * @param headers HTTP headers * @return an input stream for reading from the URL connection. - * @throws IOException if an I/O exception occurs + * @throws IOException if an I/O exception occurs, or the destination or scheme is not allowed */ public static InputStream openUrl(URL url, Map headers) throws IOException { + // Remote fetches are limited to public http(s) destinations by default, and redirects are + // resolved explicitly so each hop is checked against the same rule. file: is a local read. + // jar: is allowed only when the URL it nests is itself local (see below). Other schemes are + // not supported. Set DJL_ALLOW_INSECURE_URL=true (or -Dai.djl.allow_insecure_url=true) to + // restore the previous unrestricted behavior. String protocol = url.getProtocol(); - if ("http".equalsIgnoreCase(protocol) || "https".equalsIgnoreCase(protocol)) { - if (isOfflineMode()) { - throw new IOException("Offline mode is enabled."); + if (isInsecureUrlAllowed()) { + if ("http".equalsIgnoreCase(protocol) || "https".equalsIgnoreCase(protocol)) { + return openHttpStream(url, headers); } + return new BufferedInputStream(url.openStream()); + } + if ("http".equalsIgnoreCase(protocol) || "https".equalsIgnoreCase(protocol)) { + return openHttpStream(url, headers); + } + if ("file".equalsIgnoreCase(protocol)) { + requireLocalFileUrl(url); + return new BufferedInputStream(url.openStream()); + } + if ("jar".equalsIgnoreCase(protocol)) { + // A jar: URL nests another URL, as jar:!/, and the connection fetches + // that nested URL, so the nested URL is what has to be judged: the outer URL reports + // protocol "jar" and an empty host, which no destination rule can act on. + requireLocalRead(nestedUrlOf(url)); + return new BufferedInputStream(url.openStream()); + } + throw new IOException( + "Blocked request using unsupported URL protocol: " + + protocol + + ". Set DJL_ALLOW_INSECURE_URL=true to override."); + } + + /** + * Returns the URL nested inside a {@code jar:} URL, or the URL itself if it cannot be read. + * + * @param url the jar URL + * @return the nested URL + * @throws IOException if the nested URL cannot be determined + */ + private static URL nestedUrlOf(URL url) throws IOException { + URLConnection conn = url.openConnection(); + if (!(conn instanceof JarURLConnection)) { + // A container may register its own handler for jar:. Without access to the nested URL + // there is nothing to judge, so refuse rather than guess. + throw new IOException( + "Blocked jar URL with an unrecognized connection type: " + + conn.getClass().getName() + + ". Set DJL_ALLOW_INSECURE_URL=true to override."); + } + return ((JarURLConnection) conn).getJarFileURL(); + } + + /** + * Verifies that a URL denotes a local read rather than a network fetch. + * + *

Judged on the destination rather than on the scheme name, because several schemes read + * locally in one form and over the network in another. {@code file:} with an authority is + * fetched over FTP by the JDK, and {@code jar:} nests whatever it is given. Container schemes + * that address a resource inside the running application ({@code nested:} for a Spring Boot + * executable jar, {@code vfs:}, {@code wsjar:}, {@code bundleresource:}) are local reads and + * are allowed, so packaging an application as an executable jar keeps working. + * + * @param url the URL to check + * @throws IOException if the URL would perform a network fetch + */ + private static void requireLocalRead(URL url) throws IOException { + String protocol = url.getProtocol(); + if ("http".equalsIgnoreCase(protocol) + || "https".equalsIgnoreCase(protocol) + || "ftp".equalsIgnoreCase(protocol) + || "ftps".equalsIgnoreCase(protocol)) { + throw new IOException( + "Blocked jar URL nesting a remote protocol: " + + protocol + + ". Set DJL_ALLOW_INSECURE_URL=true to override."); + } + if ("file".equalsIgnoreCase(protocol)) { + requireLocalFileUrl(url); + } else if ("jar".equalsIgnoreCase(protocol)) { + requireLocalRead(nestedUrlOf(url)); + } + } + + /** + * Verifies that a {@code file:} URL has no remote authority. + * + *

{@code file://host/path} is not a local read: the JDK's file handler falls back to FTP for + * any authority that is not empty, {@code localhost} or a {@code ~} user reference, so such a + * URL performs an outbound fetch to {@code host}. + * + * @param url the file URL to check + * @throws IOException if the URL carries a remote authority + */ + private static void requireLocalFileUrl(URL url) throws IOException { + String host = url.getHost(); + if (host != null + && !host.isEmpty() + && !"localhost".equalsIgnoreCase(host) + && !host.startsWith("~")) { + throw new IOException( + "Blocked file URL with a remote authority: " + + host + + ". Set DJL_ALLOW_INSECURE_URL=true to override."); + } + } + + private static InputStream openHttpStream(URL url, Map headers) + throws IOException { + HttpURLConnection conn = openHttpConnection(url, "GET", headers); + try { + return new BufferedInputStream(conn.getInputStream()); + } catch (IOException e) { + // An error response leaves the connection open, so release it before propagating. + conn.disconnect(); + throw e; + } + } + + /** + * Opens an http(s) connection using the specified request method, resolving redirects + * explicitly so that every hop is checked against the same destination rule. + * + *

This is the shared entry point for remote reads: {@link #openUrl(URL, Map)} uses it for + * {@code GET}, and callers that only need response metadata can use it for {@code HEAD}. The + * caller owns the returned connection and should disconnect it when finished. + * + * @param url the URL to open + * @param method the HTTP request method, for example {@code GET} or {@code HEAD} + * @param headers HTTP headers + * @return the connection for the final, validated hop + * @throws IOException if an I/O exception occurs or a hop is not allowed + */ + public static HttpURLConnection openHttpConnection( + URL url, String method, Map headers) throws IOException { + if (isOfflineMode()) { + throw new IOException("Offline mode is enabled."); + } + if (headers == null) { + // No headers is a reasonable thing for a caller to mean, and an NPE from inside the + // redirect loop would be hard to trace back to the argument. + headers = Collections.emptyMap(); + } + if (isInsecureUrlAllowed()) { HttpURLConnection conn = (HttpURLConnection) url.openConnection(); + conn.setRequestMethod(method); for (Map.Entry entry : headers.entrySet()) { conn.addRequestProperty(entry.getKey(), entry.getValue()); } - return conn.getInputStream(); + return conn; + } + return openHttpConnection(url, method, headers, Utils::isPublicHost); + } + + /** + * Opens an http(s) connection, resolving redirects explicitly so that every hop is checked + * against {@code hostAllowed} rather than being followed by the connection itself. + * + *

The host check is a parameter so the redirect handling can be exercised against a local + * test server. + * + * @param url the URL to open + * @param method the HTTP request method + * @param headers HTTP headers + * @param hostAllowed the check applied to the host of every hop + * @return the connection for the final, validated hop + * @throws IOException if an I/O exception occurs or a hop is not allowed + */ + static HttpURLConnection openHttpConnection( + URL url, String method, Map headers, Predicate hostAllowed) + throws IOException { + // Also normalised here rather than relying on the public overload: this one is called + // directly, so it should not depend on a caller having done it. + Map requestHeaders = headers == null ? Collections.emptyMap() : headers; + URL current = url; + // Set once a redirect leaves the original origin, so credentials are not carried to it. + boolean stripCredentials = false; + // One initial request plus at most MAX_REDIRECTS redirect hops. + for (int i = 0; i <= MAX_REDIRECTS; ++i) { + String protocol = current.getProtocol(); + if (!"http".equalsIgnoreCase(protocol) && !"https".equalsIgnoreCase(protocol)) { + throw new IOException("Blocked redirect to unsupported URL protocol: " + protocol); + } + if (!hostAllowed.test(current.getHost())) { + throw new IOException( + "Blocked request to non-public address: " + current.getHost()); + } + HttpURLConnection conn = (HttpURLConnection) current.openConnection(); + conn.setInstanceFollowRedirects(false); + // Bound the connect phase. This loop can open up to MAX_REDIRECTS + 1 connections, so + // without it a black-holed destination multiplies the OS SYN timeout by the hop count. + // Only the connect phase is bounded: a read timeout would also apply to the body, and + // artifact downloads are large and legitimately slow. + conn.setConnectTimeout(CONNECT_TIMEOUT_MILLIS); + conn.setRequestMethod(method); + for (Map.Entry entry : requestHeaders.entrySet()) { + if (stripCredentials && isCredentialHeader(entry.getKey())) { + continue; + } + conn.addRequestProperty(entry.getKey(), entry.getValue()); + } + int code; + try { + code = conn.getResponseCode(); + } catch (IOException e) { + conn.disconnect(); + throw e; + } + if (!isRedirect(code)) { + if (code / 100 == 3) { + // A 3xx that is not followed (300, 304, 305, ...) still carries a body, and + // getInputStream() succeeds for any code below 400, so returning it would let + // a negotiation page be written out as the model or the native library. + conn.disconnect(); + throw new IOException( + "Unexpected redirect status " + code + " for URL: " + url); + } + // Other statuses are the caller's to interpret: openUrl surfaces an error when the + // stream is opened, and the content-length probe treats non-200 as unknown size. + return conn; + } + // Redirect: re-validate the target on the next loop iteration instead of letting + // HttpURLConnection follow it blindly (which would defeat the host check above). + String location = conn.getHeaderField("Location"); + conn.disconnect(); + if (location == null) { + throw new IOException("Redirect (" + code + ") with no Location header"); + } + // Resolve relative redirects against the current URL, then re-validate the target. + URL next = new URL(current, location); + if (isSchemeDowngrade(current, next)) { + throw new IOException( + "Blocked redirect downgrading https to " + next.getProtocol()); + } + if (!sameOrigin(current, next)) { + // Following a redirect to another origin must not carry credentials with it. Once + // dropped they stay dropped, even if a later hop returns to the original origin. + stripCredentials = true; + } + current = next; + } + throw new IOException("Too many redirects (>" + MAX_REDIRECTS + ") for URL: " + url); + } + + /** + * Returns whether insecure (non-public) URL access is explicitly allowed. + * + * @return true if {@code DJL_ALLOW_INSECURE_URL} / {@code ai.djl.allow_insecure_url} is set + * true + */ + static boolean isInsecureUrlAllowed() { + String mode = + getenv("DJL_ALLOW_INSECURE_URL", System.getProperty("ai.djl.allow_insecure_url")); + return Boolean.parseBoolean(mode); + } + + /** + * Returns whether a host resolves exclusively to public (non-internal) IP addresses. + * + *

A host is rejected if it is empty or if any resolved address is loopback, link-local, + * site-local (RFC 1918 private), wildcard, multicast, or an IPv6 unique local address + * (fc00::/7). The latter is checked explicitly because {@link InetAddress#isSiteLocalAddress()} + * only covers the deprecated fec0::/10 range. IPv4-mapped IPv6 literals resolve to an {@code + * Inet4Address} and are covered by the IPv4 checks. + * + *

Note that this resolves the host name and the subsequent connection resolves it again, so + * a name that resolves to an allowed address here could resolve to a different address when the + * connection is made. Pinning the connection to a validated address would be needed to close + * that gap; environments that require it should restrict egress at the network layer. + * + * @param host the host name to validate + * @return true if every resolved address is a public address + */ + static boolean isPublicHost(String host) { + if (host == null || host.isEmpty()) { + return false; + } + try { + InetAddress[] addresses = InetAddress.getAllByName(host); + if (addresses.length == 0) { + return false; + } + for (InetAddress addr : addresses) { + if (addr.isLoopbackAddress() + || addr.isLinkLocalAddress() + || addr.isSiteLocalAddress() + || addr.isAnyLocalAddress() + || addr.isMulticastAddress() + || isIpv6UniqueLocal(addr) + || isCarrierGradeNat(addr)) { + return false; + } + } + } catch (UnknownHostException e) { + // The host name is the actionable part; the stack trace adds nothing and would be + // emitted once per unresolvable host. + logger.warn("Cannot resolve host, treating as not public: {}", host); + return false; + } + return true; + } + + /** + * Returns whether following {@code to} from {@code from} would drop TLS. + * + *

{@link HttpURLConnection} refuses to follow a redirect that changes scheme, so resolving + * hops explicitly has to refuse the downgrade too. Otherwise an {@code https} fetch answered + * with a plaintext {@code Location} would be retrieved in the clear, and nothing downstream + * would notice: {@code SimpleUrlRepository} never records a checksum, so checksum validation is + * a no-op for these artifacts. + * + * @param from the current hop + * @param to the redirect target + * @return true if the hop drops https + */ + static boolean isSchemeDowngrade(URL from, URL to) { + return "https".equalsIgnoreCase(from.getProtocol()) + && !"https".equalsIgnoreCase(to.getProtocol()); + } + + /** + * Returns whether two URLs share a scheme, host and effective port. + * + * @param a the first URL + * @param b the second URL + * @return true if both URLs have the same origin + */ + static boolean sameOrigin(URL a, URL b) { + return a.getProtocol().equalsIgnoreCase(b.getProtocol()) + && a.getHost().equalsIgnoreCase(b.getHost()) + && effectivePort(a) == effectivePort(b); + } + + /** + * Returns the port of a URL, substituting the protocol default when the port is implicit. + * + *

{@link URL#getPort()} returns -1 for an implicit port, so comparing it directly would + * treat {@code https://host/a} and {@code https://host:443/b} as different origins. + * + * @param url the URL to read the port from + * @return the explicit port, or the protocol default when the port is implicit + */ + private static int effectivePort(URL url) { + int port = url.getPort(); + return port == -1 ? url.getDefaultPort() : port; + } + + /** + * Returns whether a header name carries credentials that must not follow a cross-origin + * redirect. + * + * @param name the header name + * @return true if the header carries credentials + */ + static boolean isCredentialHeader(String name) { + return "Authorization".equalsIgnoreCase(name) + || "Cookie".equalsIgnoreCase(name) + || "Proxy-Authorization".equalsIgnoreCase(name); + } + + private static boolean isRedirect(int code) { + // 300, 304 and 305 are 3xx but are not redirects to follow. + return code == 301 || code == 302 || code == 303 || code == 307 || code == 308; + } + + /** + * Returns whether an IPv4 address is in the RFC 6598 carrier-grade NAT range (100.64.0.0/10). + * + *

Checked explicitly because {@link InetAddress} has no predicate for it and because it is + * routinely used for Kubernetes secondary pod CIDRs, so an internal service can be addressed + * there. Other reserved ranges are deliberately not included: no deployment serves an internal + * endpoint from IETF protocol assignments or unallocated space, so blocking them would add + * surface without removing risk. + * + * @param addr the address to classify + * @return true if the address is in 100.64.0.0/10 + */ + private static boolean isCarrierGradeNat(InetAddress addr) { + byte[] b = addr.getAddress(); + if (b.length != 4 || (b[0] & 0xFF) != 100) { + return false; } - return new BufferedInputStream(url.openStream()); + int second = b[1] & 0xFF; + return second >= 64 && second <= 127; + } + + private static boolean isIpv6UniqueLocal(InetAddress addr) { + byte[] bytes = addr.getAddress(); + // fc00::/7 - the first seven bits are 1111110 + return bytes.length == 16 && (bytes[0] & 0xFE) == 0xFC; } /** diff --git a/api/src/test/java/ai/djl/repository/RepositoryDownloadTest.java b/api/src/test/java/ai/djl/repository/RepositoryDownloadTest.java new file mode 100644 index 00000000000..54820ac3aa4 --- /dev/null +++ b/api/src/test/java/ai/djl/repository/RepositoryDownloadTest.java @@ -0,0 +1,139 @@ +/* + * Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance + * with the License. A copy of the License is located at + * + * http://aws.amazon.com/apache2.0/ + * + * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES + * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions + * and limitations under the License. + */ +package ai.djl.repository; + +import ai.djl.util.Utils; + +import org.testng.Assert; +import org.testng.SkipException; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.net.URI; +import java.nio.file.Files; +import java.nio.file.Path; + +/** Tests the URL handling applied when downloading artifacts declared by repository metadata. */ +public class RepositoryDownloadTest { + + private String savedOffline; + + @BeforeMethod + public void clearOfflineMode() { + savedOffline = System.getProperty("ai.djl.offline"); + if (Utils.getenv("DJL_OFFLINE") != null) { + throw new SkipException("DJL_OFFLINE is set in the environment"); + } + // Offline mode is refused before the destination is examined, so these tests would assert + // on the wrong message under `gradlew --offline`, which sets this property. + System.clearProperty("ai.djl.offline"); + } + + @AfterMethod + public void restoreOfflineMode() { + if (savedOffline == null) { + System.clearProperty("ai.djl.offline"); + } else { + System.setProperty("ai.djl.offline", savedOffline); + } + } + + @Test + public void testNonPublicArtifactUriIsRejected() throws IOException { + // An artifact URI declared in repository metadata may be absolute, which bypasses base-URI + // resolution. It must still be subject to the same destination rule as the other download + // paths, so a non-public destination is refused. + IOException e = downloadItemUri("http://169.254.169.254/latest/meta-data/x.tar.gz"); + Assert.assertTrue( + e.getMessage().contains("non-public"), "unexpected message: " + e.getMessage()); + } + + @Test + public void testPrivateArtifactUriIsRejected() throws IOException { + IOException e = downloadItemUri("http://10.0.0.5/model.tar.gz"); + Assert.assertTrue( + e.getMessage().contains("non-public"), "unexpected message: " + e.getMessage()); + } + + @Test + public void testLoopbackArtifactUriIsRejected() throws IOException { + IOException e = downloadItemUri("http://127.0.0.1:1/model.tar.gz"); + Assert.assertTrue( + e.getMessage().contains("non-public"), "unexpected message: " + e.getMessage()); + } + + @Test + public void testUnsupportedSchemeArtifactUriIsRejected() throws IOException { + IOException e = downloadItemUri("ftp://ftp.example.com/model.tar.gz"); + Assert.assertTrue( + e.getMessage().contains("URL protocol"), "unexpected message: " + e.getMessage()); + } + + /** + * Invokes the shared artifact-download path with an absolute item URI and returns the thrown + * exception. + */ + private IOException downloadItemUri(String itemUri) throws IOException { + Path tmp = Files.createTempDirectory("djl-download-probe"); + try { + AbstractRepository repository = + new RemoteRepository( + "test", URI.create("https://resources.djl.ai/test-models/")); + Artifact.Item item = new Artifact.Item(); + item.setName("data"); + item.setUri(itemUri); + + return Assert.expectThrows( + IOException.class, () -> repository.download(tmp, URI.create(""), item, null)); + } finally { + Files.deleteIfExists(tmp); + } + } + + @Test + public void testSimpleUrlRepositoryRejectsNonPublicHost() { + // The content-length HEAD probe is a second entry point into the same download flow, so it + // applies the same destination rule. Assert on the thrown cause rather than on + // getResources(), which swallows every IOException and would pass either way. + SimpleUrlRepository repository = + new SimpleUrlRepository( + "test", + URI.create("http://169.254.169.254/latest/meta-data/x.tar.gz"), + "x.tar.gz"); + IOException e = Assert.expectThrows(IOException.class, () -> repository.locate(null)); + Assert.assertTrue(e.getMessage().contains("non-public"), "unexpected: " + e.getMessage()); + } + + @Test + public void testSimpleUrlRepositoryDownloadRejectsNonPublicHost() throws IOException { + SimpleUrlRepository repository = + new SimpleUrlRepository( + "test", URI.create("http://10.0.0.5/model.tar.gz"), "model.tar.gz"); + Path tmp = Files.createTempDirectory("djl-simpleurl-probe"); + try { + Artifact.Item item = new Artifact.Item(); + item.setName("data"); + item.setUri("http://10.0.0.5/model.tar.gz"); + IOException e = + Assert.expectThrows( + IOException.class, + () -> repository.download(tmp, URI.create(""), item, null)); + Assert.assertTrue( + e.getMessage().contains("non-public"), "unexpected: " + e.getMessage()); + } finally { + Files.deleteIfExists(tmp); + } + } +} diff --git a/api/src/test/java/ai/djl/util/UrlAccessAndCompilationTest.java b/api/src/test/java/ai/djl/util/UrlAccessAndCompilationTest.java new file mode 100644 index 00000000000..eef6d2178da --- /dev/null +++ b/api/src/test/java/ai/djl/util/UrlAccessAndCompilationTest.java @@ -0,0 +1,851 @@ +/* + * Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance + * with the License. A copy of the License is located at + * + * http://aws.amazon.com/apache2.0/ + * + * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES + * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions + * and limitations under the License. + */ +package ai.djl.util; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; + +import org.testng.Assert; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.io.InputStream; +import java.net.HttpURLConnection; +import java.net.InetSocketAddress; +import java.net.MalformedURLException; +import java.net.URL; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.jar.JarOutputStream; +import java.util.stream.Stream; +import java.util.zip.ZipEntry; + +/** + * Tests for the configurable URL-access and bundled-source-compilation defaults in {@link + * Utils#openUrl} and {@link ClassLoaderUtils#compileJavaClass}. + */ +public class UrlAccessAndCompilationTest { + + private String savedInsecureUrl; + private String savedCompileJava; + private String savedOffline; + + @BeforeMethod + public void saveProperties() { + // The build forwards every ai.djl.* system property into the test JVM, so these may already + // be set by the developer running the suite. Record them and put them back afterwards + // instead of clearing them, which would change behavior for every later test in the module. + savedInsecureUrl = System.getProperty("ai.djl.allow_insecure_url"); + savedCompileJava = System.getProperty("ai.djl.compile_java"); + savedOffline = System.getProperty("ai.djl.offline"); + if (Utils.getenv("DJL_OFFLINE") != null) { + throw new org.testng.SkipException("DJL_OFFLINE is set in the environment"); + } + if (Utils.getenv("DJL_ALLOW_INSECURE_URL") != null) { + throw new org.testng.SkipException("DJL_ALLOW_INSECURE_URL is set in the environment"); + } + // The opt-out disables everything these tests assert, and the build forwards ai.djl.* + // properties into the test JVM, so clear it here and restore it in cleanup(). + System.clearProperty("ai.djl.allow_insecure_url"); + // Offline mode is refused before the destination is examined, which is correct but would + // make every assertion below read "Offline mode is enabled" instead. `gradlew --offline` + // sets this property, so clear it for these tests and restore it in cleanup(). + System.clearProperty("ai.djl.offline"); + } + + @AfterMethod + public void cleanup() { + restore("ai.djl.allow_insecure_url", savedInsecureUrl); + restore("ai.djl.compile_java", savedCompileJava); + restore("ai.djl.offline", savedOffline); + } + + private static void restore(String key, String value) { + if (value == null) { + System.clearProperty(key); + } else { + System.setProperty(key, value); + } + } + + // ----- URL access: Utils.openUrl ----- + + @Test + public void testFileProtocolAllowedByDefault() throws IOException { + // file:// performs no network fetch, so it is allowed by default. + Path tmp = Files.createTempFile("djl-local", ".txt"); + Files.write(tmp, "secret".getBytes(StandardCharsets.UTF_8)); + try (InputStream is = Utils.openUrl(tmp.toUri().toURL())) { + Assert.assertEquals(new String(is.readAllBytes(), StandardCharsets.UTF_8), "secret"); + } finally { + Files.deleteIfExists(tmp); + } + } + + @Test + public void testJarProtocolAllowedByDefault() throws IOException { + // jar:file: with no authority reads a resource bundled inside a local jar (e.g. + // djl-serving's plugin.definition). The nested URL is what decides: a remote nested scheme + // is covered by testJarUrlNestingRemoteProtocolIsBlocked and a nested file: carrying an + // authority by testJarUrlNestingRemoteFileAuthorityIsBlocked. Build a real one-entry jar + // and read it back through a jar: URL. + Path jar = Files.createTempFile("djl-jar", ".jar"); + try (JarOutputStream jos = new JarOutputStream(Files.newOutputStream(jar))) { + jos.putNextEntry(new ZipEntry("META-INF/probe.txt")); + jos.write("jarok".getBytes(StandardCharsets.UTF_8)); + jos.closeEntry(); + } + try { + URL jarUrl = new URL("jar:" + jar.toUri() + "!/META-INF/probe.txt"); + try (InputStream is = Utils.openUrl(jarUrl)) { + Assert.assertEquals(new String(is.readAllBytes(), StandardCharsets.UTF_8), "jarok"); + } + } finally { + // The JVM caches open jar files, so a handle can remain on the temp jar after the + // stream is closed. On Windows that makes an immediate delete fail, so fall back to + // deleting on JVM exit. + jar.toFile().deleteOnExit(); + try { + Files.deleteIfExists(jar); + } catch (IOException ignore) { + // cleaned up by deleteOnExit + } + } + } + + @Test + public void testJarUrlNestingRemoteProtocolIsBlocked() throws IOException { + // A jar: URL nests another URL and the connection fetches it, but reports protocol "jar" + // with an empty host, so nothing about the outer URL reveals the destination. Without a + // check on the nested URL, jar:http://host/x.jar reaches the network unseen. + AtomicInteger hits = new AtomicInteger(); + HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + server.createContext( + "/", + exchange -> { + hits.incrementAndGet(); + exchange.sendResponseHeaders(200, -1); + exchange.close(); + }); + server.start(); + try { + int port = server.getAddress().getPort(); + URL nested = new URL("jar:http://127.0.0.1:" + port + "/x.jar!/entry.txt"); + IOException e = Assert.expectThrows(IOException.class, () -> Utils.openUrl(nested)); + Assert.assertTrue( + e.getMessage().contains("nesting a remote protocol"), + "unexpected: " + e.getMessage()); + Assert.assertEquals(hits.get(), 0, "the nested destination must not be contacted"); + } finally { + server.stop(0); + } + } + + @Test + public void testFileUrlWithRemoteAuthorityIsBlocked() throws IOException { + // file: is only a local read when it has no authority. The JDK's file handler falls back to + // FTP for any other authority, so file://host/path performs an outbound fetch to host -- + // which the scheme name alone does not reveal. Verified: openConnection() on these returns + // sun.net.www.protocol.ftp.FtpURLConnection. + for (String u : + new String[] { + "file://10.0.0.5/pub/model.tar.gz", "file://169.254.169.254/latest/meta-data/x" + }) { + IOException e = Assert.expectThrows(IOException.class, () -> Utils.openUrl(new URL(u))); + Assert.assertTrue( + e.getMessage().contains("remote authority"), u + " -> " + e.getMessage()); + } + // A genuinely local file: URL still works, with and without an explicit localhost. + Path tmp = Files.createTempFile("djl-localauth", ".txt"); + Files.write(tmp, "local".getBytes(StandardCharsets.UTF_8)); + try { + try (InputStream is = Utils.openUrl(tmp.toUri().toURL())) { + Assert.assertEquals(new String(is.readAllBytes(), StandardCharsets.UTF_8), "local"); + } + // Build the authority form from the URI path, not from the filesystem path: on Windows + // an absolute path is "D:\..." and concatenating it produces a malformed URL. + URL viaLocalhost = new URL("file://localhost" + tmp.toUri().getPath()); + try (InputStream is = Utils.openUrl(viaLocalhost)) { + Assert.assertEquals(new String(is.readAllBytes(), StandardCharsets.UTF_8), "local"); + } + } finally { + Files.deleteIfExists(tmp); + } + } + + @Test + public void testJarUrlNestingRemoteFileAuthorityIsBlocked() { + // The nested protocol being "file" is not sufficient: jar:file://host/x.jar!/e is fetched + // over FTP exactly like the bare file: form, so the nested URL's authority has to be + // checked + // too. This is the shape a nested-protocol-name check alone lets through. + IOException e = + Assert.expectThrows( + IOException.class, + () -> + Utils.openUrl( + new URL("jar:file://10.0.0.5/pub/evil.jar!/entry.txt"))); + Assert.assertTrue( + e.getMessage().contains("remote authority"), "unexpected: " + e.getMessage()); + } + + @Test + public void testFtpProtocolBlockedByDefault() { + // Non-http(s), non-local schemes (ftp, gopher, ...) remain blocked by default: they are + // neither trusted local reads nor host-validated http(s), so loosening jar/file must not + // reopen them. + // Assert on the message: unpatched, an ftp: URL to a non-existent host also throws an + // IOException (UnknownHostException), so a bare assertThrows would pass either way. + IOException e = + Assert.expectThrows( + IOException.class, + () -> Utils.openUrl(new URL("ftp://ftp.example.com/pub/model.tar.gz"))); + Assert.assertTrue( + e.getMessage().contains("unsupported URL protocol"), + "unexpected: " + e.getMessage()); + } + + @Test + public void testNonPublicHostBlockedBeforeConnecting() throws IOException { + // The destination check runs before any request, so the listener must never be contacted. + // Redirect re-validation is covered by testRedirectToDisallowedHostIsRejected. + AtomicInteger hits = new AtomicInteger(); + HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + server.createContext( + "/", + exchange -> { + hits.incrementAndGet(); + exchange.sendResponseHeaders(200, -1); + exchange.close(); + }); + server.start(); + try { + URL url = + new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/model.tar.gz"); + IOException ex = Assert.expectThrows(IOException.class, () -> Utils.openUrl(url)); + Assert.assertTrue( + ex.getMessage().contains("Blocked request to non-public address"), + "expected non-public-address block, got: " + ex.getMessage()); + Assert.assertEquals(hits.get(), 0, "the destination must not be contacted"); + } finally { + server.stop(0); + } + } + + @Test + public void testLinkLocalAddressBlockedByDefault() { + IOException e = + Assert.expectThrows( + IOException.class, + () -> Utils.openUrl(new URL("http://169.254.169.254/latest/meta-data/"))); + Assert.assertTrue(e.getMessage().contains("non-public"), "unexpected: " + e.getMessage()); + } + + @Test + public void testLoopbackBlockedByDefault() { + // Unpatched, both of these throw ConnectException (connection refused), which is also an + // IOException, so the message is what distinguishes the fix from the pre-existing failure. + for (String u : new String[] {"http://127.0.0.1:1/", "http://localhost:1/"}) { + IOException e = Assert.expectThrows(IOException.class, () -> Utils.openUrl(new URL(u))); + Assert.assertTrue(e.getMessage().contains("non-public"), u + " -> " + e.getMessage()); + } + } + + @Test + public void testPrivateAddressBlockedByDefault() { + // Without the message assertion these would eventually pass on unpatched code too, after + // the OS SYN timeout rather than immediately. + for (String u : new String[] {"http://10.0.0.5/", "http://192.168.1.1/"}) { + IOException e = Assert.expectThrows(IOException.class, () -> Utils.openUrl(new URL(u))); + Assert.assertTrue(e.getMessage().contains("non-public"), u + " -> " + e.getMessage()); + } + } + + @Test + public void testHostValidationPredicate() { + // Non-public hosts must be rejected. + Assert.assertFalse(Utils.isPublicHost("127.0.0.1")); + Assert.assertFalse(Utils.isPublicHost("localhost")); + Assert.assertFalse(Utils.isPublicHost("169.254.169.254")); + Assert.assertFalse(Utils.isPublicHost("10.0.0.5")); + Assert.assertFalse(Utils.isPublicHost("192.168.0.1")); + Assert.assertFalse(Utils.isPublicHost("172.16.0.1")); + Assert.assertFalse(Utils.isPublicHost("")); + Assert.assertFalse(Utils.isPublicHost(null)); + // A well-known public address must pass. + Assert.assertTrue(Utils.isPublicHost("8.8.8.8")); + } + + @Test + public void testInsecureUrlAllowedFlagDefaultsFalse() { + Assert.assertFalse(Utils.isInsecureUrlAllowed()); + } + + // ----- Bundled-source compilation: ClassLoaderUtils.compileJavaClass ----- + + @Test + public void testCompilationDisabledByDefault() throws IOException { + // The environment variable takes precedence over the system property, matching + // DJL_OFFLINE, so skip rather than fail if a build host has it set. + if (Utils.getenv("DJL_COMPILE_JAVA") != null) { + throw new org.testng.SkipException("DJL_COMPILE_JAVA is set in the environment"); + } + // The build forwards ai.djl.* properties into the test JVM, so clear it for this test; the + // original value is restored by cleanup(). + System.clearProperty("ai.djl.compile_java"); + Path dir = Files.createTempDirectory("djl-compile"); + Path classes = Files.createDirectories(dir.resolve("classes")); + Files.write( + classes.resolve("Bundled.java"), + ("public class Bundled { static { System.setProperty(\"djl.test.canary\", \"ran\");" + + " } }") + .getBytes(StandardCharsets.UTF_8)); + System.clearProperty("djl.test.canary"); + try { + // Scanning must not throw: a model may still ship a usable .class or .jar, or the + // sources may be incidental, so failing here would break loads that never needed them. + ClassLoaderUtils.compileJavaClass(classes); + // Default (disabled): nothing compiled and the static initializer never ran. + Assert.assertFalse(Files.exists(classes.resolve("Bundled.class"))); + Assert.assertNull( + System.getProperty("djl.test.canary"), + "the bundled static initializer must not have run"); + Assert.assertFalse(ClassLoaderUtils.isDynamicCompilationEnabled()); + } finally { + deleteTree(dir); + } + } + + @Test + public void testScanSurvivesUnreadableSubdirectory() throws IOException { + // Files.walk surfaces traversal errors from the terminal operation as an unchecked + // UncheckedIOException, which must not escape as a failed model load. + Path dir = Files.createTempDirectory("djl-walkfail"); + Path classes = Files.createDirectories(dir.resolve("classes")); + Path sub = Files.createDirectories(classes.resolve("sub")); + try { + Files.write( + sub.resolve("Hidden.java"), "class Hidden {}".getBytes(StandardCharsets.UTF_8)); + if (!sub.toFile().setReadable(false) || sub.toFile().canRead()) { + // Running as root bypasses the permission, so the walk would succeed and the + // assertion below would hold for the wrong reason. + sub.toFile().setReadable(true); + throw new org.testng.SkipException("cannot make a directory unreadable here"); + } + // Must return normally rather than propagating the UncheckedIOException. + ClassLoaderUtils.compileJavaClass(classes); + } finally { + sub.toFile().setReadable(true); + deleteTree(dir); + } + } + + @Test + public void testCompilationEnabledWithOptIn() throws IOException { + if (Utils.getenv("DJL_COMPILE_JAVA") != null) { + throw new org.testng.SkipException("DJL_COMPILE_JAVA is set in the environment"); + } + System.setProperty("ai.djl.compile_java", "true"); + Assert.assertTrue(ClassLoaderUtils.isDynamicCompilationEnabled()); + Path dir = Files.createTempDirectory("djl-compile-optin"); + Path classes = Files.createDirectories(dir.resolve("classes")); + Files.write( + classes.resolve("Hello.java"), + "public class Hello { public int v() { return 1; } }" + .getBytes(StandardCharsets.UTF_8)); + try { + ClassLoaderUtils.compileJavaClass(classes); + // Opt-in: compilation runs (a system compiler must be present in the test JDK). + if (javax.tools.ToolProvider.getSystemJavaCompiler() != null) { + Assert.assertTrue(Files.exists(classes.resolve("Hello.class"))); + } + } finally { + deleteTree(dir); + } + } + + private static void deleteTree(Path root) throws IOException { + if (!Files.exists(root)) { + return; + } + try (java.util.stream.Stream walk = Files.walk(root)) { + walk.sorted(java.util.Comparator.reverseOrder()) + .forEach( + p -> { + try { + Files.deleteIfExists(p); + } catch (IOException ignore) { + // best effort + } + }); + } + } + + // ----- Redirect handling (exercised against a local server via the host-check seam) ----- + + /** Serves a redirect chain / final body on loopback for redirect tests. */ + private static HttpServer startServer(Map redirects, String body) + throws IOException { + HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + server.createContext( + "/", + exchange -> { + String path = exchange.getRequestURI().getPath(); + String target = redirects.get(path); + if (target != null) { + if (!target.isEmpty()) { + exchange.getResponseHeaders().add("Location", target); + } + exchange.sendResponseHeaders(302, -1); + } else { + byte[] out = body.getBytes(StandardCharsets.UTF_8); + exchange.sendResponseHeaders(200, out.length); + exchange.getResponseBody().write(out); + } + exchange.close(); + }); + server.start(); + return server; + } + + @Test + public void testRedirectIsFollowedAndRevalidated() throws IOException { + Map r = new HashMap<>(); + r.put("/start", "/final"); + HttpServer server = startServer(r, "arrived"); + try { + URL url = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/start"); + // allow loopback so the redirect loop itself runs + try (InputStream is = + Utils.openHttpConnection(url, "GET", Collections.emptyMap(), h -> true) + .getInputStream()) { + Assert.assertEquals( + new String(is.readAllBytes(), StandardCharsets.UTF_8), "arrived"); + } + } finally { + server.stop(0); + } + } + + @Test + public void testRedirectToDisallowedHostIsRejected() throws IOException { + Map r = new HashMap<>(); + r.put("/start", "http://169.254.169.254/latest/meta-data/x.tar.gz"); + HttpServer server = startServer(r, "unused"); + try { + URL url = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/start"); + // first hop allowed, redirect target must still be re-checked and refused + IOException e = + Assert.expectThrows( + IOException.class, + () -> + Utils.openHttpConnection( + url, + "GET", + Collections.emptyMap(), + h -> !"169.254.169.254".equals(h))); + Assert.assertTrue( + e.getMessage().contains("non-public"), "unexpected: " + e.getMessage()); + } finally { + server.stop(0); + } + } + + @Test + public void testRedirectToUnsupportedSchemeIsRejected() throws IOException { + Map r = new HashMap<>(); + r.put("/start", "ftp://ftp.example.com/model.tar.gz"); + HttpServer server = startServer(r, "unused"); + try { + URL url = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/start"); + IOException e = + Assert.expectThrows( + IOException.class, + () -> + Utils.openHttpConnection( + url, "GET", Collections.emptyMap(), h -> true)); + Assert.assertTrue( + e.getMessage().contains("URL protocol"), "unexpected: " + e.getMessage()); + } finally { + server.stop(0); + } + } + + @Test + public void testRedirectWithoutLocationIsRejected() throws IOException { + Map r = new HashMap<>(); + r.put("/start", ""); + HttpServer server = startServer(r, "unused"); + try { + URL url = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/start"); + IOException e = + Assert.expectThrows( + IOException.class, + () -> + Utils.openHttpConnection( + url, "GET", Collections.emptyMap(), h -> true)); + Assert.assertTrue( + e.getMessage().contains("no Location"), "unexpected: " + e.getMessage()); + } finally { + server.stop(0); + } + } + + @Test + public void testTooManyRedirectsIsRejected() throws IOException { + Map r = new HashMap<>(); + for (int i = 0; i < 10; ++i) { + r.put("/hop" + i, "/hop" + (i + 1)); + } + HttpServer server = startServer(r, "unused"); + try { + URL url = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/hop0"); + IOException e = + Assert.expectThrows( + IOException.class, + () -> + Utils.openHttpConnection( + url, "GET", Collections.emptyMap(), h -> true)); + Assert.assertTrue( + e.getMessage().contains("Too many redirects"), "unexpected: " + e.getMessage()); + } finally { + server.stop(0); + } + } + + @Test + public void testInsecureOptOutAllowsNonPublicHttpHost() throws IOException { + HttpServer server = startServer(new HashMap<>(), "opted-out"); + try { + System.setProperty("ai.djl.allow_insecure_url", "true"); + URL url = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/data"); + try (InputStream is = Utils.openUrl(url)) { + Assert.assertEquals( + new String(is.readAllBytes(), StandardCharsets.UTF_8), "opted-out"); + } + } finally { + server.stop(0); + } + } + + @Test + public void testUnknownHostIsNotPublic() { + Assert.assertFalse(Utils.isPublicHost("no-such-host.invalid")); + } + + @Test + public void testIpv6UniqueLocalAddressIsNotPublic() { + // fc00::/7 is the IPv6 unique-local range. InetAddress#isSiteLocalAddress only covers the + // deprecated fec0::/10, so these must be rejected explicitly. + Assert.assertFalse(Utils.isPublicHost("fd00::1")); + Assert.assertFalse(Utils.isPublicHost("fc00::1")); + Assert.assertFalse(Utils.isPublicHost("fdff:ffff::1")); + // link-local and loopback IPv6 remain rejected, and a public IPv6 address is allowed + Assert.assertFalse(Utils.isPublicHost("fe80::1")); + Assert.assertFalse(Utils.isPublicHost("::1")); + Assert.assertTrue(Utils.isPublicHost("2001:4860:4860::8888")); + } + + @Test + public void testOptOutAppliesToHeadConnection() throws IOException { + // The shared connection helper must honor the opt-out too, otherwise callers that only need + // response metadata (the content-length probe) stay blocked when the flag is set. + HttpServer server = startServer(new HashMap<>(), "head-ok"); + try { + URL url = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/data"); + // blocked by default: loopback is not public + Assert.assertThrows( + IOException.class, + () -> Utils.openHttpConnection(url, "HEAD", Collections.emptyMap())); + System.setProperty("ai.djl.allow_insecure_url", "true"); + HttpURLConnection conn = Utils.openHttpConnection(url, "HEAD", Collections.emptyMap()); + try { + Assert.assertEquals(conn.getResponseCode(), 200); + } finally { + conn.disconnect(); + } + } finally { + server.stop(0); + } + } + + @Test + public void testSameOriginTreatsImplicitAndExplicitDefaultPortAsEqual() + throws MalformedURLException { + // URL#getPort returns -1 for an implicit port, so comparing raw ports would classify a + // redirect that only makes the default port explicit as cross-origin and drop credentials. + Assert.assertTrue( + Utils.sameOrigin( + new URL("https://example.com/a"), new URL("https://example.com:443/b"))); + Assert.assertTrue( + Utils.sameOrigin( + new URL("http://example.com/a"), new URL("http://example.com:80/b"))); + // A genuinely different port, scheme or host is still cross-origin. + Assert.assertFalse( + Utils.sameOrigin( + new URL("https://example.com/a"), new URL("https://example.com:8443/b"))); + Assert.assertFalse( + Utils.sameOrigin( + new URL("https://example.com/a"), new URL("http://example.com/b"))); + Assert.assertFalse( + Utils.sameOrigin( + new URL("https://example.com/a"), new URL("https://other.example/b"))); + } + + @Test + public void testCrossOriginRedirectDropsCredentialsAndAllowsNullHeaderValue() + throws IOException { + // Two things at once, because they exercise the same hop: the credential header must not + // reach the redirect target, and a caller's header map may legally hold a null value. + Map r = new HashMap<>(); + // Written by the server's handler thread and read by this thread, so it must be safe for + // concurrent access. Absent headers are left out rather than stored as a "null" string, so + // the assertions below test presence directly. + Map seen = new ConcurrentHashMap<>(); + HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + int port = server.getAddress().getPort(); + // 127.0.0.1 -> localhost is a different host string, so the hop is cross-origin. + r.put("/start", "http://localhost:" + port + "/final"); + server.createContext( + "/", + exchange -> { + String path = exchange.getRequestURI().getPath(); + String target = r.get(path); + if (target != null) { + exchange.getResponseHeaders().add("Location", target); + exchange.sendResponseHeaders(302, -1); + } else { + // Record what the final hop actually received. + record(seen, exchange, "auth", "Authorization"); + record(seen, exchange, "cookie", "Cookie"); + record(seen, exchange, "keep", "X-Keep"); + byte[] out = "arrived".getBytes(StandardCharsets.UTF_8); + exchange.sendResponseHeaders(200, out.length); + exchange.getResponseBody().write(out); + } + exchange.close(); + }); + server.start(); + try { + Map headers = new HashMap<>(); + headers.put("Authorization", "Bearer secret"); + headers.put("Cookie", "session=abc"); + headers.put("X-Keep", "kept"); + headers.put("X-Optional", null); + URL url = new URL("http://127.0.0.1:" + port + "/start"); + try (InputStream is = + Utils.openHttpConnection(url, "GET", headers, h -> true).getInputStream()) { + Assert.assertEquals( + new String(is.readAllBytes(), StandardCharsets.UTF_8), "arrived"); + } + // "keep" arriving proves the final hop was reached, so the absence checks below are + // not passing merely because the redirect never completed. + Assert.assertEquals(seen.get("keep"), "kept", "other headers must still be sent"); + Assert.assertFalse( + seen.containsKey("auth"), "Authorization must not be sent across origins"); + Assert.assertFalse( + seen.containsKey("cookie"), "Cookie must not be sent across origins"); + // Proxy-Authorization is deliberately not asserted here: HttpURLConnection treats it as + // a restricted header and never transmits it, so an absence check would pass whether or + // not the strip works. It is covered by testCredentialHeadersAreRecognized instead. + } finally { + server.stop(0); + } + } + + @Test + public void testSameOriginRedirectKeepsCredentials() throws IOException { + // The counterpart to the test above: a same-origin hop must not lose the header. + Map r = new HashMap<>(); + Map seen = new ConcurrentHashMap<>(); + HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + r.put("/start", "/final"); + server.createContext( + "/", + exchange -> { + String target = r.get(exchange.getRequestURI().getPath()); + if (target != null) { + exchange.getResponseHeaders().add("Location", target); + exchange.sendResponseHeaders(302, -1); + } else { + record(seen, exchange, "auth", "Authorization"); + record(seen, exchange, "cookie", "Cookie"); + exchange.sendResponseHeaders(200, -1); + } + exchange.close(); + }); + server.start(); + try { + Map headers = new HashMap<>(); + headers.put("Authorization", "Bearer secret"); + headers.put("Cookie", "session=abc"); + URL url = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/start"); + HttpURLConnection conn = Utils.openHttpConnection(url, "GET", headers, h -> true); + try { + Assert.assertEquals(conn.getResponseCode(), 200); + } finally { + conn.disconnect(); + } + // All three must survive a same-origin hop. This is also the control that keeps the + // cross-origin test honest: it proves the connection can send these headers at all, so + // their absence there is the strip working rather than the header never being set. + Assert.assertEquals(seen.get("auth"), "Bearer secret"); + Assert.assertEquals(seen.get("cookie"), "session=abc"); + } finally { + server.stop(0); + } + } + + @Test + public void testNullHeadersAreTreatedAsEmpty() throws IOException { + // openHttpConnection is new public API; null headers must not surface as an NPE. There are + // two independent places the header map is read -- the validated redirect loop and the + // opt-out branch -- and both must tolerate it. A default-path call to a non-public host is + // NOT one of them: it is rejected before any header is read, so it would pass either way. + HttpServer server = startServer(new HashMap<>(), "no-headers"); + try { + URL url = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/data"); + // Validated path, reached with a permissive host check so the loop body actually runs. + // This is the branch production traffic takes. + HttpURLConnection validated = Utils.openHttpConnection(url, "HEAD", null, h -> true); + try { + Assert.assertEquals(validated.getResponseCode(), 200); + } finally { + validated.disconnect(); + } + // Opt-out branch: a separate header loop, so cover it separately. + System.setProperty("ai.djl.allow_insecure_url", "true"); + HttpURLConnection conn = Utils.openHttpConnection(url, "HEAD", null); + try { + Assert.assertEquals(conn.getResponseCode(), 200); + } finally { + conn.disconnect(); + } + } finally { + server.stop(0); + } + } + + @Test + public void testSchemeDowngradeIsRecognized() throws MalformedURLException { + // Checked directly rather than over a socket: pointing an https URL at a plaintext listener + // blocks in the TLS handshake, since only the connect phase is bounded by a timeout. + Assert.assertTrue( + Utils.isSchemeDowngrade(new URL("https://host/a"), new URL("http://host/b"))); + Assert.assertTrue( + Utils.isSchemeDowngrade(new URL("https://host/a"), new URL("http://other.host/b"))); + // Same scheme, and the http -> https upgrade, must both be allowed. + Assert.assertFalse( + Utils.isSchemeDowngrade( + new URL("https://host/a"), new URL("https://other.host/b"))); + Assert.assertFalse( + Utils.isSchemeDowngrade(new URL("http://host/a"), new URL("https://host/b"))); + Assert.assertFalse( + Utils.isSchemeDowngrade(new URL("http://host/a"), new URL("http://host/b"))); + } + + @Test + public void testUnfollowedRedirectStatusIsRejected() throws IOException { + // 300/304/305 are 3xx but are not followed. getInputStream() succeeds for any code below + // 400, so handing the connection back would let the negotiation page be written out as the + // artifact. + HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + server.createContext( + "/", + exchange -> { + byte[] out = "CHOICES-PAGE".getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Location", "/other"); + exchange.sendResponseHeaders(300, out.length); + exchange.getResponseBody().write(out); + exchange.close(); + }); + server.start(); + try { + URL url = new URL("http://127.0.0.1:" + server.getAddress().getPort() + "/x.tar.gz"); + IOException e = + Assert.expectThrows( + IOException.class, + () -> Utils.openHttpConnection(url, "GET", null, h -> true)); + Assert.assertTrue( + e.getMessage().contains("Unexpected redirect status 300"), + "unexpected: " + e.getMessage()); + } finally { + server.stop(0); + } + } + + @Test + public void testCarrierGradeNatIsNotPublic() { + // 100.64.0.0/10 is standard Kubernetes secondary pod CIDR space, so an internal service can + // be addressed there, and InetAddress has no predicate for it. + Assert.assertFalse(Utils.isPublicHost("100.64.0.1")); + Assert.assertFalse(Utils.isPublicHost("100.127.255.254")); + // Addresses either side of the range must not be caught by the arithmetic. + Assert.assertTrue(Utils.isPublicHost("100.63.255.255")); + Assert.assertTrue(Utils.isPublicHost("100.128.0.1")); + } + + @Test + public void testCredentialHeadersAreRecognized() { + // All three names must be recognized, case-insensitively. Proxy-Authorization cannot be + // asserted over a real connection because HttpURLConnection refuses to transmit it, so this + // is where that branch is covered. + Assert.assertTrue(Utils.isCredentialHeader("Authorization")); + Assert.assertTrue(Utils.isCredentialHeader("authorization")); + Assert.assertTrue(Utils.isCredentialHeader("Cookie")); + Assert.assertTrue(Utils.isCredentialHeader("COOKIE")); + Assert.assertTrue(Utils.isCredentialHeader("Proxy-Authorization")); + Assert.assertTrue(Utils.isCredentialHeader("proxy-authorization")); + // Headers that must keep following redirects. + Assert.assertFalse(Utils.isCredentialHeader("Accept")); + Assert.assertFalse(Utils.isCredentialHeader("User-Agent")); + Assert.assertFalse(Utils.isCredentialHeader("X-Authorization")); + } + + @Test + public void testCompilationSkippedQuietlyWithoutBundledSources() throws IOException { + // No bundled sources is the common case and must be a silent no-op. + Path dir = Files.createTempDirectory("djl-nosrc"); + Path classes = Files.createDirectories(dir.resolve("classes")); + try { + ClassLoaderUtils.compileJavaClass(classes); + // Files.list rather than File.list, which returns null instead of throwing. + try (Stream entries = Files.list(classes)) { + Assert.assertEquals(entries.count(), 0L); + } + } finally { + Files.deleteIfExists(classes); + Files.deleteIfExists(dir); + } + } + + /** + * Records a request header under {@code key} only when it was actually sent, so a caller can + * distinguish absent from present-but-empty. + */ + private static void record( + Map seen, HttpExchange exchange, String key, String header) { + String value = exchange.getRequestHeaders().getFirst(header); + if (value != null) { + seen.put(key, value); + } + } +} diff --git a/docs/create_serving_ready_model.md b/docs/create_serving_ready_model.md index a36f086b20d..16193340fe1 100644 --- a/docs/create_serving_ready_model.md +++ b/docs/create_serving_ready_model.md @@ -33,7 +33,11 @@ DJL can load Translator from the following source: - from jar files directly locate in `libs` folder - from compiled java .class file in `libs/classes` folder -- DJL can compile .java files in `libs/classes` folder at runtime and load compiled class +- DJL can compile .java files in `libs/classes` folder at runtime and load compiled class. + Compiling sources at model-load time is **opt-in**: set `DJL_COMPILE_JAVA=true` or + `-Dai.djl.compile_java=true` when loading models from a trusted source. Without it, a model that + bundles `.java` sources fails to load with a message naming the flag. Shipping a precompiled + `.class` or `.jar` instead needs no flag and is the recommended option. ## Configure data processing based on standard Translator DJL provides several built-in Translator for well-know ML applications, such as `Image Classification` diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/CustomTranslatorTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/CustomTranslatorTest.java index d651311f02c..9ed8dc18375 100644 --- a/integration/src/main/java/ai/djl/integration/tests/model_zoo/CustomTranslatorTest.java +++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/CustomTranslatorTest.java @@ -40,6 +40,7 @@ import com.google.gson.reflect.TypeToken; import org.testng.Assert; +import org.testng.SkipException; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -146,15 +147,29 @@ public void testImageClassificationTranslator() Path destFile = classesDir.resolve("MyTranslator.java"); Files.copy(srcFile, destFile, StandardCopyOption.REPLACE_EXISTING); - // load translator from classes folder - runImageClassification(Application.UNDEFINED, null, "MyTranslator"); - - Path jarFile = libsDir.resolve("example.jar"); - ZipUtils.zip(classesDir, jarFile, false); - Utils.deleteQuietly(classesDir); - - // load translator from jar file - runImageClassification(Application.UNDEFINED, null, "MyTranslator"); + // Compiling a bundled .java translator is opt-in; enable it for the cases below, which + // exercise that path (the jar case is built from the compiled output). The environment + // variable takes precedence over the system property, so an environment that pins it off + // would make the property below a no-op and the assertions fail for an unrelated reason. + String envOptIn = Utils.getenv("DJL_COMPILE_JAVA"); + if (envOptIn != null && !Boolean.parseBoolean(envOptIn)) { + throw new SkipException( + "DJL_COMPILE_JAVA is set to " + envOptIn + " in the environment"); + } + System.setProperty("ai.djl.compile_java", "true"); + try { + // load translator from classes folder + runImageClassification(Application.UNDEFINED, null, "MyTranslator"); + + Path jarFile = libsDir.resolve("example.jar"); + ZipUtils.zip(classesDir, jarFile, false); + Utils.deleteQuietly(classesDir); + + // load translator from jar file + runImageClassification(Application.UNDEFINED, null, "MyTranslator"); + } finally { + System.clearProperty("ai.djl.compile_java"); + } } @Test