diff --git a/docs/changelog.ko.md b/docs/changelog.ko.md index 5839a1b..e07f381 100644 --- a/docs/changelog.ko.md +++ b/docs/changelog.ko.md @@ -6,6 +6,17 @@ ssrf-guard의 주요 변경 사항을 기록합니다. ## [Unreleased] +v3.1.0 작업 추적. 릴리즈된 항목은 태그 시 별도 섹션으로 이동. + +### Added + +- **`ssrf-guard-llm`** — 신규. 모든 LLM 툴 어댑터가 공유하는 JSON 트리 walk + URL 추출 + 정책 검증을 담는 framework-agnostic 코어 모듈. `ToolInputGuard` (인터페이스)와 `JsonToolInputGuard` (기본 구현) 노출. 이전에 `SsrfGuardedToolCallback` 내부에 있던 ~200줄 로직을 이쪽으로 이동. +- **`ssrf-guard-langchain4j`** — 신규. LangChain4j의 `ToolExecutor`를 wrap하는 어댑터 모듈. `ssrf-guard-springai`와 동일한 LLM 에이전트 SSRF 표면을 LangChain4j 생태계 (Java LLM 양대 프레임워크 중 다른 하나) 대응. Spring Boot 사용자는 `BeanPostProcessor`로 자동 wrap; 비-Spring/프로그래매틱 사용자용 `SsrfGuardedToolExecutors.wrap(...)` 헬퍼 제공. + +### Changed + +- **`ssrf-guard-springai`가 thin adapter (~30줄)로 리팩토링.** 새 `-llm` 모듈의 `JsonToolInputGuard`에 위임. Public API 변경 없음 — `SsrfGuardedToolCallback`의 모든 생성자/메서드 시그니처 유지. v3.0.x 소비자는 API 변경 못 느낌 — 단지 `ssrf-guard-llm`을 transitive로 받게 됨. + ## [3.0.1] — 메트릭 빈 classpath 게이트 수정 ### Fixed diff --git a/docs/changelog.md b/docs/changelog.md index e82bc19..f042098 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -6,6 +6,17 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and ## [Unreleased] +Tracking v3.1.0 work. Released items will move into their own section below when the tag lands. + +### Added + +- **`ssrf-guard-llm`** — new framework-agnostic core module holding the JSON walking + URL extraction + policy validation that every LLM tool adapter shares. Exposes `ToolInputGuard` (interface) and `JsonToolInputGuard` (default impl). Replaces ~200 lines of logic that previously lived inside `SsrfGuardedToolCallback`. +- **`ssrf-guard-langchain4j`** — new adapter module wrapping LangChain4j's `ToolExecutor`. Closes the same LLM-agent SSRF surface as `ssrf-guard-springai` but for the LangChain4j ecosystem (the other major Java LLM framework). Auto-wrap via `BeanPostProcessor` for Spring Boot users; `SsrfGuardedToolExecutors.wrap(...)` helpers for non-Spring / programmatic users. + +### Changed + +- **`ssrf-guard-springai` refactored to a thin adapter (~30 lines).** Delegates to `JsonToolInputGuard` from the new `-llm` module. Public API unchanged — every constructor and method on `SsrfGuardedToolCallback` keeps the same shape. v3.0.x consumers see no API change; they just pick up `ssrf-guard-llm` transitively. + ## [3.0.1] — Fix metrics bean classpath gate ### Fixed diff --git a/gradle.properties b/gradle.properties index c9dd6fa..bd3fbac 100644 --- a/gradle.properties +++ b/gradle.properties @@ -1,6 +1,6 @@ # Project coordinates GROUP=kr.devslab -VERSION=3.0.1-SNAPSHOT +VERSION=3.1.0-SNAPSHOT # Project metadata for Maven Central POM POM_NAME=SSRF Guard diff --git a/settings.gradle.kts b/settings.gradle.kts index f09327e..832cbd6 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -26,7 +26,9 @@ rootProject.name = "ssrf-guard" // ssrf-guard-resttemplate — Spring RestTemplate interceptor + auto-config // ssrf-guard-webclient — Spring WebFlux ExchangeFilterFunction + auto-config // ssrf-guard-feign — Spring Cloud OpenFeign interceptor + auto-config -// ssrf-guard-springai — Spring AI tool URL validation +// ssrf-guard-llm — framework-agnostic JSON tool input validation (v3.1+) +// ssrf-guard-springai — Spring AI tool URL validation (thin adapter over -llm) +// ssrf-guard-langchain4j — LangChain4j tool URL validation (thin adapter over -llm, v3.1+) // ssrf-guard-jdkhttp — java.net.http.HttpClient wrapper (no Spring) // ssrf-guard-okhttp — OkHttp interceptor + DNS // ssrf-guard — v2.0.0-compatible meta artifact (restclient + httpclient5) @@ -37,7 +39,9 @@ include( "ssrf-guard-resttemplate", "ssrf-guard-webclient", "ssrf-guard-feign", + "ssrf-guard-llm", "ssrf-guard-springai", + "ssrf-guard-langchain4j", "ssrf-guard-jdkhttp", "ssrf-guard-okhttp", "ssrf-guard", diff --git a/ssrf-guard-langchain4j/build.gradle.kts b/ssrf-guard-langchain4j/build.gradle.kts new file mode 100644 index 0000000..b459326 --- /dev/null +++ b/ssrf-guard-langchain4j/build.gradle.kts @@ -0,0 +1,80 @@ +// ssrf-guard-langchain4j — thin adapter wiring LangChain4j's ToolExecutor +// to ssrf-guard-llm. Same security model as ssrf-guard-springai; different +// framework abstraction. +// +// Closes the same SSRF surface (LLM agent fetch_url-style tools) for the +// LangChain4j community — the other major Java LLM framework. + +base { + archivesName.set("ssrf-guard-langchain4j") +} + +dependencies { + // The framework-agnostic core. Jackson follows transitively. + api(project(":ssrf-guard-llm")) + api("org.springframework.boot:spring-boot-autoconfigure") + + // LangChain4j 1.15.0 — first stable line. ToolExecutor lives in the + // main `langchain4j` artifact (the AiServices implementation), + // ToolExecutionRequest in `langchain4j-core` follows transitively. + compileOnly("dev.langchain4j:langchain4j:1.15.0") + + // Micrometer — compileOnly. The MetricsConfiguration inner class is + // gated by @ConditionalOnClass(MeterRegistry.class), so this dep is + // only ever needed at compile time (the class reference inside the + // gated inner config) — never required at runtime by consumers. + compileOnly("io.micrometer:micrometer-core") + + compileOnly("org.projectlombok:lombok") + annotationProcessor("org.projectlombok:lombok") + annotationProcessor("org.springframework.boot:spring-boot-configuration-processor") + compileOnly("com.google.code.findbugs:jsr305:3.0.2") + + testImplementation("org.springframework.boot:spring-boot-starter-test") + testImplementation("dev.langchain4j:langchain4j:1.15.0") + testRuntimeOnly("org.junit.platform:junit-platform-launcher") +} + +mavenPublishing { + coordinates( + providers.gradleProperty("GROUP").get(), + "ssrf-guard-langchain4j", + providers.gradleProperty("VERSION").get() + ) + pom { + name.set("SSRF Guard — LangChain4j") + description.set("SSRF Guard adapter for LangChain4j tool execution. Validates URL-shaped arguments in ToolExecutionRequest before the underlying ToolExecutor runs. Thin wrapper over ssrf-guard-llm.") + url.set(providers.gradleProperty("POM_URL")) + inceptionYear.set(providers.gradleProperty("POM_INCEPTION_YEAR")) + licenses { + license { + name.set(providers.gradleProperty("POM_LICENSE_NAME")) + url.set(providers.gradleProperty("POM_LICENSE_URL")) + distribution.set(providers.gradleProperty("POM_LICENSE_DIST")) + } + } + developers { + developer { + id.set(providers.gradleProperty("POM_DEVELOPER_ID")) + name.set(providers.gradleProperty("POM_DEVELOPER_NAME")) + url.set(providers.gradleProperty("POM_DEVELOPER_URL")) + email.set(providers.gradleProperty("POM_DEVELOPER_EMAIL")) + organization.set(providers.gradleProperty("POM_ORGANIZATION_NAME")) + organizationUrl.set(providers.gradleProperty("POM_ORGANIZATION_URL")) + } + } + organization { + name.set(providers.gradleProperty("POM_ORGANIZATION_NAME")) + url.set(providers.gradleProperty("POM_ORGANIZATION_URL")) + } + scm { + url.set(providers.gradleProperty("POM_SCM_URL")) + connection.set(providers.gradleProperty("POM_SCM_CONNECTION")) + developerConnection.set(providers.gradleProperty("POM_SCM_DEV_CONNECTION")) + } + issueManagement { + system.set(providers.gradleProperty("POM_ISSUE_SYSTEM")) + url.set(providers.gradleProperty("POM_ISSUE_URL")) + } + } +} diff --git a/ssrf-guard-langchain4j/src/main/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardLangchain4jAutoConfiguration.java b/ssrf-guard-langchain4j/src/main/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardLangchain4jAutoConfiguration.java new file mode 100644 index 0000000..b8f2a1e --- /dev/null +++ b/ssrf-guard-langchain4j/src/main/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardLangchain4jAutoConfiguration.java @@ -0,0 +1,105 @@ +package kr.devslab.ssrfguard.langchain4j; + +import dev.langchain4j.service.tool.ToolExecutor; +import kr.devslab.ssrfguard.core.HostPolicy; +import kr.devslab.ssrfguard.core.NoOpSsrfGuardMetrics; +import kr.devslab.ssrfguard.core.SsrfGuardMetrics; +import kr.devslab.ssrfguard.core.SsrfGuardProperties; +import kr.devslab.ssrfguard.core.UrlPolicy; +import org.springframework.beans.factory.ObjectProvider; +import org.springframework.beans.factory.config.BeanPostProcessor; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +/** + * Auto-configuration that wires the SSRF defenses into LangChain4j tool + * execution. + * + *

Activated when: + *

+ * + *

The {@link BeanPostProcessor} below wraps every {@link ToolExecutor} + * bean Spring registers. A consumer's hand-rolled {@code @Bean ToolExecutor} + * (or a {@code @Tool}-annotated class registered as a bean) gets the URL + * policy applied without any wiring code. + * + *

Mirrors {@code SsrfGuardSpringAiAutoConfiguration} — same model, + * different framework. Both share the same {@link UrlPolicy} bean shape + * so a Spring AI app and a LangChain4j app configured against the same + * `ssrf.guard.*` properties block identical attacks. + */ +@AutoConfiguration +@ConditionalOnClass(ToolExecutor.class) +@EnableConfigurationProperties(SsrfGuardProperties.class) +@ConditionalOnProperty(prefix = "ssrf.guard", name = "enabled", havingValue = "true", matchIfMissing = true) +public class SsrfGuardLangchain4jAutoConfiguration { + + @Bean + @ConditionalOnMissingBean(SsrfGuardMetrics.class) + SsrfGuardMetrics ssrfGuardMetricsLangchain4j() { + return NoOpSsrfGuardMetrics.INSTANCE; + } + + @Bean + @ConditionalOnMissingBean + UrlPolicy ssrfUrlPolicyLangchain4j(SsrfGuardProperties props, + ObjectProvider hostPolicy, + SsrfGuardMetrics metrics) { + HostPolicy hp = hostPolicy.getIfAvailable( + () -> new HostPolicy(props.getExactHosts(), props.getSuffixes())); + return new UrlPolicy( + props.getAllowedSchemes(), + props.getAllowedPorts(), + hp, + props.isRejectIpLiteralHosts(), + props.isRejectUserInfo(), + metrics + ); + } + + /** + * Wraps every {@link ToolExecutor} bean with {@link SsrfGuardedToolExecutor}. + * Idempotent — already-wrapped executors pass through. + */ + @Bean + @ConditionalOnProperty(prefix = "ssrf.guard.langchain4j", name = "wrap-tool-executors", + havingValue = "true", matchIfMissing = true) + BeanPostProcessor ssrfGuardToolExecutorBeanPostProcessor(UrlPolicy policy) { + return new BeanPostProcessor() { + @Override + public Object postProcessAfterInitialization(Object bean, String beanName) { + if (bean instanceof ToolExecutor exec && !(bean instanceof SsrfGuardedToolExecutor)) { + return new SsrfGuardedToolExecutor(exec, policy); + } + return bean; + } + }; + } + + @Configuration(proxyBeanMethods = false) + @ConditionalOnClass(name = "io.micrometer.core.instrument.MeterRegistry") + static class MetricsConfiguration { + + @Bean + @ConditionalOnMissingBean(SsrfGuardMetrics.class) + SsrfGuardMetrics micrometerSsrfGuardMetricsLangchain4j( + ObjectProvider meterRegistry) { + var reg = meterRegistry.getIfAvailable(); + return reg == null + ? NoOpSsrfGuardMetrics.INSTANCE + : new kr.devslab.ssrfguard.core.MicrometerSsrfGuardMetrics(reg); + } + } +} diff --git a/ssrf-guard-langchain4j/src/main/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardedToolExecutor.java b/ssrf-guard-langchain4j/src/main/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardedToolExecutor.java new file mode 100644 index 0000000..a42215f --- /dev/null +++ b/ssrf-guard-langchain4j/src/main/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardedToolExecutor.java @@ -0,0 +1,97 @@ +package kr.devslab.ssrfguard.langchain4j; + +import dev.langchain4j.agent.tool.ToolExecutionRequest; +import dev.langchain4j.service.tool.ToolExecutor; +import kr.devslab.ssrfguard.core.SsrfGuardException; +import kr.devslab.ssrfguard.core.UrlPolicy; +import kr.devslab.ssrfguard.llm.JsonToolInputGuard; +import kr.devslab.ssrfguard.llm.ToolInputGuard; + +/** + * Wraps a LangChain4j {@link ToolExecutor} so URL-shaped arguments in the + * incoming {@link ToolExecutionRequest} are validated against an + * {@link UrlPolicy} before the underlying executor runs. + * + *

Sibling to {@code SsrfGuardedToolCallback} in {@code ssrf-guard-springai}. + * Same model, same core ({@link JsonToolInputGuard}) — only the framework + * abstraction differs. + * + *

The threat — recap

+ * LangChain4j {@code @Tool} methods and programmatic {@code ToolExecutor}s + * both end up receiving a JSON {@code arguments} string from the LLM. A + * tool like: + * 
{@code
+ * @Tool("Fetch a URL and return its body")
+ * String fetchUrl(String url) {
+ *     return restClient.get().uri(url).retrieve().body(String.class);
+ * }
+ * }
+ * is one prompt away from SSRF if {@code url} is attacker-controlled + * (directly via user message, or indirectly via RAG-injected instructions). + * + *

What the wrap does

+ *
    + *
  1. Reads {@code request.arguments()} — the JSON the LLM emitted.
    2. + *
  2. Delegates to {@link JsonToolInputGuard} — walks the whole tree, finds + * URL-shaped strings, validates each through {@link UrlPolicy}.
    3. + *
  3. On rejection, returns the structured JSON error directly (the LLM + * reads it on its next turn and recovers). On approval, delegates to + * the wrapped executor.
    4. + *
+ * + *

Usage

+ * Manual wrap (works without Spring): + * 
{@code
+ * UrlPolicy policy = ...;
+ * ToolExecutor raw = (request, memoryId) -> myFetchUrl(request.arguments());
+ * ToolExecutor safe = new SsrfGuardedToolExecutor(raw, policy);
+ *
+ * AiServices.builder(MyAssistant.class)
+ *     .chatModel(model)
+ *     .tools(Map.of(toolSpec, safe))
+ *     .build();
+ * }
+ * + * Auto-wrap (Spring): drop {@code ssrf-guard-langchain4j} on the classpath — + * {@link SsrfGuardLangchain4jAutoConfiguration} registers a + * {@code BeanPostProcessor} that wraps every {@link ToolExecutor} bean in + * the context. + */ +public final class SsrfGuardedToolExecutor implements ToolExecutor { + + private final ToolExecutor delegate; + private final ToolInputGuard guard; + + public SsrfGuardedToolExecutor(ToolExecutor delegate, UrlPolicy policy) { + this(delegate, policy, false); + } + + public SsrfGuardedToolExecutor(ToolExecutor delegate, UrlPolicy policy, boolean throwOnViolation) { + this.delegate = delegate; + this.guard = new JsonToolInputGuard(policy, throwOnViolation); + } + + /** + * Advanced constructor — supply a custom {@link ToolInputGuard}. Useful + * for non-JSON tool argument shapes, custom error payloads, or + * fixed-decision guards in tests. + */ + public SsrfGuardedToolExecutor(ToolExecutor delegate, ToolInputGuard guard) { + this.delegate = delegate; + this.guard = guard; + } + + @Override + public String execute(ToolExecutionRequest request, Object memoryId) { + String violation = guard.checkOrFormatError(request.arguments()); + return violation != null ? violation : delegate.execute(request, memoryId); + } + + public ToolExecutor delegate() { + return delegate; + } + + public ToolInputGuard guard() { + return guard; + } +} diff --git a/ssrf-guard-langchain4j/src/main/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardedToolExecutors.java b/ssrf-guard-langchain4j/src/main/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardedToolExecutors.java new file mode 100644 index 0000000..abc17bd --- /dev/null +++ b/ssrf-guard-langchain4j/src/main/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardedToolExecutors.java @@ -0,0 +1,61 @@ +package kr.devslab.ssrfguard.langchain4j; + +import dev.langchain4j.service.tool.ToolExecutor; +import kr.devslab.ssrfguard.core.UrlPolicy; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; + +/** + * Convenience helpers for bulk-wrapping LangChain4j tool executors. + * + * 
{@code
+ * Map raw = ...;
+ * Map safe = SsrfGuardedToolExecutors.wrap(raw, urlPolicy);
+ *
+ * AiServices.builder(MyAssistant.class)
+ *     .chatModel(model)
+ *     .tools(safe)
+ *     .build();
+ * }
+ * + *

The Spring auto-config wraps {@code ToolExecutor} beans via a + * {@code BeanPostProcessor}; these helpers are for the non-Spring case + * (plain LangChain4j) and the "I know exactly which executors I want + * wrapped" case. + */ +public final class SsrfGuardedToolExecutors { + private SsrfGuardedToolExecutors() {} + + /** Wrap a single executor. Idempotent — already-wrapped passes through. */ + public static ToolExecutor wrapOne(ToolExecutor executor, UrlPolicy policy) { + if (executor == null) return null; + if (executor instanceof SsrfGuardedToolExecutor) return executor; + return new SsrfGuardedToolExecutor(executor, policy); + } + + /** Wrap a collection. Preserves order. */ + public static List wrap(Collection executors, UrlPolicy policy) { + if (executors == null) return List.of(); + List out = new ArrayList<>(executors.size()); + for (ToolExecutor e : executors) out.add(wrapOne(e, policy)); + return out; + } + + /** + * Wrap a {@code Map} — the shape + * {@code AiServices.builder(...).tools(...)} expects when registering + * tools programmatically (instead of via {@code @Tool} annotations). + * Keys (specs) are returned unchanged; values are wrapped. + */ + public static Map wrap(Map specToExecutor, UrlPolicy policy) { + if (specToExecutor == null) return Map.of(); + java.util.LinkedHashMap out = new java.util.LinkedHashMap<>(specToExecutor.size()); + for (var entry : specToExecutor.entrySet()) { + out.put(entry.getKey(), wrapOne(entry.getValue(), policy)); + } + return out; + } +} diff --git a/ssrf-guard-langchain4j/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/ssrf-guard-langchain4j/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports new file mode 100644 index 0000000..94335a2 --- /dev/null +++ b/ssrf-guard-langchain4j/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -0,0 +1 @@ +kr.devslab.ssrfguard.langchain4j.SsrfGuardLangchain4jAutoConfiguration diff --git a/ssrf-guard-langchain4j/src/test/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardedToolExecutorTest.java b/ssrf-guard-langchain4j/src/test/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardedToolExecutorTest.java new file mode 100644 index 0000000..a269f8d --- /dev/null +++ b/ssrf-guard-langchain4j/src/test/java/kr/devslab/ssrfguard/langchain4j/SsrfGuardedToolExecutorTest.java @@ -0,0 +1,137 @@ +package kr.devslab.ssrfguard.langchain4j; + +import dev.langchain4j.agent.tool.ToolExecutionRequest; +import dev.langchain4j.service.tool.ToolExecutor; +import kr.devslab.ssrfguard.core.BlockReason; +import kr.devslab.ssrfguard.core.HostPolicy; +import kr.devslab.ssrfguard.core.NoOpSsrfGuardMetrics; +import kr.devslab.ssrfguard.core.SsrfGuardException; +import kr.devslab.ssrfguard.core.UrlPolicy; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatExceptionOfType; + +/** + * Tests the LangChain4j adapter against the same threat model as Spring AI: + * the LLM hands over a {@code ToolExecutionRequest} with a URL in the + * {@code arguments} JSON, the wrap inspects it through {@code UrlPolicy}, + * and the underlying executor never runs if the URL is rejected. + */ +class SsrfGuardedToolExecutorTest { + + private static UrlPolicy policy(List exact) { + return new UrlPolicy( + Set.of("http", "https"), + Set.of(-1, 80, 443), + new HostPolicy(exact, List.of()), + true, + true, + NoOpSsrfGuardMetrics.INSTANCE + ); + } + + private static ToolExecutionRequest request(String name, String arguments) { + return ToolExecutionRequest.builder() + .name(name) + .arguments(arguments) + .build(); + } + + /** A fake fetch_url tool that counts invocations so we can assert the wrap blocks. */ + static class CountingFetchUrlTool implements ToolExecutor { + final AtomicInteger invocations = new AtomicInteger(); + + @Override + public String execute(ToolExecutionRequest request, Object memoryId) { + invocations.incrementAndGet(); + return "PRETEND-FETCHED " + request.arguments(); + } + } + + @Test + void permits_whitelisted_url() { + CountingFetchUrlTool tool = new CountingFetchUrlTool(); + SsrfGuardedToolExecutor safe = new SsrfGuardedToolExecutor(tool, policy(List.of("api.example.com"))); + + String result = safe.execute(request("fetch_url", "{\"url\":\"https://api.example.com/v1\"}"), null); + + assertThat(result).startsWith("PRETEND-FETCHED"); + assertThat(tool.invocations).hasValue(1); + } + + @Test + void blocks_aws_metadata_url_returns_structured_error() { + CountingFetchUrlTool tool = new CountingFetchUrlTool(); + SsrfGuardedToolExecutor safe = new SsrfGuardedToolExecutor(tool, policy(List.of("api.example.com"))); + + String result = safe.execute( + request("fetch_url", "{\"url\":\"http://169.254.169.254/latest/meta-data/iam/security-credentials/\"}"), + null); + + assertThat(result) + .contains("\"error\":\"ssrf_blocked\"") + .contains("\"reason\":\"blocked_ip_literal\""); + assertThat(tool.invocations) + .as("underlying tool must NOT have been invoked") + .hasValue(0); + } + + @Test + void blocks_nested_url_field() { + CountingFetchUrlTool tool = new CountingFetchUrlTool(); + SsrfGuardedToolExecutor safe = new SsrfGuardedToolExecutor(tool, policy(List.of("api.example.com"))); + + String result = safe.execute( + request("fetch_url", "{\"request\":{\"target\":\"https://evil.com/\"}}"), + null); + + assertThat(result).contains("\"error\":\"ssrf_blocked\""); + assertThat(tool.invocations).hasValue(0); + } + + @Test + void blocks_obfuscated_ip_literal() { + CountingFetchUrlTool tool = new CountingFetchUrlTool(); + SsrfGuardedToolExecutor safe = new SsrfGuardedToolExecutor(tool, policy(List.of("api.example.com"))); + + String result = safe.execute(request("fetch_url", "{\"url\":\"http://2130706433/\"}"), null); + + assertThat(result).contains("\"reason\":\"blocked_ip_literal\""); + assertThat(tool.invocations).hasValue(0); + } + + @Test + void throws_when_throw_on_violation_is_true() { + CountingFetchUrlTool tool = new CountingFetchUrlTool(); + SsrfGuardedToolExecutor safe = new SsrfGuardedToolExecutor(tool, policy(List.of("api.example.com")), true); + + assertThatExceptionOfType(SsrfGuardException.class) + .isThrownBy(() -> safe.execute(request("fetch_url", "{\"url\":\"http://10.0.0.5/\"}"), null)); + assertThat(tool.invocations).hasValue(0); + } + + @Test + void passes_through_non_json_input() { + CountingFetchUrlTool tool = new CountingFetchUrlTool(); + SsrfGuardedToolExecutor safe = new SsrfGuardedToolExecutor(tool, policy(List.of("api.example.com"))); + + // Some LangChain4j tools take plain-text args; non-JSON should pass through. + String result = safe.execute(request("translate", "just a plain string"), null); + + assertThat(result).startsWith("PRETEND-FETCHED"); + assertThat(tool.invocations).hasValue(1); + } + + @Test + void wrapOne_is_idempotent() { + CountingFetchUrlTool tool = new CountingFetchUrlTool(); + ToolExecutor once = SsrfGuardedToolExecutors.wrapOne(tool, policy(List.of("api.example.com"))); + ToolExecutor twice = SsrfGuardedToolExecutors.wrapOne(once, policy(List.of("api.example.com"))); + assertThat(twice).isSameAs(once); + } +} diff --git a/ssrf-guard-llm/build.gradle.kts b/ssrf-guard-llm/build.gradle.kts new file mode 100644 index 0000000..3215d77 --- /dev/null +++ b/ssrf-guard-llm/build.gradle.kts @@ -0,0 +1,77 @@ +// ssrf-guard-llm — framework-agnostic core for LLM-agent tool input validation. +// +// Holds the JSON-walking / URL-extraction / policy-validation logic that +// every framework adapter (-springai, -langchain4j, future -mcp, …) shares. +// Adapters become ~30-line wrappers that delegate to ToolInputGuard and +// translate between this module's neutral API and the framework's own +// tool callback interface. +// +// Dependencies: +// - ssrf-guard-core for UrlPolicy / SsrfGuardException / BlockReason +// - Jackson (api) for JSON parsing — exposed so adapter modules don't +// each have to add Jackson themselves + +base { + archivesName.set("ssrf-guard-llm") +} + +dependencies { + api(project(":ssrf-guard-core")) + + // Jackson for parsing the JSON tool input. `api` so adapter modules pick + // it up transitively — keeps Spring AI / LangChain4j adapters from + // having to declare Jackson themselves. + api("com.fasterxml.jackson.core:jackson-databind") + + compileOnly("org.projectlombok:lombok") + annotationProcessor("org.projectlombok:lombok") + compileOnly("com.google.code.findbugs:jsr305:3.0.2") + + testImplementation("org.springframework.boot:spring-boot-starter-test") + testImplementation("org.assertj:assertj-core") + testRuntimeOnly("org.junit.platform:junit-platform-launcher") +} + +mavenPublishing { + coordinates( + providers.gradleProperty("GROUP").get(), + "ssrf-guard-llm", + providers.gradleProperty("VERSION").get() + ) + pom { + name.set("SSRF Guard — LLM Tool Input Validation") + description.set("Framework-agnostic core for validating URL-shaped arguments in LLM tool inputs (JSON). Used by ssrf-guard-springai, ssrf-guard-langchain4j, and any custom tool dispatcher.") + url.set(providers.gradleProperty("POM_URL")) + inceptionYear.set(providers.gradleProperty("POM_INCEPTION_YEAR")) + licenses { + license { + name.set(providers.gradleProperty("POM_LICENSE_NAME")) + url.set(providers.gradleProperty("POM_LICENSE_URL")) + distribution.set(providers.gradleProperty("POM_LICENSE_DIST")) + } + } + developers { + developer { + id.set(providers.gradleProperty("POM_DEVELOPER_ID")) + name.set(providers.gradleProperty("POM_DEVELOPER_NAME")) + url.set(providers.gradleProperty("POM_DEVELOPER_URL")) + email.set(providers.gradleProperty("POM_DEVELOPER_EMAIL")) + organization.set(providers.gradleProperty("POM_ORGANIZATION_NAME")) + organizationUrl.set(providers.gradleProperty("POM_ORGANIZATION_URL")) + } + } + organization { + name.set(providers.gradleProperty("POM_ORGANIZATION_NAME")) + url.set(providers.gradleProperty("POM_ORGANIZATION_URL")) + } + scm { + url.set(providers.gradleProperty("POM_SCM_URL")) + connection.set(providers.gradleProperty("POM_SCM_CONNECTION")) + developerConnection.set(providers.gradleProperty("POM_SCM_DEV_CONNECTION")) + } + issueManagement { + system.set(providers.gradleProperty("POM_ISSUE_SYSTEM")) + url.set(providers.gradleProperty("POM_ISSUE_URL")) + } + } +} diff --git a/ssrf-guard-llm/src/main/java/kr/devslab/ssrfguard/llm/JsonToolInputGuard.java b/ssrf-guard-llm/src/main/java/kr/devslab/ssrfguard/llm/JsonToolInputGuard.java new file mode 100644 index 0000000..72f2fd5 --- /dev/null +++ b/ssrf-guard-llm/src/main/java/kr/devslab/ssrfguard/llm/JsonToolInputGuard.java @@ -0,0 +1,177 @@ +package kr.devslab.ssrfguard.llm; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import kr.devslab.ssrfguard.core.SsrfGuardException; +import kr.devslab.ssrfguard.core.UrlPolicy; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.URI; +import java.net.URISyntaxException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +/** + * Default {@link ToolInputGuard}. Treats the tool input as a JSON object, + * walks the whole tree, finds every {@code http(s)://...} string at any + * depth, and validates each through the supplied {@link UrlPolicy}. + * + *

Why walk the whole tree

+ * Naïve guards only inspect a top-level {@code "url"} field. Real-world + * LLM tool schemas nest URLs in surprising places: + * + * + * + * Walking the whole tree means a single bad URL anywhere in the structure + * trips the guard, not just one at a fixed location. + * + *

Failure mode — error string vs. exception

+ * Default is to return a structured JSON error string — the LLM sees it on + * its next turn and can recover gracefully ("I can't fetch that URL"). + * Set {@code throwOnViolation = true} for CI / test contexts that want a + * thrown {@link SsrfGuardException} instead. + */ +public final class JsonToolInputGuard implements ToolInputGuard { + + private static final Logger log = LoggerFactory.getLogger(JsonToolInputGuard.class); + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private final UrlPolicy policy; + private final boolean throwOnViolation; + + public JsonToolInputGuard(UrlPolicy policy) { + this(policy, false); + } + + public JsonToolInputGuard(UrlPolicy policy, boolean throwOnViolation) { + this.policy = policy; + this.throwOnViolation = throwOnViolation; + } + + /** + * {@inheritDoc} + * + *

Concrete error-payload shape on rejection: + * + * 

{@code
+     * {
+     *   "error": "ssrf_blocked",
+     *   "reason": "blocked_private_ip",
+     *   "url": "http://169.254.169.254/...",
+     *   "message": "DNS resolved to a private/loopback address: /169.254.169.254",
+     *   "guidance": "Refuse the request or ask the user for a different URL. ..."
+     * }
+     * }
+ * + *

The {@code reason} field is the stable + * {@link kr.devslab.ssrfguard.core.BlockReason#label() BlockReason label} + * — safe to surface in metrics tags and grep against in log search. + */ + @Override + public String checkOrFormatError(String toolInput) { + if (toolInput == null || toolInput.isBlank()) return null; + + JsonNode root; + try { + root = MAPPER.readTree(toolInput); + } catch (Exception e) { + // Not JSON — pass through. Many tools accept plain text or other + // shapes; the HTTP-client-level guard catches anything that does + // eventually become a URL. + log.debug("ssrf-guard: tool input not parseable as JSON, skipping URL scan"); + return null; + } + + List urls = collectUrlLikeStrings(root); + for (String url : urls) { + URI uri; + try { + uri = new URI(url); + } catch (URISyntaxException ignored) { + continue; + } + String scheme = uri.getScheme(); + if (scheme == null) continue; + // Only http/https. file://, gopher://, etc. would be rejected by + // the URL policy anyway, but we don't want to false-positive on + // strings like "mailto:..." or "urn:uuid:...". + if (!scheme.equalsIgnoreCase("http") && !scheme.equalsIgnoreCase("https")) continue; + + try { + policy.validate(uri); + } catch (SsrfGuardException e) { + if (throwOnViolation) throw e; + return formatErrorPayload(e, url); + } + } + return null; + } + + private static List collectUrlLikeStrings(JsonNode node) { + List out = new ArrayList<>(); + collectUrlLikeStrings(node, out); + return out; + } + + private static void collectUrlLikeStrings(JsonNode node, List out) { + if (node == null) return; + if (node.isTextual()) { + String v = node.asText(); + if (looksLikeUrl(v)) out.add(v); + return; + } + if (node.isArray()) { + for (JsonNode child : node) collectUrlLikeStrings(child, out); + return; + } + if (node.isObject()) { + // Jackson 2.18+ deprecated fields(); .properties() returns a + // Set>. Iterate the values directly. + for (JsonNode child : node.properties().stream() + .map(Map.Entry::getValue).toList()) { + collectUrlLikeStrings(child, out); + } + } + } + + private static boolean looksLikeUrl(String s) { + if (s == null) return false; + String trimmed = s.trim(); + return trimmed.startsWith("http://") || trimmed.startsWith("https://"); + } + + private String formatErrorPayload(SsrfGuardException e, String url) { + try { + return MAPPER.writeValueAsString(Map.of( + "error", "ssrf_blocked", + "reason", e.reason().label(), + "url", url, + "message", e.getMessage(), + "guidance", "Refuse the request or ask the user for a different URL. " + + "The blocked URL targets a private/internal network or violates the application's SSRF policy." + )); + } catch (Exception jsonErr) { + // Fallback — minimal hand-rolled JSON. We control all the field + // values here so this stays safe to concatenate. + return "{\"error\":\"ssrf_blocked\",\"reason\":\"" + e.reason().label() + "\"}"; + } + } + + public UrlPolicy policy() { + return policy; + } + + public boolean throwOnViolation() { + return throwOnViolation; + } +} diff --git a/ssrf-guard-llm/src/main/java/kr/devslab/ssrfguard/llm/ToolInputGuard.java b/ssrf-guard-llm/src/main/java/kr/devslab/ssrfguard/llm/ToolInputGuard.java new file mode 100644 index 0000000..4a9fa2b --- /dev/null +++ b/ssrf-guard-llm/src/main/java/kr/devslab/ssrfguard/llm/ToolInputGuard.java @@ -0,0 +1,57 @@ +package kr.devslab.ssrfguard.llm; + +import kr.devslab.ssrfguard.core.SsrfGuardException; + +/** + * Validates URL-shaped arguments inside LLM tool inputs. Framework-agnostic — + * Spring AI, LangChain4j, custom dispatchers, MCP servers, and anyone else + * who hands a JSON blob to an underlying tool can route that blob through a + * single {@link ToolInputGuard} implementation before letting the tool run. + * + *

The contract on the return value is deliberately simple — the caller + * either gets {@code null} (everything OK, proceed) or a JSON-shaped error + * string ready to be returned to the LLM verbatim: + * + * 

{@code
+ * {"error":"ssrf_blocked","reason":"blocked_private_ip","url":"http://169.254.169.254/...","message":"...","guidance":"..."}
+ * }
+ * + *

That payload is what well-behaved chat models read on their next turn + * and use to apologise / ask for a different URL rather than crashing the + * agent loop with an unhandled exception. Implementations that prefer + * fail-loud semantics (CI tests, dev-mode harnesses) can throw + * {@link SsrfGuardException} instead — see {@link JsonToolInputGuard}'s + * {@code throwOnViolation} flag. + * + *

Why an interface rather than a single concrete type

+ * Most adapters take the default {@link JsonToolInputGuard} and depend on + * the interface. Two reasons we still expose the abstraction: + * + * + */ +public interface ToolInputGuard { + + /** + * Walk the tool input, find every URL-shaped argument, and validate + * each through the configured policy. + * + * @param toolInput the raw input the framework would hand to the tool — + * typically a JSON object string. May be null or blank. + * @return {@code null} if every URL passed; a JSON error payload string + * if any URL was blocked. Concrete shape documented on + * {@link JsonToolInputGuard#checkOrFormatError(String)}. + * @throws SsrfGuardException if the implementation is configured to + * throw on violations (see + * {@link JsonToolInputGuard}'s + * {@code throwOnViolation} flag) and a URL + * failed the policy. + */ + String checkOrFormatError(String toolInput); +} diff --git a/ssrf-guard-llm/src/test/java/kr/devslab/ssrfguard/llm/JsonToolInputGuardTest.java b/ssrf-guard-llm/src/test/java/kr/devslab/ssrfguard/llm/JsonToolInputGuardTest.java new file mode 100644 index 0000000..856a80e --- /dev/null +++ b/ssrf-guard-llm/src/test/java/kr/devslab/ssrfguard/llm/JsonToolInputGuardTest.java @@ -0,0 +1,104 @@ +package kr.devslab.ssrfguard.llm; + +import kr.devslab.ssrfguard.core.BlockReason; +import kr.devslab.ssrfguard.core.HostPolicy; +import kr.devslab.ssrfguard.core.NoOpSsrfGuardMetrics; +import kr.devslab.ssrfguard.core.SsrfGuardException; +import kr.devslab.ssrfguard.core.UrlPolicy; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Set; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatExceptionOfType; + +/** + * Tests the framework-agnostic core of LLM tool input validation. These + * tests intentionally do not touch any LLM framework — that's the whole + * point of extracting this module. + */ +class JsonToolInputGuardTest { + + private static UrlPolicy policy(List exact) { + return new UrlPolicy( + Set.of("http", "https"), + Set.of(-1, 80, 443), + new HostPolicy(exact, List.of()), + true, + true, + NoOpSsrfGuardMetrics.INSTANCE + ); + } + + @Test + void returns_null_when_input_is_blank_or_not_json() { + var guard = new JsonToolInputGuard(policy(List.of("api.example.com"))); + assertThat(guard.checkOrFormatError(null)).isNull(); + assertThat(guard.checkOrFormatError("")).isNull(); + assertThat(guard.checkOrFormatError(" ")).isNull(); + assertThat(guard.checkOrFormatError("not json at all")).isNull(); + } + + @Test + void allows_input_without_urls() { + var guard = new JsonToolInputGuard(policy(List.of("api.example.com"))); + assertThat(guard.checkOrFormatError("{\"query\":\"weather today\",\"limit\":5}")).isNull(); + } + + @Test + void allows_whitelisted_url_in_top_level_field() { + var guard = new JsonToolInputGuard(policy(List.of("api.example.com"))); + assertThat(guard.checkOrFormatError("{\"url\":\"https://api.example.com/v1\"}")).isNull(); + } + + @Test + void blocks_aws_metadata_url_returns_structured_error() { + var guard = new JsonToolInputGuard(policy(List.of("api.example.com"))); + String err = guard.checkOrFormatError("{\"url\":\"http://169.254.169.254/latest/meta-data/\"}"); + assertThat(err) + .isNotNull() + .contains("\"error\":\"ssrf_blocked\"") + .contains("\"reason\":\"blocked_ip_literal\"") + .contains("\"url\":\"http://169.254.169.254/latest/meta-data/\"") + .contains("\"guidance\":"); + } + + @Test + void blocks_nested_url_field() { + var guard = new JsonToolInputGuard(policy(List.of("api.example.com"))); + String err = guard.checkOrFormatError("{\"request\":{\"target\":\"https://evil.com/\"},\"timeout\":5}"); + assertThat(err).contains("\"error\":\"ssrf_blocked\""); + } + + @Test + void blocks_url_inside_array() { + var guard = new JsonToolInputGuard(policy(List.of("api.example.com"))); + String err = guard.checkOrFormatError("{\"urls\":[\"https://api.example.com/ok\", \"https://evil.com/bad\"]}"); + assertThat(err).contains("\"error\":\"ssrf_blocked\""); + } + + @Test + void blocks_obfuscated_ip_literal() { + var guard = new JsonToolInputGuard(policy(List.of("api.example.com"))); + String err = guard.checkOrFormatError("{\"url\":\"http://2130706433/\"}"); + assertThat(err).contains("\"reason\":\"blocked_ip_literal\""); + } + + @Test + void throws_when_throw_on_violation_is_true() { + var guard = new JsonToolInputGuard(policy(List.of("api.example.com")), true); + assertThatExceptionOfType(SsrfGuardException.class) + .isThrownBy(() -> guard.checkOrFormatError("{\"url\":\"http://10.0.0.5/\"}")) + .matches(e -> e.reason() == BlockReason.BLOCKED_IP_LITERAL); + } + + @Test + void ignores_non_http_schemes() { + // mailto:, urn:uuid:, file:// — should not trip the URL detector + // because looksLikeUrl() only matches http(s):// + var guard = new JsonToolInputGuard(policy(List.of("api.example.com"))); + assertThat(guard.checkOrFormatError("{\"to\":\"mailto:user@example.com\"}")).isNull(); + assertThat(guard.checkOrFormatError("{\"id\":\"urn:uuid:abc\"}")).isNull(); + } +} diff --git a/ssrf-guard-springai/build.gradle.kts b/ssrf-guard-springai/build.gradle.kts index b191624..b735da3 100644 --- a/ssrf-guard-springai/build.gradle.kts +++ b/ssrf-guard-springai/build.gradle.kts @@ -14,19 +14,19 @@ base { } dependencies { - api(project(":ssrf-guard-core")) + // ssrf-guard-llm carries the framework-agnostic JSON walking + URL + // validation logic (extracted in v3.1). Jackson follows transitively + // through -llm's `api` dep — this module no longer declares Jackson + // directly. Behaviour for v3.0.x consumers is unchanged because the + // public Spring AI types (`SsrfGuardedToolCallback`, the auto-config, + // `SsrfGuardedToolCallbacks`) keep the same signatures. + api(project(":ssrf-guard-llm")) api("org.springframework.boot:spring-boot-autoconfigure") - // Spring AI Tool API — 1.0 stable. The starter pulls in spring-ai-core - // transitively, where @Tool / ToolCallback live. - // Spring AI 1.0 GA artifacts. The original spring-ai-core was renamed - // — model abstractions (ToolCallback, ToolDefinition, @Tool) now live - // in spring-ai-model, and ChatClient/ChatModel are in spring-ai-client-chat. + // Spring AI Tool API — 1.0 stable. spring-ai-core was renamed to + // spring-ai-model where @Tool / ToolCallback / ToolDefinition live. compileOnly("org.springframework.ai:spring-ai-model:1.0.7") - // Jackson for parsing the JSON tool input. - compileOnly("com.fasterxml.jackson.core:jackson-databind") - compileOnly("io.micrometer:micrometer-core") compileOnly("org.projectlombok:lombok") annotationProcessor("org.projectlombok:lombok") diff --git a/ssrf-guard-springai/src/main/java/kr/devslab/ssrfguard/springai/SsrfGuardedToolCallback.java b/ssrf-guard-springai/src/main/java/kr/devslab/ssrfguard/springai/SsrfGuardedToolCallback.java index 450ad10..2e25b94 100644 --- a/ssrf-guard-springai/src/main/java/kr/devslab/ssrfguard/springai/SsrfGuardedToolCallback.java +++ b/ssrf-guard-springai/src/main/java/kr/devslab/ssrfguard/springai/SsrfGuardedToolCallback.java @@ -1,71 +1,75 @@ package kr.devslab.ssrfguard.springai; -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; -import kr.devslab.ssrfguard.core.BlockReason; import kr.devslab.ssrfguard.core.SsrfGuardException; import kr.devslab.ssrfguard.core.UrlPolicy; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; +import kr.devslab.ssrfguard.llm.JsonToolInputGuard; +import kr.devslab.ssrfguard.llm.ToolInputGuard; import org.springframework.ai.tool.ToolCallback; import org.springframework.ai.tool.definition.ToolDefinition; import org.springframework.ai.tool.metadata.ToolMetadata; import org.springframework.ai.chat.model.ToolContext; -import java.net.URI; -import java.net.URISyntaxException; -import java.util.ArrayList; -import java.util.List; -import java.util.Map; - /** * Wraps a Spring AI {@link ToolCallback} so URL-shaped tool inputs are * validated against an {@link UrlPolicy} before the underlying tool executes. * - *

The threat this closes: an LLM agent's "fetch_url" / "scrape_page" / - * "summarize_link" tool. The model accepts a URL from the user (sometimes - * directly, sometimes indirectly through RAG search results), passes it as - * a tool argument, and your code does {@code restClient.get().uri(url)…}. - * That's a one-line SSRF vector if the URL isn't validated — and the SSRF - * Guard interceptors at the HTTP-client layer DO catch it. But the wrap - * here gives the LLM a structured error at tool-call time, before - * any HTTP traffic is generated, so the model can recover gracefully - * ("I'm not allowed to fetch that URL, let me ask the user for a - * different one") instead of seeing an HTTP exception it can't reason - * about. + *

v3.1 note. The JSON walking + URL extraction + policy validation + * lives in {@link JsonToolInputGuard} (in the {@code ssrf-guard-llm} module). + * This class is a 30-line adapter that translates between Spring AI's + * {@code ToolCallback} interface and the framework-agnostic + * {@link ToolInputGuard} contract. Same applies to + * {@code SsrfGuardedTool} in {@code ssrf-guard-langchain4j} — the two + * adapters share the same core, so a fix to URL detection or error + * formatting lands in one place. * - *

Heuristic URL detection

- * The tool input arrives as a JSON string ({@code {"url":"http://...","timeout":5}}). - * We walk the JSON tree and, for every string-valued leaf, try to parse it - * as a URI. Anything that parses AND has an http/https scheme runs through - * the policy. False positives (a string field that happens to start with - * {@code "http://..."} but isn't meant as a URL) are rare and would be - * blocked anyway if they'd be used as URLs downstream. + *

The public API is unchanged from v3.0.x — same constructors, same + * behaviour. Existing consumers don't see the refactor. * - *

Failure mode — error string vs. exception

- * We return a structured JSON error string rather than throwing. - * Spring AI surfaces tool exceptions as failed tool calls and the LLM has - * to interpret them; returning an error string keeps the conversation - * flowing and is what most production agents want. Consumers who'd rather - * fail loud can pass {@code throwOnViolation = true} to the constructor. + *

What this still does (delegated to the core)

+ * */ public final class SsrfGuardedToolCallback implements ToolCallback { - private static final Logger log = LoggerFactory.getLogger(SsrfGuardedToolCallback.class); - private static final ObjectMapper MAPPER = new ObjectMapper(); - private final ToolCallback delegate; - private final UrlPolicy policy; - private final boolean throwOnViolation; + private final ToolInputGuard guard; + /** + * @param delegate the raw {@link ToolCallback} this wraps + * @param policy the URL policy applied to every URL in the tool input + */ public SsrfGuardedToolCallback(ToolCallback delegate, UrlPolicy policy) { this(delegate, policy, false); } + /** + * @param throwOnViolation when {@code true}, the wrap rethrows + * {@link SsrfGuardException} on a violation instead + * of returning a JSON error string. Useful for + * CI and dev-mode harnesses. + */ public SsrfGuardedToolCallback(ToolCallback delegate, UrlPolicy policy, boolean throwOnViolation) { this.delegate = delegate; - this.policy = policy; - this.throwOnViolation = throwOnViolation; + this.guard = new JsonToolInputGuard(policy, throwOnViolation); + } + + /** + * Advanced constructor — supply a custom {@link ToolInputGuard}. Useful + * if you want non-JSON tool input handling, a different error payload + * shape, or to plug in a fixed-decision guard for tests. + */ + public SsrfGuardedToolCallback(ToolCallback delegate, ToolInputGuard guard) { + this.delegate = delegate; + this.guard = guard; } @Override @@ -80,122 +84,21 @@ public ToolMetadata getToolMetadata() { @Override public String call(String toolInput) { - String violation = checkUrlsInJson(toolInput); - if (violation != null) { - return violation; - } - return delegate.call(toolInput); + String violation = guard.checkOrFormatError(toolInput); + return violation != null ? violation : delegate.call(toolInput); } @Override public String call(String toolInput, ToolContext toolContext) { - String violation = checkUrlsInJson(toolInput); - if (violation != null) { - return violation; - } - return delegate.call(toolInput, toolContext); - } - - /** - * Walks {@code toolInput} as JSON, finds every URL-looking string, and - * runs each through {@link UrlPolicy}. Returns {@code null} if every URL - * passes; otherwise returns a JSON error string the LLM can read and - * react to. If {@code throwOnViolation = true}, re-throws the - * SsrfGuardException instead. - */ - private String checkUrlsInJson(String toolInput) { - if (toolInput == null || toolInput.isBlank()) return null; - JsonNode root; - try { - root = MAPPER.readTree(toolInput); - } catch (Exception e) { - // Not JSON — pass through. Most tools that don't take URL args - // accept plain-text or other shapes; the HTTP-client-level - // guard catches anything that does eventually become a URL. - log.debug("ssrf-guard: tool input not parseable as JSON, skipping URL scan"); - return null; - } - - List urls = collectUrlLikeStrings(root); - for (String url : urls) { - URI uri; - try { - uri = new URI(url); - } catch (URISyntaxException ignored) { - continue; - } - String scheme = uri.getScheme(); - if (scheme == null) continue; - // Only http/https — file://, gopher://, etc. would be rejected by - // policy.validate anyway, but we don't want to false-positive on - // strings like "mailto:..." or "urn:uuid:...". - if (!scheme.equalsIgnoreCase("http") && !scheme.equalsIgnoreCase("https")) continue; - - try { - policy.validate(uri); - } catch (SsrfGuardException e) { - if (throwOnViolation) throw e; - return formatErrorPayload(e, url); - } - } - return null; - } - - private static List collectUrlLikeStrings(JsonNode node) { - List out = new ArrayList<>(); - collectUrlLikeStrings(node, out); - return out; - } - - private static void collectUrlLikeStrings(JsonNode node, List out) { - if (node == null) return; - if (node.isTextual()) { - String v = node.asText(); - if (looksLikeUrl(v)) out.add(v); - return; - } - if (node.isArray()) { - for (JsonNode child : node) collectUrlLikeStrings(child, out); - return; - } - if (node.isObject()) { - // Jackson 2.18+ deprecated fields(); .properties() returns a - // Set>. Iterate the values directly. - for (JsonNode child : node.properties().stream() - .map(Map.Entry::getValue).toList()) { - collectUrlLikeStrings(child, out); - } - } - } - - private static boolean looksLikeUrl(String s) { - if (s == null) return false; - String trimmed = s.trim(); - return trimmed.startsWith("http://") || trimmed.startsWith("https://"); - } - - private String formatErrorPayload(SsrfGuardException e, String url) { - // Structured JSON the LLM can parse if it wants to; readable as plain - // text either way. Keys mirror the metric tags so log-search "find - // every blocked SSRF for tool=…" works. - try { - return MAPPER.writeValueAsString(Map.of( - "error", "ssrf_blocked", - "reason", e.reason().label(), - "url", url, - "message", e.getMessage(), - "guidance", "Refuse the request or ask the user for a different URL. The blocked URL targets a private/internal network or violates the application's SSRF policy." - )); - } catch (Exception jsonErr) { - return "{\"error\":\"ssrf_blocked\",\"reason\":\"" + e.reason().label() + "\"}"; - } + String violation = guard.checkOrFormatError(toolInput); + return violation != null ? violation : delegate.call(toolInput, toolContext); } public ToolCallback delegate() { return delegate; } - public UrlPolicy policy() { - return policy; + public ToolInputGuard guard() { + return guard; } }