From 862675bc4bd2b3190163c9e81462036fc755aa07 Mon Sep 17 00:00:00 2001 From: PG1204 Date: Sun, 25 Oct 2026 10:29:28 -0700 Subject: [PATCH] fix(amber): re-validate every redirect hop in the Hugging Face media proxy --- .../resource/HuggingFaceModelResource.scala | 205 ++++++++-- .../HuggingFaceModelResourceSpec.scala | 377 +++++++++++++++++- 2 files changed, 545 insertions(+), 37 deletions(-) diff --git a/amber/src/main/scala/org/apache/texera/web/resource/HuggingFaceModelResource.scala b/amber/src/main/scala/org/apache/texera/web/resource/HuggingFaceModelResource.scala index 0decacab5fa..18594d92f6d 100644 --- a/amber/src/main/scala/org/apache/texera/web/resource/HuggingFaceModelResource.scala +++ b/amber/src/main/scala/org/apache/texera/web/resource/HuggingFaceModelResource.scala @@ -22,16 +22,19 @@ package org.apache.texera.web.resource import com.fasterxml.jackson.core.`type`.TypeReference import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} import com.google.common.cache.{Cache, CacheBuilder} -import kong.unirest.Unirest +import kong.unirest.{RawResponse, Unirest, UnirestInstance} import org.slf4j.{Logger, LoggerFactory} -import java.net.URI +import java.net.{InetAddress, URI} +import java.util.Locale import java.util.concurrent.{Callable, ForkJoinPool, TimeUnit} import java.util.stream.Collectors import javax.annotation.security.RolesAllowed import javax.ws.rs._ import javax.ws.rs.core.{MediaType, Response} +import scala.annotation.tailrec import scala.jdk.CollectionConverters._ +import scala.util.{Failure, Success, Try} /** * REST resource that proxies the Hugging Face Hub API for the HuggingFace operator. @@ -96,8 +99,9 @@ class HuggingFaceModelResource { /** * Proxies a remote media URL to bypass browser CORS for HF inference responses. - * Only http(s) URLs whose host is in ALLOWED_MEDIA_HOST_SUFFIXES are accepted, - * blocking SSRF probes against internal services. + * Only https URLs whose host is in ALLOWED_MEDIA_HOST_SUFFIXES and resolves to + * public addresses are fetched, blocking SSRF probes against internal services. + * Redirects are followed by hand and every hop is held to the same rules. */ @GET @Path("/media-proxy") @@ -107,32 +111,8 @@ class HuggingFaceModelResource { if (trimmedUrl.isEmpty) { return errorResponse(Response.Status.BAD_REQUEST, "Media URL is required.") } - if (!trimmedUrl.startsWith("http://") && !trimmedUrl.startsWith("https://")) { - return errorResponse( - Response.Status.BAD_REQUEST, - "Only http(s) media URLs are supported." - ) - } - - val parsed = - try { new URI(trimmedUrl) } - catch { case _: Exception => null } - if (parsed == null || parsed.getHost == null || !isAllowedMediaHost(parsed.getHost)) { - return errorResponse(Response.Status.FORBIDDEN, "Media URL host is not allowed.") - } - - // Stream the upstream response via asObject(Function) so - // we never have to materialise it into a heap-resident byte[] before we - // can enforce the size cap. The function returns a MediaProxyOutcome - // that this method then converts into a Jersey Response. - val outcome = Unirest - .get(trimmedUrl) - .connectTimeout(CONNECT_TIMEOUT_MS) - .socketTimeout(SOCKET_TIMEOUT_LONG_MS) - .asObject((raw: kong.unirest.RawResponse) => streamMediaWithCap(raw)) - .getBody - outcome match { + fetchMedia(trimmedUrl) match { case MediaProxyOk(bytes, contentType) => Response.ok(bytes, contentType.getOrElse(MediaType.APPLICATION_OCTET_STREAM)).build() case MediaProxyError(status, message) => @@ -145,12 +125,91 @@ class HuggingFaceModelResource { } } + /** + * Fetch `url`, following at most MAX_MEDIA_PROXY_REDIRECTS redirects. The + * client never follows a redirect itself: each one comes back here, and the + * next hop is checked by `checkMediaUrl` before it is requested. + */ + @tailrec + private def fetchMedia(url: String, redirectsFollowed: Int = 0): MediaProxyOutcome = + checkMediaUrl(url) match { + case Some(rejection) => rejection + case None => + // Stream the upstream response via asObject(Function) so + // we never have to materialise it into a heap-resident byte[] before we + // can enforce the size cap. + val hop = mediaProxyHttp + .get(url) + .connectTimeout(CONNECT_TIMEOUT_MS) + .socketTimeout(SOCKET_TIMEOUT_LONG_MS) + .asObject((raw: RawResponse) => readMediaHop(raw)) + .getBody + + hop match { + case Right(outcome) => outcome + case Left(_) if redirectsFollowed >= MAX_MEDIA_PROXY_REDIRECTS => + logger.warn(s"Upstream media exceeded $MAX_MEDIA_PROXY_REDIRECTS redirects; rejecting.") + UPSTREAM_MEDIA_FAILED + case Left(location) => + resolveRedirect(url, location) match { + case Some(next) => fetchMedia(next, redirectsFollowed + 1) + case None => + logger.warn(s"Upstream media sent an unusable redirect Location '$location'.") + UPSTREAM_MEDIA_FAILED + } + } + } + + /** A redirect's Location (Left), or the streamed media outcome (Right). */ + private def readMediaHop(raw: RawResponse): Either[String, MediaProxyOutcome] = + if (REDIRECT_STATUSES.contains(raw.getStatus)) Left(raw.getHeaders.getFirst("Location")) + else Right(streamMediaWithCap(raw)) + + /** Resolve a redirect's Location, which may be relative, against the URL that sent it. */ + private def resolveRedirect(current: String, location: String): Option[String] = + Option(location).map(_.trim).filter(_.nonEmpty).flatMap { loc => + Try(new URI(current).resolve(loc).toString).toOption + } + + /** + * Check one URL before it is requested: the original and every redirect hop. + * Returns the error to send back, or None when the URL may be fetched: https, + * an allowlisted host, and every address the host resolves to public. The + * address check runs before the request, so it narrows but does not close + * DNS rebinding (the HTTP client resolves the host again to connect). + */ + private def checkMediaUrl(url: String): Option[MediaProxyError] = { + val lower = url.toLowerCase(Locale.ROOT) + val httpsOnly = + MediaProxyError( + Response.Status.BAD_REQUEST.getStatusCode, + "Only https media URLs are supported." + ) + val hostNotAllowed = + MediaProxyError(Response.Status.FORBIDDEN.getStatusCode, "Media URL host is not allowed.") + + if (!lower.startsWith("http://") && !lower.startsWith("https://")) return Some(httpsOnly) + val host = Try(new URI(url)).toOption.flatMap(uri => Option(uri.getHost)) + if (!host.exists(isAllowedMediaHost)) return Some(hostNotAllowed) + if (!lower.startsWith("https://")) return Some(httpsOnly) + + Try(resolveMediaHost(host.get)) match { + case Failure(e) => + logger.warn(s"Could not resolve media host '${host.get}': ${e.getMessage}") + Some(UPSTREAM_MEDIA_FAILED) + case Success(addresses) if addresses.isEmpty || !addresses.forall(isPublicAddress) => + logger.warn(s"Media host '${host.get}' resolves to a non-public address; rejecting.") + Some(hostNotAllowed) + case Success(_) => None + } + } + /** * Read an upstream media response with a hard size cap. Aborts early both * when the declared Content-Length exceeds the cap and when the body crosses * it mid-read (in case the upstream lies about Content-Length or omits it). */ - private def streamMediaWithCap(raw: kong.unirest.RawResponse): MediaProxyOutcome = { + private def streamMediaWithCap(raw: RawResponse): MediaProxyOutcome = { if (raw.getStatus != 200) { logger.warn(s"Upstream media fetch returned ${raw.getStatus}: ${raw.getStatusText}") return MediaProxyError(raw.getStatus, "Failed to fetch remote media.") @@ -532,15 +591,101 @@ object HuggingFaceModelResource { private[resource] case class MediaProxyError(status: Int, message: String) extends MediaProxyOutcome + /** The upstream could not be reached or answered with an unusable redirect. */ + private val UPSTREAM_MEDIA_FAILED = + MediaProxyError(Response.Status.BAD_GATEWAY.getStatusCode, "Failed to fetch remote media.") + + // ── /media-proxy HTTP client ── + // Its own Unirest instance, so it can refuse to follow redirects without + // changing the global config the Hugging Face Hub calls use. fetchMedia + // follows redirects by hand and re-checks every hop. + private[resource] val mediaProxyHttp: UnirestInstance = { + val instance = Unirest.spawnInstance() + instance.config().followRedirects(false) + instance + } + + private[resource] val MAX_MEDIA_PROXY_REDIRECTS = 5 + private val REDIRECT_STATUSES = Set(301, 302, 303, 307, 308) + + // How media hosts are resolved. A var only so tests can resolve without DNS. + @volatile private[resource] var resolveMediaHost: String => Seq[InetAddress] = + host => InetAddress.getAllByName(host).toSeq + // ── /media-proxy allowlist (SSRF protection) ── // Add new hosts here when integrating with a new HF inference provider. private val ALLOWED_MEDIA_HOST_SUFFIXES: Set[String] = Set( "huggingface.co", + "hf.co", // Hub file downloads (/resolve/...) redirect to its CDN "fal.media", "replicate.delivery", "replicate.com" ) + // ── /media-proxy resolved-address check (SSRF protection) ── + // Address blocks a media host must not resolve to: every range that is not + // globally routable (the IANA special-purpose registries) plus multicast, + // mirroring the is_global check the operator's generated Python applies to + // its own remote fetches. + private val NON_PUBLIC_ADDRESS_BLOCKS: Seq[(Array[Byte], Int)] = Seq( + "0.0.0.0/8", // "this" network + "10.0.0.0/8", // private + "100.64.0.0/10", // shared address space (carrier-grade NAT) + "127.0.0.0/8", // loopback + "169.254.0.0/16", // link-local, including the cloud metadata endpoint + "172.16.0.0/12", // private + "192.0.0.0/24", // IETF protocol assignments + "192.0.2.0/24", // documentation + "192.168.0.0/16", // private + "198.18.0.0/15", // benchmarking + "198.51.100.0/24", // documentation + "203.0.113.0/24", // documentation + "224.0.0.0/4", // multicast + "240.0.0.0/4", // reserved, including broadcast + "64:ff9b:1::/48", // local-use IPv4/IPv6 translation + "100::/64", // discard-only + "2001::/23", // IETF protocol assignments + "2001:db8::/32", // documentation + "2002::/16", // 6to4 + "fc00::/7", // unique-local + "fe80::/10", // link-local + "fec0::/10", // site-local (deprecated) + "ff00::/8" // multicast + ).map { block => + val parts = block.split('/') + (InetAddress.getByName(parts(0)).getAddress, parts(1).toInt) + } + + private val NAT64_PREFIX: Array[Byte] = InetAddress.getByName("64:ff9b::").getAddress.take(12) + + /** True when `address` is globally routable: outside every NON_PUBLIC_ADDRESS_BLOCKS entry. */ + private[resource] def isPublicAddress(address: InetAddress): Boolean = { + val bytes = embeddedIpv4(address.getAddress).getOrElse(address.getAddress) + !NON_PUBLIC_ADDRESS_BLOCKS.exists { + case (network, prefixLength) => inBlock(bytes, network, prefixLength) + } + } + + /** + * The IPv4 address inside an IPv4-mapped (::ffff:a.b.c.d), IPv4-compatible + * (::a.b.c.d) or NAT64 (64:ff9b::a.b.c.d) IPv6 address, so that it is judged + * by the IPv4 rules. This also maps :: and ::1 into 0.0.0.0/8. + */ + private def embeddedIpv4(bytes: Array[Byte]): Option[Array[Byte]] = { + if (bytes.length != 16) return None + val zeroPrefix = bytes.take(10).forall(_ == 0) + val mapped = zeroPrefix && bytes(10) == -1 && bytes(11) == -1 + val compatible = zeroPrefix && bytes(10) == 0 && bytes(11) == 0 + val nat64 = bytes.take(12).sameElements(NAT64_PREFIX) + if (mapped || compatible || nat64) Some(bytes.drop(12)) else None + } + + private def inBlock(address: Array[Byte], network: Array[Byte], prefixLength: Int): Boolean = + address.length == network.length && (0 until prefixLength).forall { bit => + val mask = 0x80 >>> (bit % 8) + (address(bit / 8) & mask) == (network(bit / 8) & mask) + } + /** Allow exact host or subdomain of any entry in ALLOWED_MEDIA_HOST_SUFFIXES. */ private[resource] def isAllowedMediaHost(host: String): Boolean = { if (host == null || host.isEmpty) return false diff --git a/amber/src/test/scala/org/apache/texera/web/resource/HuggingFaceModelResourceSpec.scala b/amber/src/test/scala/org/apache/texera/web/resource/HuggingFaceModelResourceSpec.scala index a4e0d14e5e0..0e0f61f3317 100644 --- a/amber/src/test/scala/org/apache/texera/web/resource/HuggingFaceModelResourceSpec.scala +++ b/amber/src/test/scala/org/apache/texera/web/resource/HuggingFaceModelResourceSpec.scala @@ -20,6 +20,7 @@ package org.apache.texera.web.resource import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} +import com.sun.net.httpserver.HttpServer import kong.unirest.{ Client, Headers, @@ -35,7 +36,9 @@ import org.scalatest.BeforeAndAfterEach import org.scalatest.funsuite.AnyFunSuite import java.io.{ByteArrayInputStream, InputStream, InputStreamReader} +import java.net.{Inet6Address, InetAddress, InetSocketAddress, UnknownHostException} import java.nio.charset.StandardCharsets +import java.util.concurrent.atomic.AtomicInteger import java.util.function.{Function => JFunction} import javax.ws.rs.core.Response import scala.collection.mutable @@ -46,9 +49,12 @@ import scala.jdk.CollectionConverters._ * caching, and filesystem behavior of the resource. * * The HF Hub URLs are hard-coded in the resource, so the upstream calls are - * exercised by swapping Unirest's global HTTP client for an in-process stub - * (`StubClient` below) that serves canned responses and records the outgoing - * requests. Nothing in this suite touches the network. + * exercised by swapping Unirest's HTTP clients (the global one and the media + * proxy's own instance) for an in-process stub (`StubClient` below) that + * serves canned responses and records the outgoing requests. Media-host DNS + * lookups go through a fake resolver. Nothing in this suite leaves the + * machine: the one test that exercises the real HTTP client talks to a + * loopback server. */ class HuggingFaceModelResourceSpec extends AnyFunSuite with BeforeAndAfterEach { @@ -56,6 +62,7 @@ class HuggingFaceModelResourceSpec extends AnyFunSuite with BeforeAndAfterEach { private val mapper = new ObjectMapper() private var resource: HuggingFaceModelResource = _ + private val realMediaHostResolver = resolveMediaHost override def beforeEach(): Unit = { resource = new HuggingFaceModelResource() @@ -67,9 +74,11 @@ class HuggingFaceModelResourceSpec extends AnyFunSuite with BeforeAndAfterEach { override def afterEach(): Unit = { modelCache.invalidateAll() taskCache.invalidateAll() - // Drop any stub HTTP client this test installed so the global Unirest - // config never leaks between tests. + // Drop any stub HTTP client or fake resolver this test installed so + // neither leaks between tests. Unirest.config().reset() + mediaProxyHttp.config().reset() + resolveMediaHost = realMediaHostResolver } // Helper: read a Response's string entity (assumes the body is a String). @@ -192,6 +201,89 @@ class HuggingFaceModelResourceSpec extends AnyFunSuite with BeforeAndAfterEach { assert(!isAllowedMediaHost("169.254.169.254")) } + test("isAllowedMediaHost accepts hf.co and its subdomains (Hub file downloads redirect there)") { + assert(isAllowedMediaHost("hf.co")) + assert(isAllowedMediaHost("us.aws.cdn.hf.co")) + assert(isAllowedMediaHost("cas-bridge.xethub.hf.co")) + } + + test("isAllowedMediaHost rejects hf.co lookalikes (leading-dot guard)") { + assert(!isAllowedMediaHost("evilhf.co")) + assert(!isAllowedMediaHost("hf.co.evil.com")) + } + + // ──────────────────────────────────────────────────────────────────────── + // isPublicAddress — what a media host may resolve to + // ──────────────────────────────────────────────────────────────────────── + + // Parsing an IP literal never touches DNS. + private def ip(literal: String): InetAddress = InetAddress.getByName(literal) + + private def assertAllPublic(literals: String*): Unit = + literals.foreach(a => assert(isPublicAddress(ip(a)), s"$a should be public")) + + private def assertNonePublic(literals: String*): Unit = + literals.foreach(a => assert(!isPublicAddress(ip(a)), s"$a should not be public")) + + test("isPublicAddress accepts globally routable IPv4 and IPv6 addresses") { + assertAllPublic("8.8.8.8", "1.1.1.1", "13.32.0.1", "2606:4700:4700::1111", "2a00:1450:4001::1") + } + + test("isPublicAddress accepts addresses just outside the non-public ranges") { + assertAllPublic("172.15.255.255", "172.32.0.0", "100.63.255.255", "100.128.0.0", "11.0.0.0") + } + + test("isPublicAddress rejects loopback and unspecified addresses") { + assertNonePublic("127.0.0.1", "127.255.255.254", "0.0.0.0", "::1", "::") + } + + test("isPublicAddress rejects the RFC 1918 private ranges") { + assertNonePublic("10.0.0.1", "172.16.0.1", "172.31.255.255", "192.168.1.1") + } + + test("isPublicAddress rejects link-local addresses, including the cloud metadata endpoint") { + assertNonePublic("169.254.169.254", "169.254.0.1", "fe80::1") + } + + test("isPublicAddress rejects the shared address space (100.64.0.0/10)") { + // 100.100.100.200 is Alibaba Cloud's metadata endpoint. + assertNonePublic("100.64.0.1", "100.100.100.200", "100.127.255.255") + } + + test( + "isPublicAddress rejects IPv6 unique-local addresses, including AWS's IPv6 metadata endpoint" + ) { + assertNonePublic("fc00::1", "fd00::1", "fd00:ec2::254") + } + + test("isPublicAddress rejects multicast, reserved and documentation ranges") { + assertNonePublic( + "224.0.0.1", + "239.255.255.250", + "240.0.0.1", + "255.255.255.255", + "192.0.2.1", + "198.51.100.1", + "203.0.113.1", + "198.18.0.1", + "ff02::1", + "2001:db8::1" + ) + } + + test("isPublicAddress judges an IPv4 address embedded in IPv6 by the IPv4 address") { + // getByName unwraps ::ffff:a.b.c.d to IPv4, so build the IPv6 form directly. + val mappedLoopback = + Inet6Address.getByAddress( + null, + Array[Byte](0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, -1, 127, 0, 0, 1), + -1 + ) + assert(!isPublicAddress(mappedLoopback)) + assertNonePublic("::7f00:1", "64:ff9b::a00:1") // IPv4-compatible 127.0.0.1; NAT64 to 10.0.0.1 + assertAllPublic("64:ff9b::808:808") // NAT64 to 8.8.8.8 + } + // ──────────────────────────────────────────────────────────────────────── // errorJson — JSON escaping // ──────────────────────────────────────────────────────────────────────── @@ -496,15 +588,29 @@ class HuggingFaceModelResourceSpec extends AnyFunSuite with BeforeAndAfterEach { override def registerShutdownHook(): Unit = () } - /** Install a stub HTTP client for the duration of `body`. */ + /** A public address for the fake resolver to hand back by default. */ + private val publicAddress = InetAddress.getByName("8.8.8.8") + + /** + * Install a stub HTTP client, on both the global Unirest config and the + * media proxy's own instance, for the duration of `body`. Media hosts + * resolve to a public address unless the test sets `resolveMediaHost`. + */ private def withStubClient[A](handler: StubRequest => StubHttpResponse)( body: StubClient => A ): A = { val client = new StubClient(handler) Unirest.config().reset() Unirest.config().httpClient(client) + mediaProxyHttp.config().reset() + mediaProxyHttp.config().httpClient(client) + resolveMediaHost = _ => Seq(publicAddress) try body(client) - finally Unirest.config().reset() + finally { + Unirest.config().reset() + mediaProxyHttp.config().reset() + resolveMediaHost = realMediaHostResolver + } } /** A raw HF model payload; omitted fields exercise the defaulting branches. */ @@ -1070,4 +1176,261 @@ class HuggingFaceModelResourceSpec extends AnyFunSuite with BeforeAndAfterEach { val node = mapper.readTree(entityString(response)) assert(node.get("error").asText() == "Failed to proxy remote media.") } + + // ──────────────────────────────────────────────────────────────────────── + // proxyRemoteMedia — scheme, resolved address and redirect re-validation + // ──────────────────────────────────────────────────────────────────────── + + private val mediaPayload = "PNG-fake-image-bytes".getBytes(StandardCharsets.UTF_8) + + private def mediaOk: StubHttpResponse = bytesResponse(mediaPayload, "image/png") + + private def redirectTo(location: String, status: Int = 302): StubHttpResponse = + bytesResponse(Array.emptyByteArray, "text/html", status, Map("Location" -> location)) + + test("proxyRemoteMedia refuses plain http even for an allowlisted host") { + val client = withStubClient(_ => mediaOk) { c => + val response = resource.proxyRemoteMedia("http://cdn-lfs.huggingface.co/generated/image.png") + assert(response.getStatus == 400) + assertErrorBody(response) + c + } + + assert(client.requests.isEmpty) + } + + test("proxyRemoteMedia accepts an upper-case https scheme") { + val response = withStubClient(_ => mediaOk) { _ => + resource.proxyRemoteMedia("HTTPS://cdn-lfs.huggingface.co/generated/image.png") + } + + assert(response.getStatus == 200) + assert(entityBytes(response).sameElements(mediaPayload)) + } + + test("proxyRemoteMedia refuses an allowlisted host that resolves to a private address") { + val client = withStubClient(_ => mediaOk) { c => + resolveMediaHost = _ => Seq(ip("10.0.0.7")) + val response = resource.proxyRemoteMedia(allowedMediaUrl) + assert(response.getStatus == 403) + assertErrorBody(response) + c + } + + assert(client.requests.isEmpty) + } + + test("proxyRemoteMedia refuses a host when any one of its addresses is not public") { + val client = withStubClient(_ => mediaOk) { c => + resolveMediaHost = _ => Seq(publicAddress, ip("127.0.0.1")) + assert(resource.proxyRemoteMedia(allowedMediaUrl).getStatus == 403) + c + } + + assert(client.requests.isEmpty) + } + + test("proxyRemoteMedia refuses a host that resolves to no addresses at all") { + val client = withStubClient(_ => mediaOk) { c => + resolveMediaHost = _ => Seq.empty + assert(resource.proxyRemoteMedia(allowedMediaUrl).getStatus == 403) + c + } + + assert(client.requests.isEmpty) + } + + test("proxyRemoteMedia returns 502 when the media host cannot be resolved") { + val client = withStubClient(_ => mediaOk) { c => + resolveMediaHost = host => throw new UnknownHostException(host) + val response = resource.proxyRemoteMedia(allowedMediaUrl) + assert(response.getStatus == 502) + assertErrorBody(response) + c + } + + assert(client.requests.isEmpty) + } + + test("proxyRemoteMedia does not follow a redirect to a host off the allowlist") { + val client = withStubClient { req => + if (req.url == allowedMediaUrl) redirectTo("https://internal.example/secret") + else bytesResponse("SECRET".getBytes(StandardCharsets.UTF_8), "text/plain") + } { c => + val response = resource.proxyRemoteMedia(allowedMediaUrl) + assert(response.getStatus == 403) + assertErrorBody(response) + c + } + + assert(client.urls == List(allowedMediaUrl)) + } + + test("proxyRemoteMedia does not follow a redirect to the cloud metadata address") { + val client = withStubClient { req => + if (req.url == allowedMediaUrl) redirectTo("http://169.254.169.254/latest/meta-data/") + else bytesResponse("SECRET".getBytes(StandardCharsets.UTF_8), "text/plain") + } { c => + assert(resource.proxyRemoteMedia(allowedMediaUrl).getStatus == 403) + c + } + + assert(client.urls == List(allowedMediaUrl)) + } + + test("proxyRemoteMedia does not follow a redirect that downgrades to http") { + val client = withStubClient { req => + if (req.url == allowedMediaUrl) + redirectTo("http://cdn-lfs.huggingface.co/generated/other.png") + else mediaOk + } { c => + assert(resource.proxyRemoteMedia(allowedMediaUrl).getStatus == 400) + c + } + + assert(client.urls == List(allowedMediaUrl)) + } + + test( + "proxyRemoteMedia does not follow a redirect to an allowlisted name that resolves privately" + ) { + val client = withStubClient { req => + if (req.url == allowedMediaUrl) redirectTo("https://internal.huggingface.co/x.png") + else mediaOk + } { c => + resolveMediaHost = { + case "internal.huggingface.co" => Seq(ip("192.168.0.10")) + case _ => Seq(publicAddress) + } + assert(resource.proxyRemoteMedia(allowedMediaUrl).getStatus == 403) + c + } + + assert(client.urls == List(allowedMediaUrl)) + } + + test("proxyRemoteMedia follows a Hub file redirect to the hf.co CDN and serves the final body") { + val hubFileUrl = "https://huggingface.co/datasets/org/repo/resolve/main/cat.png" + val cdnUrl = "https://us.aws.cdn.hf.co/xet-bridge-us/abc?Expires=1&Signature=x%7Ey" + val client = withStubClient { req => + if (req.url == hubFileUrl) redirectTo(cdnUrl) else mediaOk + } { c => + val response = resource.proxyRemoteMedia(hubFileUrl) + assert(response.getStatus == 200) + assert(entityBytes(response).sameElements(mediaPayload)) + assert(response.getMediaType.toString == "image/png") + c + } + + assert(client.urls == List(hubFileUrl, cdnUrl)) + } + + test("proxyRemoteMedia follows every redirect status code") { + val finalUrl = "https://cdn-lfs.huggingface.co/generated/final.png" + Seq(301, 302, 303, 307, 308).foreach { code => + val response = withStubClient { req => + if (req.url == allowedMediaUrl) redirectTo(finalUrl, code) else mediaOk + } { _ => resource.proxyRemoteMedia(allowedMediaUrl) } + assert(response.getStatus == 200, s"status $code was not followed") + } + } + + test("proxyRemoteMedia resolves a relative redirect against the current URL") { + val client = withStubClient { req => + if (req.url == allowedMediaUrl) redirectTo("/generated/final.png") else mediaOk + } { c => + assert(resource.proxyRemoteMedia(allowedMediaUrl).getStatus == 200) + c + } + + assert( + client.urls == List(allowedMediaUrl, "https://cdn-lfs.huggingface.co/generated/final.png") + ) + } + + private def hopUrl(n: Int): String = s"https://cdn-lfs.huggingface.co/hop/$n.png" + + /** Each hop redirects to the next one until `lastHop`, which serves the media. */ + private def redirectChain(lastHop: Int)(req: StubRequest): StubHttpResponse = { + val n = req.url.stripPrefix("https://cdn-lfs.huggingface.co/hop/").stripSuffix(".png").toInt + if (n < lastHop) redirectTo(hopUrl(n + 1)) else mediaOk + } + + test("proxyRemoteMedia follows up to MAX_MEDIA_PROXY_REDIRECTS redirects") { + val client = withStubClient(redirectChain(MAX_MEDIA_PROXY_REDIRECTS)) { c => + assert(resource.proxyRemoteMedia(hopUrl(0)).getStatus == 200) + c + } + + assert(client.requests.size == MAX_MEDIA_PROXY_REDIRECTS + 1) + } + + test("proxyRemoteMedia gives up after MAX_MEDIA_PROXY_REDIRECTS redirects") { + val client = withStubClient(redirectChain(Int.MaxValue)) { c => + val response = resource.proxyRemoteMedia(hopUrl(0)) + assert(response.getStatus == 502) + assertErrorBody(response) + c + } + + assert(client.requests.size == MAX_MEDIA_PROXY_REDIRECTS + 1) + } + + test("proxyRemoteMedia returns 502 for a redirect without a Location header") { + val client = withStubClient(_ => bytesResponse(Array.emptyByteArray, "text/html", 302)) { c => + val response = resource.proxyRemoteMedia(allowedMediaUrl) + assert(response.getStatus == 502) + assertErrorBody(response) + c + } + + assert(client.requests.size == 1) + } + + test("proxyRemoteMedia returns 502 for a redirect whose Location is not a valid URL") { + val client = withStubClient(_ => redirectTo("https://cdn-lfs.huggingface.co/a b.png")) { c => + assert(resource.proxyRemoteMedia(allowedMediaUrl).getStatus == 502) + c + } + + assert(client.requests.size == 1) + } + + // ──────────────────────────────────────────────────────────────────────── + // The media proxy's own HTTP client (real client, loopback server) + // ──────────────────────────────────────────────────────────────────────── + + test("the media proxy's HTTP client does not follow redirects on its own") { + // Every hop has to come back to proxyRemoteMedia to be re-validated, so the + // real client must hand a 302 back instead of following it. + val targetHits = new AtomicInteger() + val server = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress, 0), 0) + server.createContext( + "/start", + exchange => { + exchange.getResponseHeaders.add("Location", "/target") + exchange.sendResponseHeaders(302, -1) + exchange.close() + } + ) + server.createContext( + "/target", + exchange => { + targetHits.incrementAndGet() + exchange.sendResponseHeaders(200, -1) + exchange.close() + } + ) + server.start() + try { + val port = server.getAddress.getPort + val status = mediaProxyHttp.get(s"http://127.0.0.1:$port/start").asEmpty().getStatus + assert(status == 302) + assert(targetHits.get() == 0) + } finally server.stop(0) + } + + test("the media proxy's redirect policy leaves the global Unirest config alone") { + assert(Unirest.config().getFollowRedirects) + } }