diff --git a/api/Dockerfile b/api/Dockerfile index 311bc51df1578c..4be449d454830f 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -99,10 +99,10 @@ ENV VIRTUAL_ENV=/app/api/.venv COPY --from=packages --chown=dify:dify ${VIRTUAL_ENV} ${VIRTUAL_ENV} ENV PATH="${VIRTUAL_ENV}/bin:${PATH}" -RUN mkdir -p /usr/local/share/nltk_data \ - && NLTK_DATA=/usr/local/share/nltk_data python -m nltk.downloader punkt_tab averaged_perceptron_tagger_eng stopwords \ - && NLTK_DATA=/usr/local/share/nltk_data python -c "import nltk; nltk.data.find('tokenizers/punkt_tab'); nltk.data.find('taggers/averaged_perceptron_tagger_eng'); nltk.data.find('corpora/stopwords')" \ - && chmod -R 755 /usr/local/share/nltk_data +# RUN mkdir -p /usr/local/share/nltk_data \ +# && NLTK_DATA=/usr/local/share/nltk_data python -m nltk.downloader punkt_tab averaged_perceptron_tagger_eng stopwords \ +# && NLTK_DATA=/usr/local/share/nltk_data python -c "import nltk; nltk.data.find('tokenizers/punkt_tab'); nltk.data.find('taggers/averaged_perceptron_tagger_eng'); nltk.data.find('corpora/stopwords')" \ +# && chmod -R 755 /usr/local/share/nltk_data ENV TIKTOKEN_CACHE_DIR=/app/api/.tiktoken_cache diff --git a/dify-agent-runtime/Makefile b/dify-agent-runtime/Makefile index fffa071cc3ba32..2d277a564921d1 100644 --- a/dify-agent-runtime/Makefile +++ b/dify-agent-runtime/Makefile @@ -112,6 +112,94 @@ integration-up: fi; \ sleep 2; \ done + @echo "Setting up egress proxy container + echo backend..." + $(eval EGRESS_NET := sandbox-rt-egress-net-$(TEST_ID)) + docker network create $(EGRESS_NET) > /dev/null + $(eval ECHO_CONTAINER := echo-backend-$(TEST_ID)) + docker run -d --name $(ECHO_CONTAINER) \ + --network $(EGRESS_NET) --network-alias echo-backend \ + mendhak/http-https-echo:31 + $(eval EGRESS_CONTAINER_NAME := sandbox-rt-egress-$(TEST_ID)) + $(eval EGRESS_HOST_PORT := $(shell python3 -c 'import socket; s=socket.socket(); s.bind(("",0)); print(s.getsockname()[1]); s.close()')) + $(eval EGRESS_AUTH_TOKEN := test-token-egress-$(TEST_ID)) + docker run -d --name $(EGRESS_CONTAINER_NAME) \ + --network $(EGRESS_NET) \ + -p $(EGRESS_HOST_PORT):5004 \ + -e SHELLCTL_AUTH_TOKEN=$(EGRESS_AUTH_TOKEN) \ + -e SHELLCTL_EGRESSPROXY_ENABLED=true \ + -e HTTP_PROXY=http://127.0.0.1:18080 \ + -e HTTPS_PROXY=http://127.0.0.1:18080 \ + -e NO_PROXY=localhost,127.0.0.1 \ + $(IMAGE_NAME) + @echo 'EGRESS_NET=$(EGRESS_NET)' >> $(STATE_FILE) + @echo 'ECHO_CONTAINER=$(ECHO_CONTAINER)' >> $(STATE_FILE) + @echo 'EGRESS_CONTAINER_NAME=$(EGRESS_CONTAINER_NAME)' >> $(STATE_FILE) + @echo 'EGRESS_HOST_PORT=$(EGRESS_HOST_PORT)' >> $(STATE_FILE) + @echo 'EGRESS_AUTH_TOKEN=$(EGRESS_AUTH_TOKEN)' >> $(STATE_FILE) + @for i in $$(seq 1 30); do \ + if curl -sf http://localhost:$(EGRESS_HOST_PORT)/healthz > /dev/null 2>&1; then \ + echo "Egress proxy runtime is ready on port $(EGRESS_HOST_PORT)"; \ + break; \ + fi; \ + if [ "$$i" -eq 30 ]; then \ + echo "ERROR: egress proxy runtime not ready after 60s" >&2; \ + docker logs $(EGRESS_CONTAINER_NAME); \ + docker rm -f $(EGRESS_CONTAINER_NAME) $(ECHO_CONTAINER) 2>/dev/null; \ + docker network rm $(EGRESS_NET) 2>/dev/null; \ + exit 1; \ + fi; \ + sleep 2; \ + done + @echo "Setting up upstream squid + egress-proxy-with-upstream container..." + @echo " NOTE: an earlier iteration tried to place this runtime container on" + @echo " a network with no direct path to echo-backend (only reachable via" + @echo " squid-upstream bridging two networks), to force a true end-to-end" + @echo " regression test of hostname passthrough mirroring production's" + @echo " local_sandbox/agent_ssrf_proxy/api topology. That approach was" + @echo " reverted: multi-homing squid-upstream across two docker networks" + @echo " via 'docker network connect' hit a Docker/host networking pitfall" + @echo " (squid answered fine on its primary network but not the secondary" + @echo " one -- asymmetric routing / rp_filter in the Docker Desktop VM)," + @echo " which is an environment limitation unrelated to the proxy code." + @echo " Hostname-passthrough itself is covered reliably (no Docker network" + @echo " involved) by TestProxyUpstreamChainingPreservesHostname in" + @echo " internal/egressproxy/proxy_test.go. This test only verifies that" + @echo " upstream chaining + credential injection work end-to-end." + $(eval SQUID_CONTAINER := squid-upstream-$(TEST_ID)) + docker run -d --name $(SQUID_CONTAINER) \ + --network $(EGRESS_NET) --network-alias squid-upstream \ + -v $(CURDIR)/tests/squid-test.conf:/etc/squid/squid.conf:ro \ + ubuntu/squid:latest + $(eval EGRESS_UPSTREAM_CONTAINER_NAME := sandbox-rt-egress-upstream-$(TEST_ID)) + $(eval EGRESS_UPSTREAM_HOST_PORT := $(shell python3 -c 'import socket; s=socket.socket(); s.bind(("",0)); print(s.getsockname()[1]); s.close()')) + $(eval EGRESS_UPSTREAM_AUTH_TOKEN := test-token-egress-upstream-$(TEST_ID)) + docker run -d --name $(EGRESS_UPSTREAM_CONTAINER_NAME) \ + --network $(EGRESS_NET) \ + -p $(EGRESS_UPSTREAM_HOST_PORT):5004 \ + -e SHELLCTL_AUTH_TOKEN=$(EGRESS_UPSTREAM_AUTH_TOKEN) \ + -e SHELLCTL_EGRESSPROXY_ENABLED=true \ + -e SHELLCTL_EGRESSPROXY_UPSTREAM=http://squid-upstream:3128 \ + -e HTTP_PROXY=http://127.0.0.1:18080 \ + -e HTTPS_PROXY=http://127.0.0.1:18080 \ + -e NO_PROXY=localhost,127.0.0.1 \ + $(IMAGE_NAME) + @echo 'SQUID_CONTAINER=$(SQUID_CONTAINER)' >> $(STATE_FILE) + @echo 'EGRESS_UPSTREAM_CONTAINER_NAME=$(EGRESS_UPSTREAM_CONTAINER_NAME)' >> $(STATE_FILE) + @echo 'EGRESS_UPSTREAM_HOST_PORT=$(EGRESS_UPSTREAM_HOST_PORT)' >> $(STATE_FILE) + @echo 'EGRESS_UPSTREAM_AUTH_TOKEN=$(EGRESS_UPSTREAM_AUTH_TOKEN)' >> $(STATE_FILE) + @for i in $$(seq 1 30); do \ + if curl -sf http://localhost:$(EGRESS_UPSTREAM_HOST_PORT)/healthz > /dev/null 2>&1; then \ + echo "Egress-with-upstream runtime is ready on port $(EGRESS_UPSTREAM_HOST_PORT)"; \ + break; \ + fi; \ + if [ "$$i" -eq 30 ]; then \ + echo "ERROR: egress-with-upstream runtime not ready after 60s" >&2; \ + docker logs $(EGRESS_UPSTREAM_CONTAINER_NAME); \ + docker rm -f $(EGRESS_UPSTREAM_CONTAINER_NAME) $(SQUID_CONTAINER) 2>/dev/null; \ + exit 1; \ + fi; \ + sleep 2; \ + done integration-test: @test -f $(STATE_FILE) || { echo "ERROR: run 'make integration-up' first" >&2; exit 1; } @@ -120,6 +208,10 @@ integration-test: SHELLCTL_TEST_TOKEN=$$AUTH_TOKEN \ SHELLCTL_GO_URL_NO_ISOLATION=http://localhost:$$HOST_PORT_NOISO \ SHELLCTL_TEST_TOKEN_NO_ISOLATION=$$AUTH_TOKEN_NOISO \ + SHELLCTL_EGRESS_GO_URL=http://localhost:$$EGRESS_HOST_PORT \ + SHELLCTL_EGRESS_TEST_TOKEN=$$EGRESS_AUTH_TOKEN \ + SHELLCTL_EGRESS_UPSTREAM_GO_URL=http://localhost:$$EGRESS_UPSTREAM_HOST_PORT \ + SHELLCTL_EGRESS_UPSTREAM_TEST_TOKEN=$$EGRESS_UPSTREAM_AUTH_TOKEN \ go test -tags=integration -v -count=1 -timeout=300s ./tests/... integration-logs: @@ -131,6 +223,11 @@ integration-down: . ./$(STATE_FILE); \ docker rm -f $$CONTAINER_NAME 2>/dev/null || true; \ docker rm -f $$CONTAINER_NAME_NOISO 2>/dev/null || true; \ + docker rm -f $$EGRESS_CONTAINER_NAME 2>/dev/null || true; \ + docker rm -f $$EGRESS_UPSTREAM_CONTAINER_NAME 2>/dev/null || true; \ + docker rm -f $$SQUID_CONTAINER 2>/dev/null || true; \ + docker rm -f $$ECHO_CONTAINER 2>/dev/null || true; \ + docker network rm $$EGRESS_NET 2>/dev/null || true; \ rm -f $(STATE_FILE); \ fi diff --git a/dify-agent-runtime/README.md b/dify-agent-runtime/README.md index 59b1d45116736e..e947d8e78d8803 100644 --- a/dify-agent-runtime/README.md +++ b/dify-agent-runtime/README.md @@ -36,6 +36,7 @@ docker build -f dify-agent-runtime/docker/Dockerfile \ ``` docker run -d --name dify-agent-runtime \ -p 15004:5004 \ + -e SHELLCTL_EGRESSPROXY_ENABLED=true \ dify-agent-runtime:latest ``` diff --git a/dify-agent-runtime/cmd/runner/main.go b/dify-agent-runtime/cmd/runner/main.go index 29f35b4867d936..fd63aa359141be 100644 --- a/dify-agent-runtime/cmd/runner/main.go +++ b/dify-agent-runtime/cmd/runner/main.go @@ -67,12 +67,12 @@ func parentMode() { env := os.Environ() // Remove internal shellctl vars from inherited env. env = filterEnv(env, []string{ - "TMUX", - "SHELLCTL_STATE_DIR", - "SHELLCTL_RUNTIME_DIR", - "SHELLCTL_TMUX_SOCKET", - "SHELLCTL_RUNNER", - "SHELLCTL_AUTH_TOKEN", + envvar.EnvTMUX, + envvar.EnvShellctlStateDir, + envvar.EnvShellctlRuntimeDir, + envvar.EnvShellctlTmuxSocket, + envvar.EnvShellctlRunner, + envvar.EnvShellctlAuthToken, }) envOverlay := loadEnvJSON(envPath) @@ -172,6 +172,13 @@ func childMode() { home := os.Getenv("HOME") jobDir := filepath.Dir(scriptPath) cfg := landlock.ConfigFromEnv(home, cwd, jobDir) + + // Grant read access to the egress proxy CA cert directory, which + // lives outside the per-binding HOME that Landlock grants. + if caCert := os.Getenv(envvar.EnvSSLCertFile); caCert != "" { + cfg.ROPaths = append(cfg.ROPaths, filepath.Dir(caCert)) + } + if err := landlock.Restrict(cfg); err != nil { // the landlock is best-effort, so we just log the error whatever it is fmt.Fprintf(os.Stderr, "shellctl-runner: WARNING: %v — running without filesystem isolation\n", err) diff --git a/dify-agent-runtime/docker/Dockerfile b/dify-agent-runtime/docker/Dockerfile index ca058727e450ba..7359e808f4f096 100644 --- a/dify-agent-runtime/docker/Dockerfile +++ b/dify-agent-runtime/docker/Dockerfile @@ -73,7 +73,8 @@ COPY --from=go-builder /bin/dify-agent /usr/local/bin/dify-agent RUN useradd --create-home --shell /bin/sh dify \ && mkdir -p /mnt/drive \ && chown dify:dify /home \ - && chown -R dify:dify /home/dify /mnt/drive + && chown -R dify:dify /home/dify /mnt/drive \ + && chown -R dify:dify /usr/local/share/ca-certificates /etc/ssl/certs /etc/ca-certificates.conf USER dify WORKDIR /home/dify diff --git a/dify-agent-runtime/docs/egress-proxy-design.md b/dify-agent-runtime/docs/egress-proxy-design.md new file mode 100644 index 00000000000000..6f9752ca0e2d50 --- /dev/null +++ b/dify-agent-runtime/docs/egress-proxy-design.md @@ -0,0 +1,46 @@ +# Egress Credential Proxy — Demo Guide + +This guide demonstrates the multi-tenant egress credential proxy system for Dify Agent's local sandbox. It walks through configuring a system-level credential manifest, understanding the architecture, and verifying that credentials are injected transparently — without ever appearing as plaintext environment variables in the job process. + +--- + +## Architecture + +``` +┌─────────────────────────────────────────────────────────────────────┐ +│ local_sandbox container │ +│ │ +│ ┌──────────────┐ ┌──────────────────────┐ ┌───────────────┐ │ +│ │ Agent Job │ │ Egress MITM Proxy │ │ Squid (SSRF) │ │ +│ │ (tmux) │───▶│ 127.0.0.1:18080 │───▶│ agent_ssrf │ │ +│ │ │ │ │ │ _proxy:3128 │ │ +│ │ env: │ │ ┌────────────────┐ │ └──────┬────────┘ │ +│ │ HTTP_PROXY │ │ │ Resolver │ │ │ │ +│ │ HTTPS_PROXY │ │ │ system tier │ │ ▼ │ +│ │ TAVILY_API_ │ │ │ session tier │ │ ┌──────────┐ │ +│ │ KEY=__secret│ │ │ (per sandbox) │ │ │ Internet │ │ +│ │ :tavily/ │ │ └────────────────┘ │ │ (e.g. │ │ +│ │ api_key__ │ │ │ │ tavily) │ │ +│ └──────────────┘ │ 1. Inject headers │ └──────────┘ │ +│ │ 2. Strip Proxy-Auth │ │ +│ └──────────────────────┘ │ +│ │ +│ system-credentials.yaml ──▶ loaded at startup into system tier │ +│ (mounted read-only via Docker volume) │ +└─────────────────────────────────────────────────────────────────────┘ +``` + +### Key components + +- **System credential manifest** (`system-credentials.yaml`): Mounted into the container via Docker volume. Parsed at startup (YAML or JSON). Credentials enter the Resolver's **system tier** — shared across all sandbox sessions, never mutated at runtime. + +- **Session credentials**: Registered per sandbox session via `PUT /v1/prepare` API (with `session_id`). Stored in the Resolver's **session tier** — isolated per sandbox, no cross-session leakage. Session credentials shadow system credentials on key conflict. + +- **Egress MITM Proxy** (`127.0.0.1:18080`): Intercepts all outbound HTTP/HTTPS traffic from agent jobs. For HTTPS, it performs TLS interception using a per-container CA (generated fresh at startup, installed into the system trust store). The proxy: + 1. Extracts `session_id` from the `Proxy-Authorization` header (embedded as Basic-Auth userinfo in the proxy URL). + 2. **Proactively injects** credential headers based on domain-matching policies (e.g. `Authorization: Bearer ` for `api.tavily.com`). + 3. Strips the `Proxy-Authorization` header before forwarding. + +- **Squid SSRF proxy** (`agent_ssrf_proxy:3128`): Upstream of the egress proxy. Enforces network-level egress restrictions (deny private networks, allow public internet). + +- **Per-container CA**: Generated at startup by `egressproxy.GenerateCA()`. Installed into the system trust store via `update-ca-certificates` (Dockerfile grants the non-root `dify` user write access to the necessary paths). This means **all** tools — including `apt-get`, `wget`, Java, etc. — trust the MITM proxy's TLS certificates without needing per-tool env vars. diff --git a/dify-agent-runtime/go.mod b/dify-agent-runtime/go.mod index 37f9e7b9e3cec2..1ece066d00846e 100644 --- a/dify-agent-runtime/go.mod +++ b/dify-agent-runtime/go.mod @@ -1,16 +1,20 @@ module github.com/langgenius/dify/dify-agent-runtime -go 1.26 +go 1.26.5 require ( + github.com/aws/aws-sdk-go-v2 v1.43.0 + github.com/elazarl/goproxy v1.8.5 github.com/landlock-lsm/go-landlock v0.9.0 github.com/spf13/cobra v1.10.2 google.golang.org/grpc v1.82.1 google.golang.org/protobuf v1.36.11 + gopkg.in/yaml.v3 v3.0.1 modernc.org/sqlite v1.37.1 ) require ( + github.com/aws/smithy-go v1.27.3 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/google/uuid v1.6.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect @@ -19,9 +23,9 @@ require ( github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/spf13/pflag v1.0.9 // indirect golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0 // indirect - golang.org/x/net v0.55.0 // indirect - golang.org/x/sys v0.45.0 // indirect - golang.org/x/text v0.37.0 // indirect + golang.org/x/net v0.57.0 // indirect + golang.org/x/sys v0.47.0 // indirect + golang.org/x/text v0.40.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect kernel.org/pub/linux/libs/security/libcap/psx v1.2.77 // indirect modernc.org/libc v1.65.7 // indirect diff --git a/dify-agent-runtime/go.sum b/dify-agent-runtime/go.sum index 253a5c6f1d857a..5fa068062ae68e 100644 --- a/dify-agent-runtime/go.sum +++ b/dify-agent-runtime/go.sum @@ -1,8 +1,18 @@ +github.com/aws/aws-sdk-go-v2 v1.43.0 h1:fharf/WhbRAVZ1du0QL7roNFxZ6T/sWr+4Ni617bwSI= +github.com/aws/aws-sdk-go-v2 v1.43.0/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= +github.com/aws/smithy-go v1.27.3 h1:F3Zb497UhhskkfpJmfkXswyo+t0sh9OTBnIHjogWbVY= +github.com/aws/smithy-go v1.27.3/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/coder/websocket v1.8.14 h1:9L0p0iKiNOibykf283eHkKUHHrpG7f65OE3BhhO7v9g= +github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/elazarl/goproxy v1.8.5 h1:33R3Q6geBd2PHmjEI82s3dQWSoBKjTktg4YAFXutIoY= +github.com/elazarl/goproxy v1.8.5/go.mod h1:b5xm6W48AUHNpRTCvlnd0YVh+JafCCtsLsJZvvNTz+E= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= @@ -23,6 +33,8 @@ github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWE github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4= github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= @@ -30,6 +42,8 @@ github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU= github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4= github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY= github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I= @@ -45,19 +59,19 @@ go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLh go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0 h1:R84qjqJb5nVJMxqWYb3np9L5ZsaDtB+a39EqjV0JSUM= golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0/go.mod h1:S9Xr4PYopiDyqSyp5NjCrhFrqg6A5zA2E/iPHPhqnS8= -golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM= -golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU= -golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= -golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ= +golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0= +golang.org/x/net v0.57.0 h1:K5+3DljvIuDG9/Jv9rvyMywYNFCQ9RSUY6OOTTkT+tE= +golang.org/x/net v0.57.0/go.mod h1:KpXc8iv+r3XplLAG/f7Jsf9RPszJzdR0f58q9vGOuEU= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= -golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= -golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= -golang.org/x/tools v0.44.0 h1:UP4ajHPIcuMjT1GqzDWRlalUEoY+uzoZKnhOjbIPD2c= -golang.org/x/tools v0.44.0/go.mod h1:KA0AfVErSdxRZIsOVipbv3rQhVXTnlU6UhKxHd1seDI= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= +golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= +golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q= +golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw= @@ -66,7 +80,10 @@ google.golang.org/grpc v1.82.1 h1:NnAxzGRA0677vCa4BUkOAnO5+FfQqVl9iUXeD0IqcGE= google.golang.org/grpc v1.82.1/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA= google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= kernel.org/pub/linux/libs/security/libcap/psx v1.2.77 h1:Z06sMOzc0GNCwp6efaVrIrz4ywGJ1v+DP0pjVkOfDuA= kernel.org/pub/linux/libs/security/libcap/psx v1.2.77/go.mod h1:+l6Ee2F59XiJ2I6WR5ObpC1utCQJZ/VLsEbQCD8RG24= modernc.org/cc/v4 v4.26.1 h1:+X5NtzVBn0KgsBCBe+xkDC7twLb/jNVj9FPgiwSQO3s= diff --git a/dify-agent-runtime/internal/egressproxy/ca.go b/dify-agent-runtime/internal/egressproxy/ca.go new file mode 100644 index 00000000000000..af5c51d56d9fea --- /dev/null +++ b/dify-agent-runtime/internal/egressproxy/ca.go @@ -0,0 +1,99 @@ +package egressproxy + +import ( + "crypto/rand" + "crypto/rsa" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "fmt" + "math/big" + "os" + "os/exec" + "path/filepath" + "time" +) + +// CAFiles holds the paths to the generated CA certificate and key. +type CAFiles struct { + CertPath string + KeyPath string +} + +// GenerateCA creates a self-signed CA certificate and private key in dir. +// The CA is used by the MITM proxy to generate per-host TLS certificates. +func GenerateCA(dir string) (*CAFiles, error) { + if err := os.MkdirAll(dir, 0700); err != nil { + return nil, fmt.Errorf("egressproxy: mkdir %s: %w", dir, err) + } + + key, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + return nil, fmt.Errorf("egressproxy: generate CA key: %w", err) + } + + serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128)) + if err != nil { + return nil, fmt.Errorf("egressproxy: generate serial: %w", err) + } + + template := &x509.Certificate{ + SerialNumber: serialNumber, + Subject: pkix.Name{ + Organization: []string{"Dify Agent Runtime"}, + CommonName: "Dify Agent Credential Proxy CA", + }, + NotBefore: time.Now().Add(-1 * time.Hour), + NotAfter: time.Now().Add(10 * 365 * 24 * time.Hour), + KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign, + BasicConstraintsValid: true, + IsCA: true, + MaxPathLen: 0, + MaxPathLenZero: true, + } + + certDER, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key) + if err != nil { + return nil, fmt.Errorf("egressproxy: create CA cert: %w", err) + } + + certPath := filepath.Join(dir, "ca.crt") + keyPath := filepath.Join(dir, "ca.key") + + certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER}) + if err := os.WriteFile(certPath, certPEM, 0644); err != nil { + return nil, fmt.Errorf("egressproxy: write CA cert: %w", err) + } + + keyPEM := pem.EncodeToMemory(&pem.Block{ + Type: "RSA PRIVATE KEY", + Bytes: x509.MarshalPKCS1PrivateKey(key), + }) + if err := os.WriteFile(keyPath, keyPEM, 0600); err != nil { + return nil, fmt.Errorf("egressproxy: write CA key: %w", err) + } + + return &CAFiles{CertPath: certPath, KeyPath: keyPath}, nil +} + +// systemTrustAnchorPath is where the CA cert is copied for +// update-ca-certificates to pick up. +const systemTrustAnchorPath = "/usr/local/share/ca-certificates/dify-agent-egress-proxy-ca.crt" + +// InstallSystemTrust copies the CA certificate into the system trust anchors +// and runs update-ca-certificates. Failures are non-fatal; callers should log +// and continue. +func InstallSystemTrust(certPath string) error { + certPEM, err := os.ReadFile(certPath) + if err != nil { + return fmt.Errorf("egressproxy: read CA cert: %w", err) + } + if err := os.WriteFile(systemTrustAnchorPath, certPEM, 0644); err != nil { + return fmt.Errorf("egressproxy: write system trust anchor: %w", err) + } + cmd := exec.Command("update-ca-certificates") + if out, err := cmd.CombinedOutput(); err != nil { + return fmt.Errorf("egressproxy: update-ca-certificates: %w (output: %s)", err, out) + } + return nil +} diff --git a/dify-agent-runtime/internal/egressproxy/certstore.go b/dify-agent-runtime/internal/egressproxy/certstore.go new file mode 100644 index 00000000000000..6ec833f57b1851 --- /dev/null +++ b/dify-agent-runtime/internal/egressproxy/certstore.go @@ -0,0 +1,38 @@ +package egressproxy + +import ( + "crypto/tls" + "sync" +) + +// memCertStore is a simple in-memory cache of MITM leaf certificates, keyed by +// hostname. goproxy recommends caching generated certificates in production to +// avoid repeated CPU-intensive signing for every intercepted CONNECT. +type memCertStore struct { + mu sync.RWMutex + certs map[string]*tls.Certificate +} + +func newMemCertStore() *memCertStore { + return &memCertStore{certs: make(map[string]*tls.Certificate)} +} + +// Fetch implements goproxy.CertStorage. +func (s *memCertStore) Fetch(hostname string, gen func() (*tls.Certificate, error)) (*tls.Certificate, error) { + s.mu.RLock() + cert, ok := s.certs[hostname] + s.mu.RUnlock() + if ok { + return cert, nil + } + + cert, err := gen() + if err != nil { + return nil, err + } + + s.mu.Lock() + s.certs[hostname] = cert + s.mu.Unlock() + return cert, nil +} diff --git a/dify-agent-runtime/internal/egressproxy/proxy.go b/dify-agent-runtime/internal/egressproxy/proxy.go new file mode 100644 index 00000000000000..445685fcda6fd6 --- /dev/null +++ b/dify-agent-runtime/internal/egressproxy/proxy.go @@ -0,0 +1,257 @@ +package egressproxy + +import ( + "crypto/tls" + "encoding/base64" + "fmt" + "log" + "net" + "net/http" + "net/url" + "os" + "regexp" + "strings" + + "github.com/elazarl/goproxy" +) + +// proxyAuthorizationHeader carries the session_id as Basic-Auth userinfo. +const proxyAuthorizationHeader = "Proxy-Authorization" + +// validSessionIDPattern restricts session_id to the same charset/length +// enforced by the server's PrepareCredentials path, so a job cannot supply +// an out-of-contract session_id (e.g. extremely long, path-traversal-shaped) +// to the resolver. Matches server.validSessionIDPattern. +var validSessionIDPattern = regexp.MustCompile(`^[a-zA-Z0-9_-]{1,128}$`) + +// errInvalidSessionID is returned when the Proxy-Authorization userinfo is +// present but does not parse into a valid session_id. Callers should reject +// the request with this error rather than silently proceeding. +var errInvalidSessionID = fmt.Errorf("invalid session_id in Proxy-Authorization") + +// sessionIDFromProxyAuth extracts the session_id embedded as the username of +// a "Proxy-Authorization: Basic ..." header. Returns ("", nil) if the header +// is absent (no session scoping requested). Returns ("", errInvalidSessionID) +// if the header is present but malformed or fails validation. Validation +// prevents cross-session confusion / DoS via out-of-contract session IDs in +// the resolver maps. +func sessionIDFromProxyAuth(h http.Header) (string, error) { + value := h.Get(proxyAuthorizationHeader) + const prefix = "Basic " + if value == "" { + return "", nil + } + if !strings.HasPrefix(value, prefix) { + return "", errInvalidSessionID + } + decoded, err := base64.StdEncoding.DecodeString(value[len(prefix):]) + if err != nil { + return "", errInvalidSessionID + } + sessionID, _, _ := strings.Cut(string(decoded), ":") + if !validSessionIDPattern.MatchString(sessionID) { + return "", errInvalidSessionID + } + return sessionID, nil +} + +const ( + // DefaultListenAddr is the loopback address for the MITM proxy. + DefaultListenAddr = "127.0.0.1:18080" +) + +// Proxy is the credential-injecting MITM forward proxy. +type Proxy struct { + resolver *Resolver + handler *goproxy.ProxyHttpServer + server *http.Server + addr string +} + +// Config holds configuration for the credential proxy. +type Config struct { + // ListenAddr is the address to listen on (default: 127.0.0.1:18080). + ListenAddr string + + // UpstreamProxy is the upstream HTTP proxy URL (e.g. http://agent_ssrf_proxy:3128). + // If empty, the proxy connects directly to upstream servers. + UpstreamProxy string + + // CACertPath is the path to the CA certificate for TLS interception. + CACertPath string + + // CAKeyPath is the path to the CA private key for TLS interception. + CAKeyPath string + + // Resolver is the credential resolver used for header injection. + Resolver *Resolver +} + +// NewProxy creates a new credential proxy but does not start it. +func NewProxy(cfg *Config) (*Proxy, error) { + if cfg.Resolver == nil { + return nil, fmt.Errorf("egressproxy: resolver is required") + } + if cfg.CACertPath == "" || cfg.CAKeyPath == "" { + return nil, fmt.Errorf("egressproxy: CA cert and key paths are required") + } + + addr := cfg.ListenAddr + if addr == "" { + addr = DefaultListenAddr + } + + resolver := cfg.Resolver + + caCertPEM, err := os.ReadFile(cfg.CACertPath) + if err != nil { + return nil, fmt.Errorf("egressproxy: read CA cert: %w", err) + } + caKeyPEM, err := os.ReadFile(cfg.CAKeyPath) + if err != nil { + return nil, fmt.Errorf("egressproxy: read CA key: %w", err) + } + caCert, err := tls.X509KeyPair(caCertPEM, caKeyPEM) + if err != nil { + return nil, fmt.Errorf("egressproxy: parse CA cert/key: %w", err) + } + + px := goproxy.NewProxyHttpServer() + px.Verbose = false + px.Logger = log.New(log.Writer(), "egressproxy: goproxy: ", 0) + px.CertStore = newMemCertStore() + + if cfg.UpstreamProxy != "" { + log.Printf("egressproxy: using upstream proxy: %s", cfg.UpstreamProxy) + upstreamURL, err := url.Parse(cfg.UpstreamProxy) + if err != nil { + return nil, fmt.Errorf("egressproxy: parse upstream proxy url: %w", err) + } + px.Tr.Proxy = http.ProxyURL(upstreamURL) + px.ConnectDial = px.NewConnectDialToProxy(cfg.UpstreamProxy) + } else { + px.Tr.Proxy = nil + px.ConnectDial = nil + } + + mitmAction := &goproxy.ConnectAction{ + Action: goproxy.ConnectMitm, + TLSConfig: goproxy.TLSConfigFromCA(&caCert), + } + rejectAction := &goproxy.ConnectAction{Action: goproxy.ConnectReject} + px.OnRequest().HandleConnectFunc(func(host string, ctx *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) { + sessionID, err := sessionIDFromProxyAuth(ctx.Req.Header) + if err != nil { + log.Printf("egressproxy: rejecting CONNECT %s: %v", host, err) + return rejectAction, host + } + ctx.UserData = sessionID + return mitmAction, host + }) + + px.OnRequest().DoFunc(makeInterceptor(resolver)) + px.OnResponse().DoFunc(makeResponseLogger()) + + return &Proxy{ + resolver: resolver, + handler: px, + addr: addr, + }, nil +} + +// makeInterceptor returns a request handler that injects credential headers +// scoped to the session_id identified for the request. The +// Proxy-Authorization header is stripped before forwarding. Requests +// carrying an invalid session_id are rejected with 400. +func makeInterceptor(resolver *Resolver) func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) { + return func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) { + sessionID, _ := ctx.UserData.(string) + if sessionID == "" { + // HTTP (non-CONNECT) requests don't go through HandleConnectFunc; + // re-extract and validate here. + sid, err := sessionIDFromProxyAuth(req.Header) + if err != nil { + log.Printf("egressproxy: rejecting %s %s: %v", req.Method, req.URL.String(), err) + return req, goproxy.NewResponse(req, goproxy.ContentTypeText, http.StatusBadRequest, "invalid session_id\n") + } + sessionID = sid + } + req.Header.Del(proxyAuthorizationHeader) + + log.Printf("egressproxy: interceptor: %s %s (host=%s, session=%q, effective_creds=%d)", + req.Method, req.URL.String(), req.Host, sessionID, resolver.LenFor(sessionID)) + + if resolver.LenFor(sessionID) == 0 { + return req, nil + } + + resolver.InjectHeadersFor(sessionID, req) + + return req, nil + } +} + +// makeResponseLogger logs the status of each forwarded response. +func makeResponseLogger() func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response { + return func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response { + if resp == nil || ctx.Req == nil { + return resp + } + log.Printf("egressproxy: invoker ok: %s %s -> %d", ctx.Req.Method, ctx.Req.URL.String(), resp.StatusCode) + return resp + } +} + +// Start begins serving the MITM proxy in a background goroutine. +// It returns once the listener is ready. +func (p *Proxy) Start() error { + ln, err := net.Listen("tcp", p.addr) + if err != nil { + return fmt.Errorf("egressproxy: listen %s: %w", p.addr, err) + } + p.addr = ln.Addr().String() + + p.server = &http.Server{ + Handler: p.handler, + } + + go func() { + log.Printf("egressproxy: MITM proxy listening on %s", p.addr) + if err := p.server.Serve(ln); err != nil && err != http.ErrServerClosed { + log.Printf("egressproxy: serve error: %v", err) + } + }() + + return nil +} + +// Stop gracefully shuts down the proxy. +func (p *Proxy) Stop() { + if p.server != nil { + _ = p.server.Close() + } +} + +// Addr returns the actual listen address (useful when port 0 is used). +func (p *Proxy) Addr() string { + return p.addr +} + +// ProxyURL returns the proxy URL without session_id. +func (p *Proxy) ProxyURL() string { + return "http://" + p.addr +} + +// ProxyURLForSession returns the proxy URL with sessionID embedded as +// Basic-Auth userinfo. If sessionID is empty, equivalent to ProxyURL. +func (p *Proxy) ProxyURLForSession(sessionID string) string { + if sessionID == "" { + return p.ProxyURL() + } + u := url.URL{ + Scheme: "http", + User: url.UserPassword(sessionID, ""), + Host: p.addr, + } + return u.String() +} diff --git a/dify-agent-runtime/internal/egressproxy/proxy_test.go b/dify-agent-runtime/internal/egressproxy/proxy_test.go new file mode 100644 index 00000000000000..d297d8eede8ce3 --- /dev/null +++ b/dify-agent-runtime/internal/egressproxy/proxy_test.go @@ -0,0 +1,420 @@ +package egressproxy + +import ( + "bufio" + "crypto/tls" + "crypto/x509" + "encoding/base64" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "os" + "strings" + "sync" + "testing" + + "github.com/langgenius/dify/dify-agent-runtime/internal/providers/simple" +) + +// newTestProxy creates and starts a Proxy backed by a freshly generated CA, +// returning it along with the CA cert pool (for trusting the proxy's MITM +// leaf certs) and a cleanup function. +func newTestProxy(t *testing.T, resolver *Resolver, upstream string) (*Proxy, *x509.CertPool) { + t.Helper() + + caFiles, err := GenerateCA(t.TempDir()) + if err != nil { + t.Fatalf("GenerateCA: %v", err) + } + + proxy, err := NewProxy(&Config{ + ListenAddr: "127.0.0.1:0", + UpstreamProxy: upstream, + CACertPath: caFiles.CertPath, + CAKeyPath: caFiles.KeyPath, + Resolver: resolver, + }) + if err != nil { + t.Fatalf("NewProxy: %v", err) + } + if err := proxy.Start(); err != nil { + t.Fatalf("Start: %v", err) + } + t.Cleanup(proxy.Stop) + + caPEM, err := os.ReadFile(caFiles.CertPath) + if err != nil { + t.Fatalf("read CA cert: %v", err) + } + pool := x509.NewCertPool() + if !pool.AppendCertsFromPEM(caPEM) { + t.Fatalf("failed to add CA cert to pool") + } + return proxy, pool +} + +// clientThroughProxy builds an http.Client that routes through the given +// proxy and trusts caPool for TLS verification of MITM'd leaf certs. +func clientThroughProxy(t *testing.T, proxy *Proxy, caPool *x509.CertPool) *http.Client { + t.Helper() + proxyURL, err := url.Parse(proxy.ProxyURL()) + if err != nil { + t.Fatalf("parse proxy url: %v", err) + } + return &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + TLSClientConfig: &tls.Config{RootCAs: caPool}, + }, + } +} + +// TestProxyHTTPCredentialInjection verifies plain-HTTP forward proxying +// injects credential headers per the Resolver's domain policy. +func TestProxyHTTPCredentialInjection(t *testing.T) { + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("X-Got-Auth", r.Header.Get("Authorization")) + w.WriteHeader(http.StatusOK) + })) + defer backend.Close() + + resolver := NewResolver() + resolver.SetSystemCredentials(map[string]*StoredCredential{ + "token": { + Value: "s3cr3t", + Inject: &simple.Policy{ + HeaderName: "Authorization", + Expr: "Bearer {{.Value}}", + }, + }, + }) + + proxy, caPool := newTestProxy(t, resolver, "") + client := clientThroughProxy(t, proxy, caPool) + + resp, err := client.Get(backend.URL + "/x") + if err != nil { + t.Fatalf("GET through proxy: %v", err) + } + defer func() { _ = resp.Body.Close() }() + + if got := resp.Header.Get("X-Got-Auth"); got != "Bearer s3cr3t" { + t.Fatalf("expected injected Authorization header %q, got %q", "Bearer s3cr3t", got) + } +} + +// TestProxyHTTPSMitmCredentialInjection verifies the proxy MITMs HTTPS +// CONNECT tunnels (decrypting, injecting credentials, and re-encrypting) +// rather than passing them through as opaque tunnels. +func TestProxyHTTPSMitmCredentialInjection(t *testing.T) { + backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("X-Got-Auth", r.Header.Get("Authorization")) + w.WriteHeader(http.StatusOK) + })) + defer backend.Close() + + resolver := NewResolver() + resolver.SetSystemCredentials(map[string]*StoredCredential{ + "token": { + Value: "s3cr3t", + Inject: &simple.Policy{ + HeaderName: "Authorization", + Expr: "Bearer {{.Value}}", + }, + }, + }) + + proxy, caPool := newTestProxy(t, resolver, "") + + // Trust the httptest TLS backend's self-signed cert for the proxy's own + // outbound connection to it (this is a test-only wiring detail; in + // production the proxy dials real destinations with normal cert + // verification). + proxy.handler.Tr.TLSClientConfig = &tls.Config{RootCAs: backend.Client().Transport.(*http.Transport).TLSClientConfig.RootCAs} + + client := clientThroughProxy(t, proxy, caPool) + + resp, err := client.Get(backend.URL + "/x") + if err != nil { + t.Fatalf("GET through proxy (MITM): %v", err) + } + defer func() { _ = resp.Body.Close() }() + + if got := resp.Header.Get("X-Got-Auth"); got != "Bearer s3cr3t" { + t.Fatalf("expected injected Authorization header %q, got %q", "Bearer s3cr3t", got) + } +} + +// fakeUpstreamCONNECTProxy is a minimal upstream proxy that only understands +// CONNECT. It records the literal, unmodified target string from each +// CONNECT request line, then blindly tunnels bytes to realBackendAddr +// regardless of what that target string was (which may not even be +// resolvable) — this is what an upstream like Squid would normally resolve +// and dial itself. +type fakeUpstreamCONNECTProxy struct { + ln net.Listener + target string // realBackendAddr the tunnel is actually wired to + + mu sync.Mutex + seenTargets []string +} + +func newFakeUpstreamCONNECTProxy(t *testing.T, realBackendAddr string) *fakeUpstreamCONNECTProxy { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + p := &fakeUpstreamCONNECTProxy{ln: ln, target: realBackendAddr} + go p.serve(t) + t.Cleanup(func() { _ = ln.Close() }) + return p +} + +func (p *fakeUpstreamCONNECTProxy) Addr() string { + return p.ln.Addr().String() +} + +func (p *fakeUpstreamCONNECTProxy) recordedTargets() []string { + p.mu.Lock() + defer p.mu.Unlock() + out := make([]string, len(p.seenTargets)) + copy(out, p.seenTargets) + return out +} + +func (p *fakeUpstreamCONNECTProxy) serve(t *testing.T) { + for { + conn, err := p.ln.Accept() + if err != nil { + return + } + go p.handle(t, conn) + } +} + +func (p *fakeUpstreamCONNECTProxy) handle(t *testing.T, conn net.Conn) { + defer func() { _ = conn.Close() }() + + br := bufio.NewReader(conn) + req, err := http.ReadRequest(br) + if err != nil { + return + } + if req.Method != http.MethodConnect { + _, _ = conn.Write([]byte("HTTP/1.1 405 Method Not Allowed\r\n\r\n")) + return + } + + // req.Host / req.RequestURI is the literal, verbatim CONNECT target + // string as sent by the client — this is what we assert on. If the + // caller (egressproxy.Proxy) had pre-resolved the hostname to an IP + // before issuing CONNECT, this would observe an IP instead of the + // original hostname. + p.mu.Lock() + p.seenTargets = append(p.seenTargets, req.Host) + p.mu.Unlock() + + backendConn, err := net.Dial("tcp", p.target) + if err != nil { + _, _ = conn.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n")) + return + } + defer func() { _ = backendConn.Close() }() + + _, _ = conn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n")) + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + _, _ = io.Copy(backendConn, br) + }() + go func() { + defer wg.Done() + _, _ = io.Copy(conn, backendConn) + }() + wg.Wait() +} + +// TestProxyUpstreamChainingPreservesHostname is a regression test for the +// bug that motivated migrating from mitmproxy-go to elazarl/goproxy: +// mitmproxy-go always resolved the destination hostname to an IP address in +// its own process before issuing CONNECT to the configured upstream proxy, +// which both discarded information the upstream needed for its own +// hostname-based ACLs/DNS view and broke entirely for hostnames this +// process itself couldn't resolve. goproxy instead forwards the literal, +// unresolved hostname string to the upstream via CONNECT, letting the +// upstream do its own resolution. +func TestProxyUpstreamChainingPreservesHostname(t *testing.T) { + backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer backend.Close() + + backendURL, err := url.Parse(backend.URL) + if err != nil { + t.Fatalf("parse backend url: %v", err) + } + + upstream := newFakeUpstreamCONNECTProxy(t, backendURL.Host) + + resolver := NewResolver() + proxy, caPool := newTestProxy(t, resolver, "http://"+upstream.Addr()) + + // The proxy's own outbound TLS handshake (after decrypting the MITM'd + // tunnel) will use SNI/hostname verification against unresolvableHost + // below, which intentionally does not match the backend's real + // certificate (issued for 127.0.0.1/localhost) — skip verification here + // since this test is only about hostname preservation through the + // upstream CONNECT chain, not about end-to-end cert validation. + proxy.handler.Tr.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} + + // Use a hostname that cannot be resolved by this test process at all. + // If the proxy tried to resolve it locally before CONNECTing upstream + // (the old mitmproxy-go bug), this request would fail outright. + const unresolvableHost = "this-host-does-not-exist.invalid" + _, port, err := net.SplitHostPort(backendURL.Host) + if err != nil { + t.Fatalf("split backend host:port: %v", err) + } + target := "https://" + unresolvableHost + ":" + port + "/x" + + client := clientThroughProxy(t, proxy, caPool) + resp, err := client.Get(target) + if err != nil { + t.Fatalf("GET through proxy chained to upstream: %v", err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode != http.StatusOK { + t.Fatalf("expected 200, got %d", resp.StatusCode) + } + + seen := upstream.recordedTargets() + if len(seen) != 1 { + t.Fatalf("expected exactly 1 CONNECT to reach the upstream, got %d: %v", len(seen), seen) + } + wantTarget := unresolvableHost + ":" + port + if seen[0] != wantTarget { + t.Fatalf("upstream CONNECT target: got %q, want literal unresolved hostname %q", seen[0], wantTarget) + } +} + +func TestSessionIDFromProxyAuthValidation(t *testing.T) { + cases := []struct { + name string + setup func() http.Header + wantID string + wantErr bool + }{ + { + name: "valid", + setup: func() http.Header { + h := http.Header{} + h.Set(proxyAuthorizationHeader, "Basic "+base64.StdEncoding.EncodeToString([]byte("sandbox-a:"))) + return h + }, + wantID: "sandbox-a", + wantErr: false, + }, + { + name: "valid with dashes and underscores", + setup: func() http.Header { + h := http.Header{} + h.Set(proxyAuthorizationHeader, "Basic "+base64.StdEncoding.EncodeToString([]byte("sandbox_a-b:"))) + return h + }, + wantID: "sandbox_a-b", + wantErr: false, + }, + { + name: "missing header (no session scoping)", + setup: func() http.Header { return http.Header{} }, + wantID: "", + wantErr: false, + }, + { + name: "empty user", + setup: func() http.Header { + h := http.Header{} + h.Set(proxyAuthorizationHeader, "Basic "+base64.StdEncoding.EncodeToString([]byte(":"))) + return h + }, + wantID: "", + wantErr: true, + }, + { + name: "non-basic scheme", + setup: func() http.Header { h := http.Header{}; h.Set(proxyAuthorizationHeader, "Bearer token"); return h }, + wantID: "", + wantErr: true, + }, + { + name: "invalid base64", + setup: func() http.Header { + h := http.Header{} + h.Set(proxyAuthorizationHeader, "Basic !!!notbase64!!!") + return h + }, + wantID: "", + wantErr: true, + }, + { + name: "path traversal attempt", + setup: func() http.Header { + h := http.Header{} + h.Set(proxyAuthorizationHeader, "Basic "+base64.StdEncoding.EncodeToString([]byte("../escape:"))) + return h + }, + wantID: "", + wantErr: true, + }, + { + name: "too long", + setup: func() http.Header { + h := http.Header{} + h.Set(proxyAuthorizationHeader, "Basic "+base64.StdEncoding.EncodeToString([]byte(strings.Repeat("a", 129)+":"))) + return h + }, + wantID: "", + wantErr: true, + }, + { + name: "invalid char dot", + setup: func() http.Header { + h := http.Header{} + h.Set(proxyAuthorizationHeader, "Basic "+base64.StdEncoding.EncodeToString([]byte("sandbox.a:"))) + return h + }, + wantID: "", + wantErr: true, + }, + { + name: "invalid char slash", + setup: func() http.Header { + h := http.Header{} + h.Set(proxyAuthorizationHeader, "Basic "+base64.StdEncoding.EncodeToString([]byte("sandbox/a:"))) + return h + }, + wantID: "", + wantErr: true, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, err := sessionIDFromProxyAuth(tc.setup()) + if got != tc.wantID { + t.Errorf("sessionIDFromProxyAuth(%q) id = %q, want %q", tc.name, got, tc.wantID) + } + if tc.wantErr && err == nil { + t.Errorf("sessionIDFromProxyAuth(%q) expected error, got nil", tc.name) + } + if !tc.wantErr && err != nil { + t.Errorf("sessionIDFromProxyAuth(%q) expected no error, got %v", tc.name, err) + } + }) + } +} diff --git a/dify-agent-runtime/internal/egressproxy/resolver.go b/dify-agent-runtime/internal/egressproxy/resolver.go new file mode 100644 index 00000000000000..0c12b6b0e721c2 --- /dev/null +++ b/dify-agent-runtime/internal/egressproxy/resolver.go @@ -0,0 +1,168 @@ +// Package egressproxy implements the in-process egress proxy that runs inside +// the sandbox. It intercepts all outbound HTTP/HTTPS requests and proactively +// injects credentials based on domain-matching policies (see package +// providers). +package egressproxy + +import ( + "log" + "net/http" + "strings" + "sync" + + "github.com/langgenius/dify/dify-agent-runtime/internal/providers" +) + +// StoredCredential holds a credential's value and optional injection policy. +// Value is interpreted by the Inject policy: simple.Policy expects a string +// (or JSON string), aws.Policy expects a structured object (see aws.Credentials). +type StoredCredential struct { + Value any + Inject providers.Policy +} + +// Resolver is a thread-safe credential store scoped by sandbox session. +// It supports proactive header injection based on domain-matching policies. +// +// Credentials live in two independent tiers: +// +// - system holds credentials seeded once at startup (see +// LoadCredentialManifest). It is set via SetSystemCredentials and is +// never touched by session operations. +// - sessions holds one independent credential set per session_id, set via +// SetSessionCredentials (from PUT /v1/prepare). Writing session N's +// credentials never touches session M's map or the system tier — there +// is no shared mutable state across sandbox sessions. +// +// Every lookup is scoped to a sessionID: it checks that session's map first +// and falls back to the system tier. An empty sessionID (no session +// identified) only ever sees the system tier. +type Resolver struct { + mu sync.RWMutex + system map[string]*StoredCredential // key: "provider/name" + sessions map[string]map[string]*StoredCredential // key: sessionID -> "provider/name" +} + +// NewResolver creates an empty credential resolver. +func NewResolver() *Resolver { + return &Resolver{ + system: make(map[string]*StoredCredential), + sessions: make(map[string]map[string]*StoredCredential), + } +} + +// SetSystemCredentials replaces the entire system-tier credential set. +func (r *Resolver) SetSystemCredentials(creds map[string]*StoredCredential) { + if creds == nil { + creds = make(map[string]*StoredCredential) + } + r.mu.Lock() + defer r.mu.Unlock() + r.system = creds +} + +// SetSessionCredentials replaces the credential set for one sandbox session, +// identified by sessionID. +func (r *Resolver) SetSessionCredentials(sessionID string, creds map[string]*StoredCredential) { + if creds == nil { + creds = make(map[string]*StoredCredential) + } + r.mu.Lock() + defer r.mu.Unlock() + r.sessions[sessionID] = creds +} + +// ClearSession removes a sandbox session's credentials. +func (r *Resolver) ClearSession(sessionID string) { + r.mu.Lock() + defer r.mu.Unlock() + delete(r.sessions, sessionID) +} + +// ResolveFor returns the effective credential for ref within sessionID's +// session, falling back to the system tier, or nil if neither has it. An +// empty sessionID only ever resolves against the system tier. +func (r *Resolver) ResolveFor(sessionID, ref string) *StoredCredential { + r.mu.RLock() + defer r.mu.RUnlock() + if sessionID != "" { + if session, ok := r.sessions[sessionID]; ok { + if cred, ok := session[ref]; ok { + return cred + } + } + } + return r.system[ref] +} + +// InjectHeadersFor proactively injects credential-derived headers into the +// request based on domain-matching injection policies, using the effective +// credential set for sessionID (session merged over system). +func (r *Resolver) InjectHeadersFor(sessionID string, req *http.Request) { + host := req.URL.Hostname() + if host == "" { + host = req.Host + } + if idx := strings.LastIndex(host, ":"); idx >= 0 { + host = host[:idx] + } + + r.mu.RLock() + defer r.mu.RUnlock() + for ref, cred := range r.effectiveCredsLocked(sessionID) { + if cred.Inject == nil { + continue + } + if !matchesDomain(host, cred.Inject.Domains()) { + continue + } + if err := cred.Inject.Apply(req, cred.Value); err != nil { + log.Printf("egressproxy: inject credential %q (session=%q): %v", ref, sessionID, err) + } + } +} + +// effectiveCredsLocked returns the merged view of the system tier and +// sessionID's session tier, with the session shadowing the system tier +// under the same ref. Callers must hold r.mu (read or write lock). +func (r *Resolver) effectiveCredsLocked(sessionID string) map[string]*StoredCredential { + session := r.sessions[sessionID] + merged := make(map[string]*StoredCredential, len(r.system)+len(session)) + for ref, cred := range r.system { + merged[ref] = cred + } + for ref, cred := range session { + merged[ref] = cred + } + return merged +} + +// LenFor returns the number of distinct effective credential refs visible to +// sessionID (system tier merged with that session's tier). +func (r *Resolver) LenFor(sessionID string) int { + r.mu.RLock() + defer r.mu.RUnlock() + return len(r.effectiveCredsLocked(sessionID)) +} + +// matchesDomain checks if host matches any of the domain patterns. +// Empty patterns list means match all. Supports "*.example.com" wildcard. +func matchesDomain(host string, patterns []string) bool { + if len(patterns) == 0 { + return true + } + host = strings.ToLower(host) + for _, p := range patterns { + p = strings.ToLower(p) + if p == host { + return true + } + if strings.HasPrefix(p, "*.") { + suffix := p[1:] // ".example.com" + if strings.HasSuffix(host, suffix) { + return true + } + } + } + return false +} diff --git a/dify-agent-runtime/internal/egressproxy/resolver_test.go b/dify-agent-runtime/internal/egressproxy/resolver_test.go new file mode 100644 index 00000000000000..f54c8bfaf7f408 --- /dev/null +++ b/dify-agent-runtime/internal/egressproxy/resolver_test.go @@ -0,0 +1,367 @@ +package egressproxy + +import ( + "io" + "net/http" + "strings" + "testing" + + "github.com/langgenius/dify/dify-agent-runtime/internal/providers/aws" + "github.com/langgenius/dify/dify-agent-runtime/internal/providers/simple" +) + +func TestResolverResolveForSystemTier(t *testing.T) { + r := NewResolver() + r.SetSystemCredentials(map[string]*StoredCredential{"openai/api_key": {Value: "sk-12345"}}) + + cred := r.ResolveFor("", "openai/api_key") + if cred == nil || cred.Value != "sk-12345" { + t.Fatalf("expected sk-12345, got %v", cred) + } + + if r.ResolveFor("", "nonexistent/key") != nil { + t.Fatal("expected nil for unknown ref") + } + // Any sessionID with no session set still sees the system tier. + if cred := r.ResolveFor("some-sandbox", "openai/api_key"); cred == nil || cred.Value != "sk-12345" { + t.Fatalf("expected system fallback for unknown sandbox, got %v", cred) + } +} + +func TestResolverInjectHeadersFor(t *testing.T) { + r := NewResolver() + r.SetSystemCredentials(map[string]*StoredCredential{ + "github/token": { + Value: "ghp_abc123", + Inject: &simple.Policy{ + HeaderName: "Authorization", + Domains_: []string{"*.github.com", "api.github.com"}, + Expr: "Bearer {{.Value}}", + }, + }, + "openai/api_key": { + Value: "sk-xyz", + Inject: &simple.Policy{ + HeaderName: "Authorization", + Domains_: []string{"api.openai.com"}, + Expr: "Bearer {{.Value}}", + }, + }, + }) + + // Request to api.github.com should get github token + req, _ := http.NewRequest("GET", "https://api.github.com/repos", nil) + r.InjectHeadersFor("", req) + if got := req.Header.Get("Authorization"); got != "Bearer ghp_abc123" { + t.Errorf("github request: got %q, want %q", got, "Bearer ghp_abc123") + } + + // Request to api.openai.com should get openai key + req2, _ := http.NewRequest("GET", "https://api.openai.com/v1/chat", nil) + r.InjectHeadersFor("", req2) + if got := req2.Header.Get("Authorization"); got != "Bearer sk-xyz" { + t.Errorf("openai request: got %q, want %q", got, "Bearer sk-xyz") + } + + // Request to unmatched domain gets nothing + req3, _ := http.NewRequest("GET", "https://example.com/api", nil) + r.InjectHeadersFor("", req3) + if got := req3.Header.Get("Authorization"); got != "" { + t.Errorf("unmatched request: got %q, want empty", got) + } +} + +func TestResolverInjectHeadersForSimpleHeaderExprAndErrors(t *testing.T) { + r := NewResolver() + r.SetSystemCredentials(map[string]*StoredCredential{ + "custom/key": { + Value: "abc123", + Inject: &simple.Policy{ + HeaderName: "X-Api-Key", + Expr: "key={{.Value}}", + }, + }, + }) + req, _ := http.NewRequest("GET", "https://example.com/x", nil) + r.InjectHeadersFor("", req) + if got := req.Header.Get("X-Api-Key"); got != "key=abc123" { + t.Errorf("got %q, want %q", got, "key=abc123") + } +} + +func TestMatchesDomain(t *testing.T) { + tests := []struct { + host string + patterns []string + want bool + }{ + {"api.github.com", []string{"*.github.com"}, true}, + {"github.com", []string{"*.github.com"}, false}, + {"api.github.com", []string{"api.github.com"}, true}, + {"evil.com", []string{"api.github.com"}, false}, + {"anything.com", nil, true}, // empty patterns = match all + {"anything.com", []string{}, true}, // empty patterns = match all + } + for _, tc := range tests { + got := matchesDomain(tc.host, tc.patterns) + if got != tc.want { + t.Errorf("matchesDomain(%q, %v) = %v, want %v", tc.host, tc.patterns, got, tc.want) + } + } +} + +func TestResolverClearSession(t *testing.T) { + r := NewResolver() + r.SetSessionCredentials("sandbox-a", map[string]*StoredCredential{"test/key": {Value: "value"}}) + r.ClearSession("sandbox-a") + if r.ResolveFor("sandbox-a", "test/key") != nil { + t.Fatal("expected session credential to be cleared") + } +} + +func TestResolverSessionsAreIsolated(t *testing.T) { + r := NewResolver() + r.SetSessionCredentials("sandbox-a", map[string]*StoredCredential{"a/x": {Value: "1"}}) + r.SetSessionCredentials("sandbox-b", map[string]*StoredCredential{"b/y": {Value: "2"}}) + + if r.ResolveFor("sandbox-a", "b/y") != nil { + t.Fatal("sandbox-a must not see sandbox-b's credentials") + } + if r.ResolveFor("sandbox-b", "a/x") != nil { + t.Fatal("sandbox-b must not see sandbox-a's credentials") + } + if r.LenFor("sandbox-a") != 1 { + t.Fatalf("expected 1 effective credential for sandbox-a, got %d", r.LenFor("sandbox-a")) + } +} + +func TestResolverSessionShadowsSystemWithoutMutation(t *testing.T) { + r := NewResolver() + r.SetSystemCredentials(map[string]*StoredCredential{"custom_saas/api_key": {Value: "sk-system-default"}}) + + if cred := r.ResolveFor("sandbox-a", "custom_saas/api_key"); cred == nil || cred.Value != "sk-system-default" { + t.Fatalf("expected system default, got %v", cred) + } + + // sandbox-a's own session credential shadows the system value. + r.SetSessionCredentials("sandbox-a", map[string]*StoredCredential{"custom_saas/api_key": {Value: "sk-sandbox-a-override"}}) + if cred := r.ResolveFor("sandbox-a", "custom_saas/api_key"); cred == nil || cred.Value != "sk-sandbox-a-override" { + t.Fatalf("expected sandbox-a override, got %v", cred) + } + + // A different, unrelated sandbox must still see only the system default. + if cred := r.ResolveFor("sandbox-b", "custom_saas/api_key"); cred == nil || cred.Value != "sk-system-default" { + t.Fatalf("expected sandbox-b to see system default, got %v", cred) + } + + // Clearing sandbox-a's session must NOT delete the system entry. + r.ClearSession("sandbox-a") + if cred := r.ResolveFor("sandbox-a", "custom_saas/api_key"); cred == nil || cred.Value != "sk-system-default" { + t.Fatalf("expected system default to survive session clear, got %v", cred) + } +} + +func TestResolverInjectHeadersForMergesSystemAndSessionTiers(t *testing.T) { + r := NewResolver() + r.SetSystemCredentials(map[string]*StoredCredential{ + "custom_saas/api_key": { + Value: "sk-system-default", + Inject: &simple.Policy{ + HeaderName: "Authorization", + Domains_: []string{"api.custom-saas.example"}, + Expr: "Bearer {{.Value}}", + }, + }, + }) + + req, _ := http.NewRequest("GET", "https://api.custom-saas.example/v1", nil) + r.InjectHeadersFor("sandbox-a", req) + if got := req.Header.Get("Authorization"); got != "Bearer sk-system-default" { + t.Errorf("got %q, want %q", got, "Bearer sk-system-default") + } + + // sandbox-a's own session credential shadows the system policy too. + r.SetSessionCredentials("sandbox-a", map[string]*StoredCredential{ + "custom_saas/api_key": { + Value: "sk-sandbox-a-override", + Inject: &simple.Policy{ + HeaderName: "Authorization", + Domains_: []string{"api.custom-saas.example"}, + Expr: "Bearer {{.Value}}", + }, + }, + }) + + req2, _ := http.NewRequest("GET", "https://api.custom-saas.example/v1", nil) + r.InjectHeadersFor("sandbox-a", req2) + if got := req2.Header.Get("Authorization"); got != "Bearer sk-sandbox-a-override" { + t.Errorf("got %q, want %q", got, "Bearer sk-sandbox-a-override") + } + + // A different sandbox with no session override still only gets the + // system default injected. + req3, _ := http.NewRequest("GET", "https://api.custom-saas.example/v1", nil) + r.InjectHeadersFor("sandbox-b", req3) + if got := req3.Header.Get("Authorization"); got != "Bearer sk-system-default" { + t.Errorf("got %q, want %q", got, "Bearer sk-system-default") + } +} + +// TestResolverInjectHeadersForAWSSigV4 verifies that the AWS SigV4 policy +// signs a request to an S3 endpoint with real credentials, stripping any +// client-supplied fake signature. +func TestResolverInjectHeadersForAWSSigV4(t *testing.T) { + r := NewResolver() + credJSON := []byte(`{"access_key_id":"AKIAIOSFODNN7EXAMPLE","secret_access_key":"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"}`) + r.SetSystemCredentials(map[string]*StoredCredential{ + "aws/s3_prod": { + Value: credJSON, + Inject: &aws.Policy{ + Domains_: []string{"*.amazonaws.com"}, + Service: "s3", + }, + }, + }) + + // Simulate a curl request (no signature) to S3. + req, _ := http.NewRequest("GET", "https://s3.us-east-1.amazonaws.com/bucket/key", nil) + r.InjectHeadersFor("", req) + + auth := req.Header.Get("Authorization") + if auth == "" || !strings.HasPrefix(auth, "AWS4-HMAC-SHA256 ") { + t.Errorf("expected AWS4-HMAC-SHA256 Authorization header, got %q", auth) + } + if req.Header.Get("X-Amz-Date") == "" { + t.Errorf("expected X-Amz-Date header to be set") + } + if req.Header.Get("X-Amz-Content-Sha256") == "" { + t.Errorf("expected X-Amz-Content-Sha256 header to be set") + } +} + +// TestResolverInjectHeadersForAWSSigV4StripsFakeSignature verifies that a +// client-supplied fake signature is stripped before re-signing with real +// credentials. +func TestResolverInjectHeadersForAWSSigV4StripsFakeSignature(t *testing.T) { + r := NewResolver() + credJSON := []byte(`{"access_key_id":"AKIAIOSFODNN7EXAMPLE","secret_access_key":"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"}`) + r.SetSystemCredentials(map[string]*StoredCredential{ + "aws/s3_prod": { + Value: credJSON, + Inject: &aws.Policy{ + Domains_: []string{"*.amazonaws.com"}, + Service: "s3", + }, + }, + }) + + // Simulate a client that signed with dummy credentials, producing a + // fake signature the proxy must overwrite. + req, _ := http.NewRequest("GET", "https://s3.us-east-1.amazonaws.com/bucket/key", nil) + req.Header.Set("Authorization", "AWS4-HMAC-SHA256 Credential=__secret:aws/s3_prod__/20260731/us-east-1/s3/aws4_request, SignedHeaders=host, Signature=fakesig") + req.Header.Set("X-Amz-Date", "20260731T120000Z") + req.Header.Set("X-Amz-Content-Sha256", "UNSIGNED-PAYLOAD") + r.InjectHeadersFor("", req) + + auth := req.Header.Get("Authorization") + if auth == "" || !strings.HasPrefix(auth, "AWS4-HMAC-SHA256 ") { + t.Errorf("expected real AWS4-HMAC-SHA256 Authorization header, got %q", auth) + } + // The fake signature must have been replaced. + if strings.Contains(auth, "__secret:aws/s3_prod__") { + t.Errorf("fake signature not stripped: %q", auth) + } + if strings.Contains(auth, "fakesig") { + t.Errorf("fake signature not stripped: %q", auth) + } +} + +// TestResolverInjectHeadersForAWSSigV4DomainFiltering verifies that requests +// to non-matching domains are not signed. +func TestResolverInjectHeadersForAWSSigV4DomainFiltering(t *testing.T) { + r := NewResolver() + credJSON := []byte(`{"access_key_id":"AKIAIOSFODNN7EXAMPLE","secret_access_key":"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"}`) + r.SetSystemCredentials(map[string]*StoredCredential{ + "aws/s3_prod": { + Value: credJSON, + Inject: &aws.Policy{ + Domains_: []string{"*.amazonaws.com"}, + Service: "s3", + }, + }, + }) + + // Request to a non-AWS domain should not be signed. + req, _ := http.NewRequest("GET", "https://example.com/api", nil) + r.InjectHeadersFor("", req) + if req.Header.Get("Authorization") != "" { + t.Errorf("expected no Authorization header for non-matching domain, got %q", req.Header.Get("Authorization")) + } +} + +// TestResolverInjectHeadersForAWSSigV4SessionToken verifies that a session +// token is included when present. +func TestResolverInjectHeadersForAWSSigV4SessionToken(t *testing.T) { + r := NewResolver() + credJSON := []byte(`{"access_key_id":"AKIAIOSFODNN7EXAMPLE","secret_access_key":"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY","session_token":"sessiontoken123"}`) + r.SetSystemCredentials(map[string]*StoredCredential{ + "aws/s3_prod": { + Value: credJSON, + Inject: &aws.Policy{ + Domains_: []string{"*.amazonaws.com"}, + Service: "s3", + }, + }, + }) + + req, _ := http.NewRequest("GET", "https://s3.us-east-1.amazonaws.com/bucket/key", nil) + r.InjectHeadersFor("", req) + + if got := req.Header.Get("X-Amz-Security-Token"); got != "sessiontoken123" { + t.Errorf("expected X-Amz-Security-Token to be set, got %q", got) + } +} + +// TestResolverInjectHeadersForAWSSigV4BodyReplay verifies that a POST body +// is correctly hashed and the body remains readable after signing. +func TestResolverInjectHeadersForAWSSigV4BodyReplay(t *testing.T) { + r := NewResolver() + credJSON := []byte(`{"access_key_id":"AKIAIOSFODNN7EXAMPLE","secret_access_key":"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"}`) + r.SetSystemCredentials(map[string]*StoredCredential{ + "aws/s3_prod": { + Value: credJSON, + Inject: &aws.Policy{ + Domains_: []string{"*.amazonaws.com"}, + Service: "s3", + }, + }, + }) + + body := "hello world" + req, _ := http.NewRequest("PUT", "https://s3.us-east-1.amazonaws.com/bucket/key", strings.NewReader(body)) + // Simulate aws cli setting a content-sha256 in "signed body" mode (a + // real 64-char hex hash — the proxy will recompute it from the actual + // body and use that for signing). + req.Header.Set("X-Amz-Content-Sha256", "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa") + r.InjectHeadersFor("", req) + + // The body should still be readable. + readBody, err := io.ReadAll(req.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + if string(readBody) != body { + t.Errorf("body not replayed correctly: got %q, want %q", readBody, body) + } + + // X-Amz-Content-Sha256 should be the real hash (recomputed from body), + // not the fake value the client sent. + sha := req.Header.Get("X-Amz-Content-Sha256") + if sha == "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" { + t.Errorf("expected recomputed SHA-256 hash, got client's fake value") + } + // Verify it's a 64-char hex string. + if len(sha) != 64 { + t.Errorf("expected 64-char hex hash, got %d chars: %q", len(sha), sha) + } +} diff --git a/dify-agent-runtime/internal/envvar/envvar.go b/dify-agent-runtime/internal/envvar/envvar.go index e93af4e828c816..30989c4172cf98 100644 --- a/dify-agent-runtime/internal/envvar/envvar.go +++ b/dify-agent-runtime/internal/envvar/envvar.go @@ -37,6 +37,42 @@ const ( DefaultDriveBase = "/mnt/drive" ) +// --- Egress Proxy --- + +const ( + // EnvEgressProxyEnabled controls whether the in-process egress proxy is started. + EnvEgressProxyEnabled = "SHELLCTL_EGRESSPROXY_ENABLED" + + // EnvEgressProxyAddr overrides the egress proxy listen address (default: 127.0.0.1:18080). + EnvEgressProxyAddr = "SHELLCTL_EGRESSPROXY_ADDR" + + // EnvEgressProxyCADir overrides the directory for the auto-generated CA cert/key. + EnvEgressProxyCADir = "SHELLCTL_EGRESSPROXY_CA_DIR" + + // EnvEgressProxyCACert is set per-job to the CA cert path for TLS trust. + EnvEgressProxyCACert = "SHELLCTL_EGRESSPROXY_CA_CERT" + + // EnvEgressProxyUpstream overrides the upstream proxy URL (empty = direct). + // IMPORTANT: always enable this in docker compose environment + EnvEgressProxyUpstream = "SHELLCTL_EGRESSPROXY_UPSTREAM" + + // EnvEgressProxySystemCredentialsDir points to a directory of credential + // manifest files loaded at startup. All .yaml/.yml/.json files are merged. + EnvEgressProxySystemCredentialsDir = "SHELLCTL_EGRESSPROXY_SYSTEM_CREDENTIALS_DIR" + + // EnvEgressProxySystemCredentialsFile is a legacy alias that points to a + // single credential manifest file. Prefer EnvEgressProxySystemCredentialsDir. + EnvEgressProxySystemCredentialsFile = "SHELLCTL_EGRESSPROXY_SYSTEM_CREDENTIALS_FILE" +) + +const ( + EnvShellctlStateDir = "SHELLCTL_STATE_DIR" + EnvShellctlRuntimeDir = "SHELLCTL_RUNTIME_DIR" + EnvShellctlTmuxSocket = "SHELLCTL_TMUX_SOCKET" + EnvShellctlRunner = "SHELLCTL_RUNNER" + EnvShellctlAuthToken = "SHELLCTL_AUTH_TOKEN" +) + // PathIsolationEnabled returns whether Landlock filesystem isolation is active. func PathIsolationEnabled() bool { v, ok := os.LookupEnv(EnvEnablePathIsolation) diff --git a/dify-agent-runtime/internal/envvar/internal_env.go b/dify-agent-runtime/internal/envvar/internal_env.go new file mode 100644 index 00000000000000..f38c87b8268c74 --- /dev/null +++ b/dify-agent-runtime/internal/envvar/internal_env.go @@ -0,0 +1,32 @@ +package envvar + +// internally used envs +// most of which are well-known ones + +const ( + EnvSSLCertFile = "SSL_CERT_FILE" +) + +// Well-known proxy env vars injected into job environments. +const ( + EnvHTTPProxy = "HTTP_PROXY" + EnvHTTPSProxy = "HTTPS_PROXY" + EnvHTTPProxyLower = "http_proxy" + EnvHTTPSProxyLower = "https_proxy" + EnvNoProxy = "NO_PROXY" + EnvNoProxyLower = "no_proxy" +) + +// Well-known CA cert / trust env vars injected into job environments. +const ( + EnvRequestsCABundle = "REQUESTS_CA_BUNDLE" + EnvNodeExtraCACerts = "NODE_EXTRA_CA_CERTS" + EnvCURLCABundle = "CURL_CA_BUNDLE" + EnvGitSSLCAInfo = "GIT_SSL_CAINFO" + EnvPIPCert = "PIP_CERT" +) + +// Internal shellctl env vars stripped from the inherited job environment. +const ( + EnvTMUX = "TMUX" +) diff --git a/dify-agent-runtime/internal/providers/aws/aws.go b/dify-agent-runtime/internal/providers/aws/aws.go new file mode 100644 index 00000000000000..253f7c9feac070 --- /dev/null +++ b/dify-agent-runtime/internal/providers/aws/aws.go @@ -0,0 +1,354 @@ +// Package aws implements the "aws-sigv4" credential injection policy: +// it re-signs matching requests with AWS Signature Version 4 using the +// credential's structured value (access key id, secret access key, optional +// session token). +// +// Any client-supplied AWS auth headers are stripped before re-signing, so +// both unsigned requests (curl) and requests signed with dummy credentials +// work transparently — the proxy overwrites the signature with real credentials. +// +// The body signing mode is auto-detected from the client's +// X-Amz-Content-Sha256 header: +// +// - hex SHA-256 hash: buffer body (≤10 MiB), compute hash, sign with it. +// - "UNSIGNED-PAYLOAD": sign headers only, no body hash. +// - "STREAMING-UNSIGNED-PAYLOAD-TRAILER": stream body through, sign headers +// only with this constant as the hash. +// - other "STREAMING-*" variants: rejected (cannot reproduce per-chunk +// signatures). +// - absent: treated as "UNSIGNED-PAYLOAD". +// +// Region is extracted from the request hostname (e.g. s3.us-east-1.amazonaws.com) +// unless Region is set explicitly. Service defaults to "s3" if empty. +package aws + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "strings" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + awssigner "github.com/aws/aws-sdk-go-v2/aws/signer/v4" + + "github.com/langgenius/dify/dify-agent-runtime/internal/providers" +) + +// Config is the JSON/YAML payload for the "aws-sigv4" inject type. +type Config struct { + Region string `json:"region,omitempty" yaml:"region,omitempty"` + Service string `json:"service,omitempty" yaml:"service,omitempty"` + Domains []string `json:"domains,omitempty" yaml:"domains,omitempty"` +} + +func init() { + providers.Register("aws-sigv4", func(config json.RawMessage) (providers.Policy, error) { + var c Config + if err := json.Unmarshal(config, &c); err != nil { + return nil, fmt.Errorf("parse aws-sigv4 config: %w", err) + } + return &Policy{ + Domains_: c.Domains, + Region: c.Region, + Service: c.Service, + }, nil + }) +} + +// Credentials holds the three fields needed for AWS Signature Version 4. +type Credentials struct { + AccessKeyID string `json:"access_key_id" yaml:"access_key_id"` + SecretAccessKey string `json:"secret_access_key" yaml:"secret_access_key"` + SessionToken string `json:"session_token,omitempty" yaml:"session_token,omitempty"` +} + +// Policy re-signs matching requests with AWS Signature Version 4. +type Policy struct { + Domains_ []string + Region string // explicit override; empty = extract from hostname + Service string // e.g. "s3", "execute-api"; empty = "s3" +} + +// Compile-time assertion that Policy implements providers.Policy. +var _ interface { + Domains() []string + Apply(*http.Request, any) error +} = (*Policy)(nil) + +// Domains returns the domain-match patterns for this policy. +func (p *Policy) Domains() []string { return p.Domains_ } + +// awsSigV4Headers are the request headers that SigV4 produces and that must +// be stripped before re-signing (whether the client signed with dummy +// credentials or a real key). +var awsSigV4Headers = []string{ + "Authorization", + "X-Amz-Date", + "X-Amz-Content-Sha256", + "X-Amz-Security-Token", +} + +// stripBeforeSign are headers that must be removed before re-signing because +// they are either AWS-SDK-internal (not meaningful to the upstream service) or +// hop-by-hop / intermediary-modifiable (an upstream proxy like Squid may alter +// or remove them, which would break the signature). +var stripBeforeSign = []string{ + "Accept-Encoding", // Squid may modify/remove this + "Amz-Sdk-Invocation-Id", + "Amz-Sdk-Request", + "Amz-Sdk-Request-Attempt", + "X-Amzn-Sdk-Version", + "X-Amzn-Trace-Id", +} + +// MaxBodyBytes is the maximum body size that will be buffered for SHA-256 +// hashing in "signed body" mode. Larger bodies in that mode are rejected. +const MaxBodyBytes = 10 * 1024 * 1024 // 10 MiB + +// Apply strips client AWS auth headers and re-signs req with real credentials. +func (p *Policy) Apply(req *http.Request, value any) error { + log.Printf("aws-sigv4: Apply called for host=%s, method=%s, url=%s", req.URL.Hostname(), req.Method, req.URL.String()) + log.Printf("aws-sigv4: before strip, Authorization=%q, X-Amz-Content-Sha256=%q", req.Header.Get("Authorization"), req.Header.Get("X-Amz-Content-Sha256")) + creds, err := DecodeCredentials(value) + if err != nil { + return fmt.Errorf("decode credentials: %w", err) + } + log.Printf("aws-sigv4: decoded creds, AccessKeyID=%q", creds.AccessKeyID) + + // Determine body signing mode from the client's x-amz-content-sha256. + contentSha := req.Header.Get("X-Amz-Content-Sha256") + mode, err := bodyMode(contentSha) + if err != nil { + return err + } + + // Strip any client-supplied AWS auth headers before re-signing. + for _, h := range awsSigV4Headers { + req.Header.Del(h) + } + // Strip headers that intermediaries (Squid) might modify or that are + // SDK-internal, so they don't end up in the signed headers list. + for _, h := range stripBeforeSign { + req.Header.Del(h) + } + + // Compute payload hash and handle body buffering. + payloadHash, err := p.handleBody(req, mode) + if err != nil { + return err + } + + // Resolve region and service. + region := p.Region + if region == "" { + region, err = ExtractRegionFromHost(req.URL.Hostname()) + if err != nil { + return err + } + } + service := p.Service + if service == "" { + service = "s3" + } + + // Build AWS credentials and signer. + awsCreds := aws.Credentials{ + AccessKeyID: creds.AccessKeyID, + SecretAccessKey: creds.SecretAccessKey, + SessionToken: creds.SessionToken, + } + signer := awssigner.NewSigner() + + // SignHTTP adds Authorization, X-Amz-Date to req.Header. It does NOT + // add X-Amz-Content-Sha256 or X-Amz-Security-Token automatically, so we + // set them explicitly after signing. + if err := signer.SignHTTP(req.Context(), awsCreds, req, payloadHash, service, region, time.Now()); err != nil { + return fmt.Errorf("sign request: %w", err) + } + // Ensure X-Amz-Content-Sha256 is present on the outgoing request so + // the upstream service can verify the payload hash. + req.Header.Set("X-Amz-Content-Sha256", payloadHash) + log.Printf("aws-sigv4: after sign, region=%q, service=%q, payloadHash=%q", region, service, payloadHash) + log.Printf("aws-sigv4: after sign, Authorization=%q", req.Header.Get("Authorization")) + log.Printf("aws-sigv4: after sign, X-Amz-Date=%q, X-Amz-Content-Sha256=%q", req.Header.Get("X-Amz-Date"), req.Header.Get("X-Amz-Content-Sha256")) + // X-Amz-Security-Token is added by the signer when SessionToken is set. + return nil +} + +// bodyMode determines how the request body should be treated during signing, +// based on the client-supplied X-Amz-Content-Sha256 header value. +func bodyMode(contentSha string) (string, error) { + switch { + case contentSha == "": + return "UNSIGNED-PAYLOAD", nil + case contentSha == "UNSIGNED-PAYLOAD": + return "UNSIGNED-PAYLOAD", nil + case contentSha == "STREAMING-UNSIGNED-PAYLOAD-TRAILER": + return "STREAMING-UNSIGNED-PAYLOAD-TRAILER", nil + case strings.HasPrefix(contentSha, "STREAMING-"): + return "", fmt.Errorf("chunk-signed streaming mode %q is not supported (use unsigned payload)", contentSha) + case isHex64(contentSha): + return "signed", nil + default: + return "", fmt.Errorf("unrecognized x-amz-content-sha256 value %q", contentSha) + } +} + +// handleBody processes the request body according to the signing mode and +// returns the payload hash to use for signing. For "signed" mode the body is +// buffered (up to MaxBodyBytes) and its SHA-256 computed. For other modes +// the body is left untouched and the mode string itself is used as the hash +// (per AWS spec). +func (p *Policy) handleBody(req *http.Request, mode string) (string, error) { + switch mode { + case "signed": + if req.Body == nil || req.Body == http.NoBody { + h := sha256.Sum256(nil) + return hex.EncodeToString(h[:]), nil + } + body, err := io.ReadAll(io.LimitReader(req.Body, MaxBodyBytes+1)) + if err != nil { + return "", fmt.Errorf("read body: %w", err) + } + if len(body) > MaxBodyBytes { + return "", fmt.Errorf("body exceeds %d bytes for signed mode", MaxBodyBytes) + } + h := sha256.Sum256(body) + req.Body = io.NopCloser(bytes.NewReader(body)) + req.ContentLength = int64(len(body)) + return hex.EncodeToString(h[:]), nil + case "UNSIGNED-PAYLOAD": + return "UNSIGNED-PAYLOAD", nil + case "STREAMING-UNSIGNED-PAYLOAD-TRAILER": + return "STREAMING-UNSIGNED-PAYLOAD-TRAILER", nil + default: + return mode, nil + } +} + +// isHex64 reports whether s is a 64-character lowercase hex string (a SHA-256 +// digest). +func isHex64(s string) bool { + if len(s) != 64 { + return false + } + for _, c := range s { + if (c < '0' || c > '9') && (c < 'a' || c > 'f') && (c < 'A' || c > 'F') { + return false + } + } + return true +} + +// ExtractRegionFromHost attempts to extract an AWS region from a hostname. +// Recognized patterns: +// - ..amazonaws.com (e.g. s3.us-east-1.amazonaws.com) +// - -.amazonaws.com (e.g. s3-us-west-2.amazonaws.com) +// - .r2.cloudflarestorage.com → "auto" (R2 is region-less) +func ExtractRegionFromHost(host string) (string, error) { + host = strings.ToLower(host) + if strings.HasSuffix(host, ".r2.cloudflarestorage.com") { + return "auto", nil + } + // ..amazonaws.com + if strings.HasSuffix(host, ".amazonaws.com") { + parts := strings.Split(host, ".") + // e.g. ["s3", "us-east-1", "amazonaws", "com"] + if len(parts) >= 4 && parts[len(parts)-2] == "amazonaws" { + region := parts[len(parts)-3] + if isAWSRegion(region) { + return region, nil + } + } + // e.g. ["s3-us-west-2", "amazonaws", "com"] + if len(parts) >= 3 { + first := parts[0] + if idx := strings.Index(first, "-"); idx >= 0 { + region := first[idx+1:] + if isAWSRegion(region) { + return region, nil + } + } + } + } + return "", fmt.Errorf("cannot extract AWS region from host %q (set region in policy)", host) +} + +// isAWSRegion does a light sanity check that the string looks like an AWS +// region (contains a digit and a dash, e.g. "us-east-1", "ap-southeast-2"). +func isAWSRegion(s string) bool { + hasDigit, hasDash := false, false + for _, c := range s { + switch { + case c >= '0' && c <= '9': + hasDigit = true + case c == '-': + hasDash = true + } + } + return hasDigit && hasDash +} + +// DecodeCredentials converts a credential Value into Credentials. Accepts +// Credentials, map[string]any, json.RawMessage, []byte, or a JSON string. +func DecodeCredentials(value any) (*Credentials, error) { + switch v := value.(type) { + case *Credentials: + return v, nil + case Credentials: + return &v, nil + case map[string]any: + return decodeCredsFromMap(v) + case json.RawMessage: + var c Credentials + if err := json.Unmarshal(v, &c); err != nil { + return nil, err + } + if c.AccessKeyID == "" || c.SecretAccessKey == "" { + return nil, fmt.Errorf("access_key_id and secret_access_key are required") + } + return &c, nil + case []byte: + var c Credentials + if err := json.Unmarshal(v, &c); err != nil { + return nil, err + } + if c.AccessKeyID == "" || c.SecretAccessKey == "" { + return nil, fmt.Errorf("access_key_id and secret_access_key are required") + } + return &c, nil + case string: + // Try JSON object first, then treat as raw access key (not supported + // for SigV4 which needs both access key and secret). + var c Credentials + if err := json.Unmarshal([]byte(v), &c); err == nil && c.AccessKeyID != "" && c.SecretAccessKey != "" { + return &c, nil + } + return nil, fmt.Errorf("aws-sigv4 requires structured credentials (access_key_id + secret_access_key)") + default: + return nil, fmt.Errorf("unsupported credential value type %T", value) + } +} + +func decodeCredsFromMap(m map[string]any) (*Credentials, error) { + c := &Credentials{} + if v, ok := m["access_key_id"].(string); ok { + c.AccessKeyID = v + } + if v, ok := m["secret_access_key"].(string); ok { + c.SecretAccessKey = v + } + if v, ok := m["session_token"].(string); ok { + c.SessionToken = v + } + if c.AccessKeyID == "" || c.SecretAccessKey == "" { + return nil, fmt.Errorf("access_key_id and secret_access_key are required") + } + return c, nil +} diff --git a/dify-agent-runtime/internal/providers/aws/aws_test.go b/dify-agent-runtime/internal/providers/aws/aws_test.go new file mode 100644 index 00000000000000..d5cf26982916a4 --- /dev/null +++ b/dify-agent-runtime/internal/providers/aws/aws_test.go @@ -0,0 +1,255 @@ +package aws + +import ( + "encoding/json" + "io" + "net/http" + "strings" + "testing" +) + +func validCredsJSON() json.RawMessage { + return json.RawMessage(`{"access_key_id":"AKIAIOSFODNN7EXAMPLE","secret_access_key":"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"}`) +} + +func validCredsWithTokenJSON() json.RawMessage { + return json.RawMessage(`{"access_key_id":"AKIAIOSFODNN7EXAMPLE","secret_access_key":"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY","session_token":"tok123"}`) +} + +func TestPolicyApplySignsRequest(t *testing.T) { + p := &Policy{Service: "s3"} + req, _ := http.NewRequest("GET", "https://s3.us-east-1.amazonaws.com/bucket/key", nil) + + if err := p.Apply(req, validCredsJSON()); err != nil { + t.Fatalf("Apply: %v", err) + } + + auth := req.Header.Get("Authorization") + if !strings.HasPrefix(auth, "AWS4-HMAC-SHA256 ") { + t.Errorf("expected AWS4-HMAC-SHA256 Authorization, got %q", auth) + } + if req.Header.Get("X-Amz-Date") == "" { + t.Error("expected X-Amz-Date to be set") + } + if req.Header.Get("X-Amz-Content-Sha256") == "" { + t.Error("expected X-Amz-Content-Sha256 to be set") + } +} + +func TestPolicyApplyStripsClientHeaders(t *testing.T) { + p := &Policy{Service: "s3"} + req, _ := http.NewRequest("GET", "https://s3.us-east-1.amazonaws.com/bucket/key", nil) + req.Header.Set("Authorization", "AWS4-HMAC-SHA256 Credential=FAKE/...") + req.Header.Set("X-Amz-Date", "20260101T000000Z") + req.Header.Set("X-Amz-Content-Sha256", "UNSIGNED-PAYLOAD") + req.Header.Set("X-Amz-Security-Token", "faketoken") + + if err := p.Apply(req, validCredsJSON()); err != nil { + t.Fatalf("Apply: %v", err) + } + + auth := req.Header.Get("Authorization") + if strings.Contains(auth, "FAKE") { + t.Errorf("fake Authorization not stripped: %q", auth) + } + // X-Amz-Security-Token should be absent (creds have no session token). + if got := req.Header.Get("X-Amz-Security-Token"); got == "faketoken" { + t.Error("fake X-Amz-Security-Token not stripped") + } +} + +func TestPolicyApplySessionToken(t *testing.T) { + p := &Policy{Service: "s3"} + req, _ := http.NewRequest("GET", "https://s3.us-east-1.amazonaws.com/bucket/key", nil) + + if err := p.Apply(req, validCredsWithTokenJSON()); err != nil { + t.Fatalf("Apply: %v", err) + } + + if got := req.Header.Get("X-Amz-Security-Token"); got != "tok123" { + t.Errorf("expected X-Amz-Security-Token=tok123, got %q", got) + } +} + +func TestPolicyApplyNoSessionToken(t *testing.T) { + p := &Policy{Service: "s3"} + req, _ := http.NewRequest("GET", "https://s3.us-east-1.amazonaws.com/bucket/key", nil) + + if err := p.Apply(req, validCredsJSON()); err != nil { + t.Fatalf("Apply: %v", err) + } + + if got := req.Header.Get("X-Amz-Security-Token"); got != "" { + t.Errorf("expected no X-Amz-Security-Token, got %q", got) + } +} + +func TestPolicyApplyBodyReplay(t *testing.T) { + p := &Policy{Service: "s3"} + body := "hello world" + req, _ := http.NewRequest("PUT", "https://s3.us-east-1.amazonaws.com/bucket/key", strings.NewReader(body)) + // Signed body mode: client provides a hex hash. + req.Header.Set("X-Amz-Content-Sha256", "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa") + + if err := p.Apply(req, validCredsJSON()); err != nil { + t.Fatalf("Apply: %v", err) + } + + // Body should still be readable. + read, err := io.ReadAll(req.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + if string(read) != body { + t.Errorf("body not replayed: got %q, want %q", read, body) + } + + // X-Amz-Content-Sha256 should be recomputed, not the client's fake value. + sha := req.Header.Get("X-Amz-Content-Sha256") + if sha == "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" { + t.Error("expected recomputed hash, got client's fake value") + } + if len(sha) != 64 { + t.Errorf("expected 64-char hex hash, got %d chars: %q", len(sha), sha) + } +} + +func TestPolicyApplyUnsignedPayload(t *testing.T) { + p := &Policy{Service: "s3"} + req, _ := http.NewRequest("GET", "https://s3.us-east-1.amazonaws.com/bucket/key", nil) + // No X-Amz-Content-Sha256 → defaults to UNSIGNED-PAYLOAD. + + if err := p.Apply(req, validCredsJSON()); err != nil { + t.Fatalf("Apply: %v", err) + } + + sha := req.Header.Get("X-Amz-Content-Sha256") + if sha != "UNSIGNED-PAYLOAD" { + t.Errorf("expected UNSIGNED-PAYLOAD, got %q", sha) + } +} + +func TestPolicyApplyRejectsChunkSignedStreaming(t *testing.T) { + p := &Policy{Service: "s3"} + req, _ := http.NewRequest("PUT", "https://s3.us-east-1.amazonaws.com/bucket/key", strings.NewReader("body")) + req.Header.Set("X-Amz-Content-Sha256", "STREAMING-AWS4-HMAC-SHA256-PAYLOAD") + + if err := p.Apply(req, validCredsJSON()); err == nil { + t.Fatal("expected error for chunk-signed streaming mode") + } +} + +func TestPolicyApplyExplicitRegion(t *testing.T) { + p := &Policy{Service: "s3", Region: "eu-west-1"} + req, _ := http.NewRequest("GET", "https://custom.endpoint.example.com/bucket/key", nil) + + if err := p.Apply(req, validCredsJSON()); err != nil { + t.Fatalf("Apply: %v", err) + } + + auth := req.Header.Get("Authorization") + if !strings.Contains(auth, "eu-west-1") { + t.Errorf("expected region eu-west-1 in Authorization, got %q", auth) + } +} + +func TestPolicyApplyR2Region(t *testing.T) { + p := &Policy{Service: "s3"} + req, _ := http.NewRequest("GET", "https://abc123.r2.cloudflarestorage.com/bucket/key", nil) + + if err := p.Apply(req, validCredsJSON()); err != nil { + t.Fatalf("Apply: %v", err) + } + + auth := req.Header.Get("Authorization") + if !strings.Contains(auth, "auto") { + t.Errorf("expected region 'auto' for R2, got %q", auth) + } +} + +func TestPolicyApplyFailsWithoutRegion(t *testing.T) { + p := &Policy{Service: "s3"} + req, _ := http.NewRequest("GET", "https://custom.endpoint.example.com/bucket/key", nil) + + if err := p.Apply(req, validCredsJSON()); err == nil { + t.Fatal("expected error when region cannot be extracted") + } +} + +func TestPolicyApplyInvalidCreds(t *testing.T) { + p := &Policy{Service: "s3"} + req, _ := http.NewRequest("GET", "https://s3.us-east-1.amazonaws.com/bucket/key", nil) + + // Missing secret_access_key. + badCreds := json.RawMessage(`{"access_key_id":"AKIA123"}`) + if err := p.Apply(req, badCreds); err == nil { + t.Fatal("expected error for missing secret_access_key") + } +} + +func TestPolicyApplyStringCredsRejected(t *testing.T) { + p := &Policy{Service: "s3"} + req, _ := http.NewRequest("GET", "https://s3.us-east-1.amazonaws.com/bucket/key", nil) + + // A plain string is not valid for SigV4. + if err := p.Apply(req, "justastring"); err == nil { + t.Fatal("expected error for string credential value") + } +} + +func TestPolicyDomains(t *testing.T) { + p := &Policy{Domains_: []string{"*.s3.amazonaws.com"}} + if got := p.Domains(); len(got) != 1 || got[0] != "*.s3.amazonaws.com" { + t.Errorf("Domains() = %v, want [*.s3.amazonaws.com]", got) + } +} + +func TestDecodeCredentialsFromMap(t *testing.T) { + m := map[string]any{ + "access_key_id": "AKIA123", + "secret_access_key": "secret", + } + c, err := DecodeCredentials(m) + if err != nil { + t.Fatalf("DecodeCredentials: %v", err) + } + if c.AccessKeyID != "AKIA123" || c.SecretAccessKey != "secret" { + t.Errorf("got %+v", c) + } +} + +func TestDecodeCredentialsMissingFields(t *testing.T) { + m := map[string]any{"access_key_id": "AKIA123"} + if _, err := DecodeCredentials(m); err == nil { + t.Fatal("expected error for missing secret_access_key") + } +} + +func TestExtractRegionFromHost(t *testing.T) { + cases := []struct { + host string + want string + err bool + }{ + {"s3.us-east-1.amazonaws.com", "us-east-1", false}, + {"s3-us-west-2.amazonaws.com", "us-west-2", false}, + {"abc123.r2.cloudflarestorage.com", "auto", false}, + {"custom.example.com", "", true}, + } + for _, tc := range cases { + got, err := ExtractRegionFromHost(tc.host) + if tc.err { + if err == nil { + t.Errorf("ExtractRegionFromHost(%q): expected error, got %q", tc.host, got) + } + continue + } + if err != nil { + t.Errorf("ExtractRegionFromHost(%q): unexpected error: %v", tc.host, err) + continue + } + if got != tc.want { + t.Errorf("ExtractRegionFromHost(%q): got %q, want %q", tc.host, got, tc.want) + } + } +} diff --git a/dify-agent-runtime/internal/providers/providers.go b/dify-agent-runtime/internal/providers/providers.go new file mode 100644 index 00000000000000..f7100b7bb40760 --- /dev/null +++ b/dify-agent-runtime/internal/providers/providers.go @@ -0,0 +1,68 @@ +// Package providers defines credential injection policies used by the +// egress proxy. Each policy knows how to interpret a credential's Value +// and inject it into an outbound HTTP request. +// +// Provider packages register themselves at init time via Register, so the +// server package never needs to import individual providers — it only +// imports this package and blank-imports the provider packages for their +// side effects. +// +// A policy is created from API-level configuration (see package server) and +// stored inside egressproxy.StoredCredential. At request time the egress +// proxy calls Apply for every policy whose Domains match the request host. +package providers + +import ( + "encoding/json" + "fmt" + "net/http" + "sync" +) + +// Policy is the interface implemented by every credential injection policy. +// +// - Domains returns the host patterns this policy applies to (empty = all). +// The egress proxy uses this to decide whether to invoke Apply. +// - Apply injects the credential into req. value is the credential's raw +// stored value (e.g. json.RawMessage); the policy is responsible for +// decoding it into the shape it needs. +type Policy interface { + Domains() []string + Apply(req *http.Request, value any) error +} + +// Factory builds a Policy from raw JSON config (the per-type payload from +// the credential manifest, e.g. the "http_header" or "aws_sigv4" object). +type Factory func(config json.RawMessage) (Policy, error) + +var ( + regMu sync.RWMutex + regFactory = map[string]Factory{} +) + +// Register associates name with factory. Called from provider package init(). +// Panics if name is already registered. +func Register(name string, factory Factory) { + regMu.Lock() + defer regMu.Unlock() + if _, exists := regFactory[name]; exists { + panic(fmt.Sprintf("providers: duplicate registration for %q", name)) + } + regFactory[name] = factory +} + +// Build looks up the factory registered under name and invokes it with config. +// Returns (nil, nil) if name is not registered and config is empty/nil. +// Returns an error if name is not registered but config is non-empty. +func Build(name string, config json.RawMessage) (Policy, error) { + regMu.RLock() + factory, ok := regFactory[name] + regMu.RUnlock() + if !ok { + if len(config) == 0 || string(config) == "null" { + return nil, nil + } + return nil, fmt.Errorf("providers: unknown inject type %q", name) + } + return factory(config) +} diff --git a/dify-agent-runtime/internal/providers/simple/simple.go b/dify-agent-runtime/internal/providers/simple/simple.go new file mode 100644 index 00000000000000..1c8eaab99d52e7 --- /dev/null +++ b/dify-agent-runtime/internal/providers/simple/simple.go @@ -0,0 +1,122 @@ +package simple + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "sync" + "text/template" + + "github.com/langgenius/dify/dify-agent-runtime/internal/providers" +) + +// Config is the JSON/YAML payload for the "http-header" inject type. +type Config struct { + Name string `json:"name" yaml:"name"` + Expr string `json:"expr,omitempty" yaml:"expr,omitempty"` + Domains []string `json:"domains,omitempty" yaml:"domains,omitempty"` +} + +func init() { + providers.Register("http-header", func(config json.RawMessage) (providers.Policy, error) { + var c Config + if err := json.Unmarshal(config, &c); err != nil { + return nil, fmt.Errorf("parse http-header config: %w", err) + } + expr := c.Expr + if expr == "" { + expr = "{{.Value}}" + } + return &Policy{ + HeaderName: c.Name, + Domains_: c.Domains, + Expr: expr, + }, nil + }) +} + +// Policy injects a single HTTP header on requests matching Domains. The +// header value is rendered from Expr, a Go text/template evaluated with the +// resolved credential value available as {{.Value}} (for string values) or +// as individual fields (for structured values decoded from JSON, e.g. +// {{.AccessKeyID}}). +type Policy struct { + HeaderName string + Domains_ []string // wildcard-capable domain patterns; empty = all + Expr string // Go text/template rendered with the value + + tmplOnce sync.Once + tmpl *template.Template + tmplErr error +} + +// Compile-time assertion that Policy implements providers.Policy. +var _ interface { + Domains() []string + Apply(*http.Request, any) error +} = (*Policy)(nil) + +// Domains returns the domain-match patterns for this policy. +func (p *Policy) Domains() []string { return p.Domains_ } + +// Apply renders the template and sets the header on req. +func (p *Policy) Apply(req *http.Request, value any) error { + rendered, err := p.render(value) + if err != nil { + return err + } + req.Header.Set(p.HeaderName, rendered) + return nil +} + +// compile lazily parses Expr into a template, caching the result (or error). +func (p *Policy) compile() (*template.Template, error) { + p.tmplOnce.Do(func() { + p.tmpl, p.tmplErr = template.New("simple-header").Parse(p.Expr) + }) + return p.tmpl, p.tmplErr +} + +// render evaluates Expr against the given credential value. The value may be +// a plain string (exposed as {{.Value}}) or a structured value decoded from +// JSON (exposed as its fields, e.g. {{.AccessKeyID}}). +func (p *Policy) render(value any) (string, error) { + tmpl, err := p.compile() + if err != nil { + return "", fmt.Errorf("parse expr %q: %w", p.Expr, err) + } + var buf bytes.Buffer + data := templateData(value) + if err := tmpl.Execute(&buf, data); err != nil { + return "", fmt.Errorf("render expr %q: %w", p.Expr, err) + } + return buf.String(), nil +} + +// templateData converts a credential value into a form suitable for Go +// text/template execution. A string value is wrapped as {.Value: s}; a +// map[string]any or json.RawMessage is decoded so individual fields are +// accessible directly (e.g. {{.AccessKeyID}}) and also via {{.Value}} if +// present. +func templateData(value any) any { + switch v := value.(type) { + case string: + return struct{ Value string }{Value: v} + case json.RawMessage: + var m map[string]any + if err := json.Unmarshal(v, &m); err != nil { + // Fall back to string. + var s string + if err2 := json.Unmarshal(v, &s); err2 == nil { + return struct{ Value string }{Value: s} + } + return struct{ Value string }{Value: string(v)} + } + return m + case map[string]any: + return v + default: + return struct{ Value string }{Value: fmt.Sprint(v)} + } +} diff --git a/dify-agent-runtime/internal/providers/simple/simple_test.go b/dify-agent-runtime/internal/providers/simple/simple_test.go new file mode 100644 index 00000000000000..0650c3114a5a58 --- /dev/null +++ b/dify-agent-runtime/internal/providers/simple/simple_test.go @@ -0,0 +1,87 @@ +package simple + +import ( + "encoding/json" + "net/http" + "testing" +) + +func TestPolicyApplyStringValue(t *testing.T) { + p := &Policy{ + HeaderName: "Authorization", + Expr: "Bearer {{.Value}}", + } + req, _ := http.NewRequest("GET", "https://example.com", nil) + if err := p.Apply(req, "mytoken"); err != nil { + t.Fatalf("Apply: %v", err) + } + if got := req.Header.Get("Authorization"); got != "Bearer mytoken" { + t.Errorf("got %q, want %q", got, "Bearer mytoken") + } +} + +func TestPolicyApplyJSONStringValue(t *testing.T) { + p := &Policy{ + HeaderName: "X-Api-Key", + Expr: "{{.Value}}", + } + req, _ := http.NewRequest("GET", "https://example.com", nil) + // Value is a JSON string literal. + val, _ := json.Marshal("sk-abc") + if err := p.Apply(req, json.RawMessage(val)); err != nil { + t.Fatalf("Apply: %v", err) + } + if got := req.Header.Get("X-Api-Key"); got != "sk-abc" { + t.Errorf("got %q, want %q", got, "sk-abc") + } +} + +func TestPolicyApplyStructuredValue(t *testing.T) { + p := &Policy{ + HeaderName: "Authorization", + Expr: "Bearer {{.access_key_id}}", + } + req, _ := http.NewRequest("GET", "https://example.com", nil) + val := json.RawMessage(`{"access_key_id":"AKIA123","secret_access_key":"secret"}`) + if err := p.Apply(req, val); err != nil { + t.Fatalf("Apply: %v", err) + } + if got := req.Header.Get("Authorization"); got != "Bearer AKIA123" { + t.Errorf("got %q, want %q", got, "Bearer AKIA123") + } +} + +func TestPolicyApplyDefaultExpr(t *testing.T) { + p := &Policy{ + HeaderName: "X-Token", + Expr: "", + } + req, _ := http.NewRequest("GET", "https://example.com", nil) + // Empty expr → template parses as empty, header set to empty. + if err := p.Apply(req, "val"); err != nil { + t.Fatalf("Apply: %v", err) + } + if got := req.Header.Get("X-Token"); got != "" { + t.Errorf("got %q, want empty", got) + } +} + +func TestPolicyApplyBadExpr(t *testing.T) { + p := &Policy{ + HeaderName: "X-Token", + Expr: "{{.Value", // malformed template + } + req, _ := http.NewRequest("GET", "https://example.com", nil) + if err := p.Apply(req, "val"); err == nil { + t.Fatal("expected error for malformed template") + } +} + +func TestPolicyDomains(t *testing.T) { + p := &Policy{ + Domains_: []string{"*.example.com"}, + } + if got := p.Domains(); len(got) != 1 || got[0] != "*.example.com" { + t.Errorf("Domains() = %v, want [*.example.com]", got) + } +} diff --git a/dify-agent-runtime/internal/server/api.go b/dify-agent-runtime/internal/server/api.go index 67c4ee98772945..a6cc8a3ca9ea06 100644 --- a/dify-agent-runtime/internal/server/api.go +++ b/dify-agent-runtime/internal/server/api.go @@ -26,6 +26,7 @@ func Handler(svc *Service, config *Config) http.Handler { mux.HandleFunc("POST /v1/jobs/{job_id}/input", auth(handleInputJob(svc, config))) mux.HandleFunc("POST /v1/jobs/{job_id}/terminate", auth(handleTerminateJob(svc, config))) mux.HandleFunc("DELETE /v1/jobs/{job_id}", auth(handleDeleteJob(svc, config))) + mux.HandleFunc("PUT /v1/prepare", auth(handlePrepare(svc))) return requestLoggingMiddleware(recoveryMiddleware(mux)) } @@ -211,6 +212,29 @@ func handleDeleteJob(svc *Service, config *Config) http.HandlerFunc { } } +func handlePrepare(svc *Service) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + var req PrepareRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeError(w, 400, "invalid_request", "Invalid JSON body") + return + } + if req.SessionID == "" { + writeError(w, 422, "validation_error", "session_id is required") + return + } + if len(req.Credentials) == 0 { + writeError(w, 400, "invalid_request", "credentials must not be empty") + return + } + if err := svc.PrepareCredentials(req.SessionID, req.Credentials); err != nil { + writeServerError(w, err) + return + } + writeJSON(w, http.StatusOK, PrepareResponse{Registered: len(req.Credentials)}) + } +} + // Middleware // statusRecorder wraps ResponseWriter to capture the status code. diff --git a/dify-agent-runtime/internal/server/config.go b/dify-agent-runtime/internal/server/config.go index 226a44b8c434ff..149afe6cb37f03 100644 --- a/dify-agent-runtime/internal/server/config.go +++ b/dify-agent-runtime/internal/server/config.go @@ -5,6 +5,8 @@ import ( "path/filepath" "runtime" "time" + + "github.com/langgenius/dify/dify-agent-runtime/internal/envvar" ) const ( @@ -25,7 +27,7 @@ const ( DefaultPipeMonitorInterval = 1 * time.Second DefaultPipeReadyTimeout = 10 * time.Second DefaultSQLiteBusyTimeoutMs = 5000 - DefaultAuthTokenEnv = "SHELLCTL_AUTH_TOKEN" + DefaultAuthTokenEnv = envvar.EnvShellctlAuthToken HealthStatus = "ok" ) @@ -54,6 +56,12 @@ type Config struct { SQLiteBusyTimeoutMs int SanitizePtyCommand []string RunnerExitCommand []string + + EgressProxyAddr string + EgressProxyCADir string + EgressProxyUpstream string + EgressProxySystemCredentialsDir string + EgressProxySystemCredentials string // legacy single-file mode } // DefaultConfig returns a Config with sensible defaults. @@ -92,9 +100,33 @@ func DefaultConfig() *Config { cfg.AuthToken = os.Getenv(DefaultAuthTokenEnv) } + if v := os.Getenv(envvar.EnvEgressProxySystemCredentialsDir); v != "" { + cfg.EgressProxySystemCredentialsDir = v + } + if v := os.Getenv(envvar.EnvEgressProxySystemCredentialsFile); v != "" { + cfg.EgressProxySystemCredentials = v + } + if v := os.Getenv(envvar.EnvEgressProxyUpstream); v != "" { + cfg.EgressProxyUpstream = v + } + if v := os.Getenv(envvar.EnvEgressProxyAddr); v != "" { + cfg.EgressProxyAddr = v + } + if v := os.Getenv(envvar.EnvEgressProxyCADir); v != "" { + cfg.EgressProxyCADir = v + } + return cfg } +// EgressProxyCAPath returns the directory used for the egress proxy CA files. +func (c *Config) EgressProxyCAPath() string { + if c.EgressProxyCADir != "" { + return c.EgressProxyCADir + } + return filepath.Join(c.RuntimeDir, "egressproxy-ca") +} + // JobsDir returns the path to the jobs artifact directory. func (c *Config) JobsDir() string { return filepath.Join(c.StateDir, "jobs") diff --git a/dify-agent-runtime/internal/server/config_test.go b/dify-agent-runtime/internal/server/config_test.go index 0b6aed9449e9a8..be0e3a04575ec9 100644 --- a/dify-agent-runtime/internal/server/config_test.go +++ b/dify-agent-runtime/internal/server/config_test.go @@ -2,6 +2,8 @@ package server import ( "testing" + + "github.com/langgenius/dify/dify-agent-runtime/internal/envvar" ) func TestDefaultConfig(t *testing.T) { @@ -40,7 +42,7 @@ func TestConfigPaths(t *testing.T) { } func TestConfigAuthTokenFromEnv(t *testing.T) { - t.Setenv("SHELLCTL_AUTH_TOKEN", "my-secret-token") + t.Setenv(envvar.EnvShellctlAuthToken, "my-secret-token") cfg := DefaultConfig() if cfg.AuthToken != "my-secret-token" { t.Errorf("expected auth token from env, got %q", cfg.AuthToken) @@ -48,9 +50,77 @@ func TestConfigAuthTokenFromEnv(t *testing.T) { } func TestConfigNoAuthToken(t *testing.T) { - t.Setenv("SHELLCTL_AUTH_TOKEN", "") + t.Setenv(envvar.EnvShellctlAuthToken, "") cfg := DefaultConfig() if cfg.AuthToken != "" { t.Errorf("expected empty auth token, got %q", cfg.AuthToken) } } + +func TestConfigEgressProxySystemCredentialsFromEnv(t *testing.T) { + t.Setenv(envvar.EnvEgressProxySystemCredentialsFile, "/etc/shellctl/system-credentials.json") + cfg := DefaultConfig() + if cfg.EgressProxySystemCredentials != "/etc/shellctl/system-credentials.json" { + t.Errorf("expected system credentials path from env, got %q", cfg.EgressProxySystemCredentials) + } +} + +func TestConfigNoEgressProxySystemCredentials(t *testing.T) { + t.Setenv(envvar.EnvEgressProxySystemCredentialsFile, "") + cfg := DefaultConfig() + if cfg.EgressProxySystemCredentials != "" { + t.Errorf("expected empty system credentials path, got %q", cfg.EgressProxySystemCredentials) + } +} + +func TestConfigEgressProxySystemCredentialsDirFromEnv(t *testing.T) { + t.Setenv(envvar.EnvEgressProxySystemCredentialsDir, "/etc/shellctl/credentials") + cfg := DefaultConfig() + if cfg.EgressProxySystemCredentialsDir != "/etc/shellctl/credentials" { + t.Errorf("expected system credentials dir from env, got %q", cfg.EgressProxySystemCredentialsDir) + } +} + +func TestConfigNoEgressProxySystemCredentialsDir(t *testing.T) { + t.Setenv(envvar.EnvEgressProxySystemCredentialsDir, "") + cfg := DefaultConfig() + if cfg.EgressProxySystemCredentialsDir != "" { + t.Errorf("expected empty system credentials dir, got %q", cfg.EgressProxySystemCredentialsDir) + } +} + +// TestConfigEgressProxyUpstreamFromEnv is a regression test: without wiring +// SHELLCTL_EGRESSPROXY_UPSTREAM into Config.EgressProxyUpstream, the credproxy +// silently falls back to direct dialing (no upstream chaining), which breaks +// resolution of hostnames only reachable through the upstream SSRF proxy. +func TestConfigEgressProxyUpstreamFromEnv(t *testing.T) { + t.Setenv(envvar.EnvEgressProxyUpstream, "http://agent_ssrf_proxy:3128") + cfg := DefaultConfig() + if cfg.EgressProxyUpstream != "http://agent_ssrf_proxy:3128" { + t.Errorf("expected upstream proxy from env, got %q", cfg.EgressProxyUpstream) + } +} + +func TestConfigNoEgressProxyUpstream(t *testing.T) { + t.Setenv(envvar.EnvEgressProxyUpstream, "") + cfg := DefaultConfig() + if cfg.EgressProxyUpstream != "" { + t.Errorf("expected empty upstream proxy, got %q", cfg.EgressProxyUpstream) + } +} + +func TestConfigEgressProxyAddrFromEnv(t *testing.T) { + t.Setenv(envvar.EnvEgressProxyAddr, "127.0.0.1:19090") + cfg := DefaultConfig() + if cfg.EgressProxyAddr != "127.0.0.1:19090" { + t.Errorf("expected egress proxy addr from env, got %q", cfg.EgressProxyAddr) + } +} + +func TestConfigEgressProxyCADirFromEnv(t *testing.T) { + t.Setenv(envvar.EnvEgressProxyCADir, "/etc/shellctl/ca") + cfg := DefaultConfig() + if cfg.EgressProxyCADir != "/etc/shellctl/ca" { + t.Errorf("expected egress proxy CA dir from env, got %q", cfg.EgressProxyCADir) + } +} diff --git a/dify-agent-runtime/internal/server/service.go b/dify-agent-runtime/internal/server/service.go index a859aa3474db61..e9075f2b05c5de 100644 --- a/dify-agent-runtime/internal/server/service.go +++ b/dify-agent-runtime/internal/server/service.go @@ -2,15 +2,27 @@ package server import ( "context" + "encoding/json" "fmt" "log" "os" "os/exec" "path/filepath" + "regexp" "strconv" "strings" "sync" "time" + + "github.com/langgenius/dify/dify-agent-runtime/internal/egressproxy" + "github.com/langgenius/dify/dify-agent-runtime/internal/envvar" + "github.com/langgenius/dify/dify-agent-runtime/internal/providers" + + // Blank imports register provider factories into the providers registry + // via their init() functions. Add new providers here. + + _ "github.com/langgenius/dify/dify-agent-runtime/internal/providers/aws" + _ "github.com/langgenius/dify/dify-agent-runtime/internal/providers/simple" ) // Service is the core job lifecycle manager backed by SQLite and tmux. @@ -22,14 +34,32 @@ type Service struct { mu sync.Mutex cancelGC context.CancelFunc cancelMon context.CancelFunc + + // Egress proxy (nil when disabled). + egressResolver *egressproxy.Resolver + egressProxy *egressproxy.Proxy + egressCAFiles *egressproxy.CAFiles + // systemCredentials mirrors the refs loaded into egressResolver's system + // tier, kept here (in addition to the resolver) solely so RunJob can + // derive placeholder env var names for them; see + // systemCredentialPlaceholderEnv. Never holds session credentials. + systemCredentials []Credential + // sessionCredentials mirrors, per session_id, the refs registered into + // egressResolver's session tier via PrepareCredentials, kept here solely + // so RunJob can derive placeholder env var names for them; see + // sessionCredentialPlaceholderEnv. Guarded by credMu, independently of mu + // (which only guards startingJobs). + credMu sync.RWMutex + sessionCredentials map[string][]Credential } // NewService creates a new shellctl service. func NewService(config *Config) *Service { return &Service{ - config: config, - tmux: NewTmuxController(config), - startingJobs: make(map[string]bool), + config: config, + tmux: NewTmuxController(config), + startingJobs: make(map[string]bool), + sessionCredentials: make(map[string][]Credential), } } @@ -38,12 +68,297 @@ func (s *Service) Initialize() error { if err := s.PrepareRuntime(); err != nil { return err } + if err := s.initEgressProxy(); err != nil { + return fmt.Errorf("egress proxy: %w", err) + } if err := s.Reconcile(); err != nil { return err } return s.GCOnce() } +// initEgressProxy generates a CA, creates the resolver, and starts the MITM proxy. +func (s *Service) initEgressProxy() error { + upstream := s.config.EgressProxyUpstream + + caDir := s.config.EgressProxyCAPath() + caFiles, err := egressproxy.GenerateCA(caDir) + if err != nil { + return err + } + s.egressCAFiles = caFiles + log.Printf("egressproxy: CA generated in %s", caDir) + + // Best-effort: also install the CA into the system trust store. + if err := egressproxy.InstallSystemTrust(caFiles.CertPath); err != nil { + log.Printf("egressproxy: system trust install failed (falling back to per-tool env vars): %v", err) + } else { + log.Printf("egressproxy: CA installed into system trust store") + } + + resolver := egressproxy.NewResolver() + s.egressResolver = resolver + + // Seed the resolver's system tier with startup-level credentials. + switch { + case s.config.EgressProxySystemCredentialsDir != "": + creds, err := LoadCredentialManifestDir(s.config.EgressProxySystemCredentialsDir) + if err != nil { + return fmt.Errorf("system credentials dir: %w", err) + } + resolver.SetSystemCredentials(credentialsToStoredMap(creds)) + s.systemCredentials = creds + log.Printf("egressproxy: loaded %d system credential(s) from dir %s", len(creds), s.config.EgressProxySystemCredentialsDir) + case s.config.EgressProxySystemCredentials != "": + creds, err := LoadCredentialManifest(s.config.EgressProxySystemCredentials) + if err != nil { + return fmt.Errorf("system credentials manifest: %w", err) + } + resolver.SetSystemCredentials(credentialsToStoredMap(creds)) + s.systemCredentials = creds + log.Printf("egressproxy: loaded %d system credential(s) from %s", len(creds), s.config.EgressProxySystemCredentials) + } + + proxy, err := egressproxy.NewProxy(&egressproxy.Config{ + ListenAddr: s.config.EgressProxyAddr, + UpstreamProxy: upstream, + CACertPath: caFiles.CertPath, + CAKeyPath: caFiles.KeyPath, + Resolver: resolver, + }) + if err != nil { + return err + } + + if err := proxy.Start(); err != nil { + return err + } + s.egressProxy = proxy + if upstream == "" { + log.Printf("egressproxy: MITM proxy started on %s (direct, no upstream)", proxy.Addr()) + } else { + log.Printf("egressproxy: MITM proxy started on %s (upstream: %s)", proxy.Addr(), upstream) + } + return nil +} + +// PrepareCredentials registers creds as the complete credential set for one +// sandbox session (sessionID) and persists them to disk. +func (s *Service) PrepareCredentials(sessionID string, creds []Credential) error { + if s.egressResolver == nil { + return NewServerError(409, "egressproxy_disabled", "Egress proxy is not enabled") + } + if !isValidSessionID(sessionID) { + return NewServerError(422, "validation_error", "session_id must be a non-empty string of letters, digits, '-', or '_' (max 128 chars)") + } + + path := s.sessionCredentialsPath(sessionID) + if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil { + return fmt.Errorf("create session credentials dir: %w", err) + } + data, err := json.Marshal(PrepareRequest{SessionID: sessionID, Credentials: creds}) + if err != nil { + return fmt.Errorf("marshal session credentials: %w", err) + } + if err := writeFileAtomic(path, data, 0600); err != nil { + return fmt.Errorf("write session credentials: %w", err) + } + + s.egressResolver.SetSessionCredentials(sessionID, credentialsToStoredMap(creds)) + + s.credMu.Lock() + if s.sessionCredentials == nil { + s.sessionCredentials = make(map[string][]Credential) + } + s.sessionCredentials[sessionID] = creds + s.credMu.Unlock() + return nil +} + +// sessionCredentialsPath returns the path to sessionID's persisted +// credential manifest under the runtime's credentials directory. +func (s *Service) sessionCredentialsPath(sessionID string) string { + return filepath.Join(s.config.RuntimeDir, "credentials", "sessions", sessionID+".json") +} + +// validSessionIDPattern restricts session_id to characters safe for use both +// as a filename component and as Basic-Auth userinfo. +var validSessionIDPattern = regexp.MustCompile(`^[a-zA-Z0-9_-]{1,128}$`) + +// isValidSessionID reports whether sessionID is safe to use as a session key, +// filename component, and Proxy-Authorization userinfo value. +func isValidSessionID(sessionID string) bool { + return validSessionIDPattern.MatchString(sessionID) +} + +// writeFileAtomic writes data to path via a uniquely-named temp file in the +// same directory + rename, so concurrent callers for the same path never +// race on a shared temp filename. Readers never observe a partially written +// file. +func writeFileAtomic(path string, data []byte, perm os.FileMode) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0700); err != nil { + return err + } + tmp, err := os.CreateTemp(dir, filepath.Base(path)+".*.tmp") + if err != nil { + return err + } + tmpName := tmp.Name() + cleanup := func() { _ = os.Remove(tmpName) } + wrote := false + defer func() { + if !wrote { + cleanup() + } + }() + if _, err := tmp.Write(data); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + if err := os.Chmod(tmpName, perm); err != nil { + return err + } + if err := os.Rename(tmpName, path); err != nil { + return err + } + wrote = true + return nil +} + +// credentialsToStoredMap converts API-level Credential values into the +// resolver's internal StoredCredential representation, keyed by ref. +func credentialsToStoredMap(creds []Credential) map[string]*egressproxy.StoredCredential { + stored := make(map[string]*egressproxy.StoredCredential, len(creds)) + for i := range creds { + c := &creds[i] + stored[c.Ref()] = &egressproxy.StoredCredential{ + Value: json.RawMessage(c.Value), + Inject: buildInjectionPolicy(c.Inject), + } + } + return stored +} + +// buildInjectionPolicy converts an API-level InjectPolicy into a +// providers.Policy via the registry. Returns nil if inject is nil. +func buildInjectionPolicy(inject *InjectPolicy) providers.Policy { + if inject == nil { + return nil + } + policy, err := providers.Build(string(inject.Type), inject.Config) + if err != nil { + log.Printf("egressproxy: build injection policy (type=%s): %v", inject.Type, err) + return nil + } + return policy +} + +// EgressProxyEnv returns the env vars for routing jobs through the egress +// proxy. Returns nil if the egress proxy is disabled. +func (s *Service) EgressProxyEnv(sessionID string) map[string]string { + if s.egressProxy == nil || s.egressCAFiles == nil { + return nil + } + proxyURL := s.egressProxy.ProxyURLForSession(sessionID) + return map[string]string{ + envvar.EnvHTTPProxy: proxyURL, + envvar.EnvHTTPSProxy: proxyURL, + envvar.EnvHTTPProxyLower: proxyURL, + envvar.EnvHTTPSProxyLower: proxyURL, + envvar.EnvNoProxy: "localhost,127.0.0.1", + envvar.EnvNoProxyLower: "localhost,127.0.0.1", + envvar.EnvSSLCertFile: s.egressCAFiles.CertPath, + envvar.EnvRequestsCABundle: s.egressCAFiles.CertPath, + envvar.EnvNodeExtraCACerts: s.egressCAFiles.CertPath, + envvar.EnvCURLCABundle: s.egressCAFiles.CertPath, + envvar.EnvGitSSLCAInfo: s.egressCAFiles.CertPath, + envvar.EnvPIPCert: s.egressCAFiles.CertPath, + } +} + +// systemCredentialPlaceholderEnv returns env var names mapped to +// __secret:provider/name__ placeholders for every system-tier credential, +// so job scripts can reference system credentials by name without seeing +// their real values. Session credentials are excluded. +func (s *Service) systemCredentialPlaceholderEnv() map[string]string { + if len(s.systemCredentials) == 0 { + return nil + } + env := make(map[string]string) + for _, c := range s.systemCredentials { + ph := "__secret:" + c.Ref() + "__" + for _, name := range credentialEnvNames(c) { + env[name] = ph + } + } + return env +} + +// sessionCredentialPlaceholderEnv returns env var names mapped to +// __secret:provider/name__ placeholders for every credential registered to +// sessionID's session. Returns nil for an unknown or empty sessionID. +func (s *Service) sessionCredentialPlaceholderEnv(sessionID string) map[string]string { + if sessionID == "" { + return nil + } + s.credMu.RLock() + creds := s.sessionCredentials[sessionID] + s.credMu.RUnlock() + if len(creds) == 0 { + return nil + } + env := make(map[string]string) + for _, c := range creds { + ph := "__secret:" + c.Ref() + "__" + for _, name := range credentialEnvNames(c) { + env[name] = ph + } + } + return env +} + +// credentialEnvNames returns the environment variable names under which a +// credential's placeholder should be exposed. If EnvNames is set, those are +// used. Otherwise, if EnvName is set, it is used. Otherwise, a name is +// derived from Provider and Name. +func credentialEnvNames(c Credential) []string { + if len(c.EnvNames) > 0 { + return c.EnvNames + } + name := c.EnvName + if name == "" { + name = defaultCredentialEnvName(c.Provider, c.Name) + } + if name == "" { + return nil + } + return []string{name} +} + +// envNameSanitizer matches runs of characters that cannot appear in a POSIX +// environment variable name. +var envNameSanitizer = regexp.MustCompile(`[^A-Za-z0-9]+`) + +// defaultCredentialEnvName derives an environment variable name from a +// credential's provider/name ref (e.g. "github"/"token" -> "GITHUB_TOKEN") +// when no explicit Credential.EnvName is configured. Returns "" if no valid +// name can be derived. +func defaultCredentialEnvName(provider, name string) string { + raw := strings.Trim(envNameSanitizer.ReplaceAllString(provider+"_"+name, "_"), "_") + if raw == "" { + return "" + } + upper := strings.ToUpper(raw) + if upper[0] >= '0' && upper[0] <= '9' { + upper = "_" + upper + } + return upper +} + // PrepareRuntime sets up directories, DB schema, runner script, and tmux server. func (s *Service) PrepareRuntime() error { if err := os.MkdirAll(s.config.StateDir, 0700); err != nil { @@ -82,6 +397,9 @@ func (s *Service) Shutdown() { if s.cancelMon != nil { s.cancelMon() } + if s.egressProxy != nil { + s.egressProxy.Stop() + } if s.db != nil { _ = s.db.Close() } @@ -173,10 +491,48 @@ func (s *Service) RunJob(req *RunJobRequest) (*JobResult, error) { s.cleanupStarting(jobID, jobDir) return nil, err } + + // Merge egress proxy env vars into the job environment. + env := req.Env + if proxyEnv := s.EgressProxyEnv(req.SessionID); proxyEnv != nil { + if env == nil { + env = make(map[string]string) + } + for k, v := range proxyEnv { + if _, exists := env[k]; !exists { + env[k] = v + } + } + } + + // Session credentials take priority over same-named system placeholders. + if placeholderEnv := s.sessionCredentialPlaceholderEnv(req.SessionID); placeholderEnv != nil { + if env == nil { + env = make(map[string]string) + } + for k, v := range placeholderEnv { + if _, exists := env[k]; !exists { + env[k] = v + } + } + } + + // System-tier credential placeholders. + if placeholderEnv := s.systemCredentialPlaceholderEnv(); placeholderEnv != nil { + if env == nil { + env = make(map[string]string) + } + for k, v := range placeholderEnv { + if _, exists := env[k]; !exists { + env[k] = v + } + } + } + envJSON := "{}" - if req.Env != nil { - pairs := make([]string, 0, len(req.Env)) - for k, v := range req.Env { + if env != nil { + pairs := make([]string, 0, len(env)) + for k, v := range env { pairs = append(pairs, fmt.Sprintf("%q:%q", k, v)) } envJSON = "{" + strings.Join(pairs, ",") + "}" diff --git a/dify-agent-runtime/internal/server/types.go b/dify-agent-runtime/internal/server/types.go index 28eecc37b8612c..6d526cc14cf55b 100644 --- a/dify-agent-runtime/internal/server/types.go +++ b/dify-agent-runtime/internal/server/types.go @@ -1,14 +1,32 @@ package server +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + + "gopkg.in/yaml.v3" +) + // RunJobRequest is the HTTP request body for POST /v1/jobs/run. +// +// Credentials are never passed here. Callers must first register them for a +// session_id via PUT /v1/prepare; the egress proxy then proactively injects +// them into outbound HTTP requests based on each credential's inject policy. type RunJobRequest struct { - Script string `json:"script"` - Cwd *string `json:"cwd,omitempty"` - Env map[string]string `json:"env,omitempty"` - Terminal *TerminalSize `json:"terminal,omitempty"` - Timeout float64 `json:"timeout,omitempty"` - OutputLimit int `json:"output_limit,omitempty"` - IdleFlushSeconds float64 `json:"idle_flush_seconds,omitempty"` + Script string `json:"script"` + Cwd *string `json:"cwd,omitempty"` + Env map[string]string `json:"env,omitempty"` + // SessionID identifies which sandbox session's credentials (registered + // via PUT /v1/prepare) apply to this job's egress traffic. Required when + // the egress proxy is enabled; ignored otherwise. + SessionID string `json:"session_id,omitempty"` + Terminal *TerminalSize `json:"terminal,omitempty"` + Timeout float64 `json:"timeout,omitempty"` + OutputLimit int `json:"output_limit,omitempty"` + IdleFlushSeconds float64 `json:"idle_flush_seconds,omitempty"` } // TerminalSize specifies the initial PTY geometry. @@ -88,6 +106,182 @@ type HealthResponse struct { Status string `json:"status"` } +// CredentialValue is the credential's secret value. For simple credentials +// this is a string; for structured credentials (e.g. AWS) this is a JSON +// object. It wraps json.RawMessage so the raw bytes are preserved and +// decoded by the injection policy at request time. It also implements +// yaml.Unmarshaler so YAML manifests can use plain strings or nested maps. +type CredentialValue json.RawMessage + +// UnmarshalJSON implements json.Unmarshaler. It accepts the raw JSON bytes +// directly (string or object), preserving them for later decoding by the +// injection policy. +func (c *CredentialValue) UnmarshalJSON(data []byte) error { + *c = CredentialValue(data) + return nil +} + +// UnmarshalYAML allows CredentialValue to be set from a YAML string or map. +// YAML strings are wrapped as JSON strings; YAML maps are re-encoded as JSON. +func (c *CredentialValue) UnmarshalYAML(node *yaml.Node) error { + // Try string first. + var s string + if err := node.Decode(&s); err == nil { + b, _ := json.Marshal(s) + *c = CredentialValue(b) + return nil + } + // Fall back to a generic map (structured credential). + var m map[string]any + if err := node.Decode(&m); err != nil { + return fmt.Errorf("credential value: expected string or map, got %v", err) + } + b, err := json.Marshal(m) + if err != nil { + return fmt.Errorf("credential value: marshal map: %w", err) + } + *c = CredentialValue(b) + return nil +} + +// Credential represents a secret with its identity and injection policy. +type Credential struct { + // Provider identifies the credential source (e.g. "github", "dify_agent_stub"). + Provider string `json:"provider" yaml:"provider"` + // Name identifies the credential within the provider (e.g. "token", "auth_jwe"). + Name string `json:"name" yaml:"name"` + // Value is the actual secret. For simple credentials this is a string; + // for structured credentials (e.g. AWS) this is a JSON object. + Value CredentialValue `json:"value" yaml:"value"` + // Inject defines how the credential is automatically injected into HTTP + // requests by the egress proxy. Required for the credential to take effect + // at the network layer. + Inject *InjectPolicy `json:"inject,omitempty" yaml:"inject,omitempty"` + // EnvName overrides the environment variable name used to expose this + // credential's __secret:provider/name__ placeholder to jobs. If empty, + // a name is derived from Provider and Name. + EnvName string `json:"env_name,omitempty" yaml:"env_name,omitempty"` + // EnvNames exposes the credential's __secret:provider/name__ placeholder + // under multiple environment variable names. This is useful for + // structured credentials that need to populate several standard env vars + // (e.g. AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN all + // pointing to the same placeholder). If both EnvName and EnvNames are + // set, all names are used. + EnvNames []string `json:"env_names,omitempty" yaml:"env_names,omitempty"` +} + +// InjectType enumerates supported credential injection strategies. +// The actual set of supported types is determined at runtime by the +// providers registry. +type InjectType string + +// InjectPolicy defines how a credential is proactively injected into +// outbound HTTP requests. Type selects the strategy; Config holds the +// raw JSON payload that the corresponding provider package decodes. +type InjectPolicy struct { + Type InjectType `json:"type" yaml:"type"` + Config json.RawMessage `json:"config,omitempty" yaml:"config,omitempty"` +} + +// UnmarshalYAML decodes Type and converts the config map to JSON bytes, +// since json.RawMessage cannot be populated directly from YAML. +func (i *InjectPolicy) UnmarshalYAML(node *yaml.Node) error { + // Walk the mapping nodes manually — yaml.v3 doesn't reliably populate + // *yaml.Node fields via struct decode. + var cfgNode *yaml.Node + for j := 0; j < len(node.Content)-1; j += 2 { + key := node.Content[j].Value + val := node.Content[j+1] + switch key { + case "type": + if err := val.Decode(&i.Type); err != nil { + return fmt.Errorf("inject: parse type: %w", err) + } + case "config": + cfgNode = val + } + } + if cfgNode == nil { + return nil + } + var v any + if err := cfgNode.Decode(&v); err != nil { + return fmt.Errorf("inject: decode config: %w", err) + } + jb, err := json.Marshal(v) + if err != nil { + return fmt.Errorf("inject: marshal config to json: %w", err) + } + i.Config = jb + return nil +} + +// Ref returns the canonical credential reference: "provider/name". +func (c *Credential) Ref() string { + return c.Provider + "/" + c.Name +} + +// PrepareRequest is the HTTP request body for PUT /v1/prepare. +// SessionID scopes these credentials to one sandbox session. +type PrepareRequest struct { + SessionID string `json:"session_id" yaml:"session_id"` + Credentials []Credential `json:"credentials" yaml:"credentials"` +} + +// LoadCredentialManifest reads a credential manifest file and returns its +// credentials. Format is chosen by file extension: .yaml/.yml as YAML, +// everything else as JSON. +func LoadCredentialManifest(path string) ([]Credential, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("read credential manifest %s: %w", path, err) + } + var req PrepareRequest + switch strings.ToLower(filepath.Ext(path)) { + case ".yaml", ".yml": + if err := yaml.Unmarshal(data, &req); err != nil { + return nil, fmt.Errorf("parse credential manifest %s: %w", path, err) + } + default: + if err := json.Unmarshal(data, &req); err != nil { + return nil, fmt.Errorf("parse credential manifest %s: %w", path, err) + } + } + return req.Credentials, nil +} + +// LoadCredentialManifestDir reads all credential manifest files from a +// directory and returns the merged credentials. Only .yaml/.yml/.json files +// are processed. Later files override earlier ones on provider/name conflicts. +func LoadCredentialManifestDir(dir string) ([]Credential, error) { + entries, err := os.ReadDir(dir) + if err != nil { + return nil, fmt.Errorf("read credential manifest dir %s: %w", dir, err) + } + + var all []Credential + for _, entry := range entries { + if entry.IsDir() { + continue + } + ext := strings.ToLower(filepath.Ext(entry.Name())) + if ext != ".yaml" && ext != ".yml" && ext != ".json" { + continue + } + creds, err := LoadCredentialManifest(filepath.Join(dir, entry.Name())) + if err != nil { + return nil, err + } + all = append(all, creds...) + } + return all, nil +} + +// PrepareResponse is the response for PUT /v1/prepare. +type PrepareResponse struct { + Registered int `json:"registered"` +} + // ErrorDetail is the machine-readable API error payload. type ErrorDetail struct { Code string `json:"code"` diff --git a/dify-agent-runtime/internal/server/types_test.go b/dify-agent-runtime/internal/server/types_test.go new file mode 100644 index 00000000000000..c78c7ff1dbeb32 --- /dev/null +++ b/dify-agent-runtime/internal/server/types_test.go @@ -0,0 +1,364 @@ +package server + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/langgenius/dify/dify-agent-runtime/internal/egressproxy" +) + +// jsonStr wraps a Go string as a CredentialValue (JSON string literal), for +// use as Credential.Value in tests. +func jsonStr(s string) CredentialValue { + b, _ := json.Marshal(s) + return CredentialValue(b) +} + +// rawStr extracts a Go string from a Credential.Value (CredentialValue) or +// StoredCredential.Value (any holding json.RawMessage). Panics on failure. +func rawStr(v any) string { + switch x := v.(type) { + case CredentialValue: + var s string + if err := json.Unmarshal(x, &s); err != nil { + panic(err) + } + return s + case json.RawMessage: + var s string + if err := json.Unmarshal(x, &s); err != nil { + panic(err) + } + return s + case string: + return x + default: + panic("unexpected value type") + } +} + +func TestLoadCredentialManifest(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "system-credentials.json") + manifest := `{ + "credentials": [ + { + "provider": "custom_saas", + "name": "api_key", + "value": "sk-system-default", + "inject": { + "type": "http-header", + "config": { + "name": "Authorization", + "expr": "Bearer {{.Value}}", + "domains": ["api.custom-saas.example"] + } + } + } + ] + }` + if err := os.WriteFile(path, []byte(manifest), 0600); err != nil { + t.Fatalf("write manifest: %v", err) + } + + creds, err := LoadCredentialManifest(path) + if err != nil { + t.Fatalf("LoadCredentialManifest: %v", err) + } + if len(creds) != 1 { + t.Fatalf("expected 1 credential, got %d", len(creds)) + } + if creds[0].Ref() != "custom_saas/api_key" || rawStr(creds[0].Value) != "sk-system-default" { + t.Errorf("unexpected credential: %+v", creds[0]) + } +} + +func TestLoadCredentialManifestYAML(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "system-credentials.yaml") + manifest := ` +credentials: + - provider: custom_saas + name: api_key + value: sk-system-default + inject: + type: http-header + config: + name: Authorization + expr: "Bearer {{.Value}}" + domains: + - api.custom-saas.example +` + if err := os.WriteFile(path, []byte(manifest), 0600); err != nil { + t.Fatalf("write manifest: %v", err) + } + + creds, err := LoadCredentialManifest(path) + if err != nil { + t.Fatalf("LoadCredentialManifest: %v", err) + } + if len(creds) != 1 { + t.Fatalf("expected 1 credential, got %d", len(creds)) + } + if creds[0].Ref() != "custom_saas/api_key" || rawStr(creds[0].Value) != "sk-system-default" { + t.Errorf("unexpected credential: %+v", creds[0]) + } + if creds[0].Inject == nil || creds[0].Inject.Type != "http-header" { + t.Errorf("expected parsed inject policy with type http-header, got %+v", creds[0].Inject) + } + // Config should contain the raw JSON for the http_header payload. + var cfg struct { + Name string `json:"name"` + } + if err := json.Unmarshal(creds[0].Inject.Config, &cfg); err != nil { + t.Fatalf("unmarshal inject config: %v", err) + } + if cfg.Name != "Authorization" { + t.Errorf("expected inject config name=Authorization, got %q", cfg.Name) + } +} + +func TestLoadCredentialManifestEmptyYAML(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "system-credentials.yaml") + if err := os.WriteFile(path, []byte("credentials: []\n"), 0600); err != nil { + t.Fatalf("write manifest: %v", err) + } + + creds, err := LoadCredentialManifest(path) + if err != nil { + t.Fatalf("LoadCredentialManifest: %v", err) + } + if len(creds) != 0 { + t.Fatalf("expected 0 credentials, got %d", len(creds)) + } +} + +func TestLoadCredentialManifestInvalidYAML(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "bad.yaml") + if err := os.WriteFile(path, []byte("credentials: [not: valid: yaml"), 0600); err != nil { + t.Fatalf("write manifest: %v", err) + } + if _, err := LoadCredentialManifest(path); err == nil { + t.Fatal("expected error for invalid YAML manifest") + } +} + +func TestLoadCredentialManifestMissingFile(t *testing.T) { + if _, err := LoadCredentialManifest("/nonexistent/path.json"); err == nil { + t.Fatal("expected error for missing manifest file") + } +} + +func TestLoadCredentialManifestInvalidJSON(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "bad.json") + if err := os.WriteFile(path, []byte("not json"), 0600); err != nil { + t.Fatalf("write manifest: %v", err) + } + if _, err := LoadCredentialManifest(path); err == nil { + t.Fatal("expected error for invalid JSON manifest") + } +} + +func TestLoadCredentialManifestDir(t *testing.T) { + dir := t.TempDir() + + // Write two manifest files and one non-manifest file (should be skipped). + if err := os.WriteFile(filepath.Join(dir, "tavily.yaml"), []byte(` +credentials: + - provider: tavily + name: api_key + value: tvly-aaa +`), 0600); err != nil { + t.Fatalf("write tavily.yaml: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "github.json"), []byte(`{ + "credentials": [ + {"provider": "github", "name": "token", "value": "ghp-bbb"} + ] +}`), 0600); err != nil { + t.Fatalf("write github.json: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "README.md"), []byte("not a manifest"), 0600); err != nil { + t.Fatalf("write README.md: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, ".gitignore"), []byte("*.cred.yaml"), 0600); err != nil { + t.Fatalf("write .gitignore: %v", err) + } + + creds, err := LoadCredentialManifestDir(dir) + if err != nil { + t.Fatalf("LoadCredentialManifestDir: %v", err) + } + if len(creds) != 2 { + t.Fatalf("expected 2 credentials, got %d", len(creds)) + } + + refs := map[string]string{} + for _, c := range creds { + refs[c.Ref()] = rawStr(c.Value) + } + if refs["tavily/api_key"] != "tvly-aaa" { + t.Errorf("tavily/api_key: got %q", refs["tavily/api_key"]) + } + if refs["github/token"] != "ghp-bbb" { + t.Errorf("github/token: got %q", refs["github/token"]) + } +} + +func TestLoadCredentialManifestDirEmpty(t *testing.T) { + dir := t.TempDir() + creds, err := LoadCredentialManifestDir(dir) + if err != nil { + t.Fatalf("LoadCredentialManifestDir on empty dir: %v", err) + } + if len(creds) != 0 { + t.Fatalf("expected 0 credentials from empty dir, got %d", len(creds)) + } +} + +func TestLoadCredentialManifestDirMissing(t *testing.T) { + if _, err := LoadCredentialManifestDir("/nonexistent/credentials"); err == nil { + t.Fatal("expected error for missing directory") + } +} + +// newTestService builds a minimal Service with an egress resolver and a +// scratch RuntimeDir, sufficient for exercising PrepareCredentials. +func newTestService(t *testing.T) *Service { + t.Helper() + return &Service{ + config: &Config{RuntimeDir: t.TempDir()}, + egressResolver: egressproxy.NewResolver(), + } +} + +// TestSessionCredentialsShadowSystemWithoutMutation verifies that once +// system-level credentials are seeded (e.g. loaded at startup from +// LoadCredentialManifest), a sandbox session's own credentials (registered +// via PrepareCredentials) take priority for resolution scoped to that +// session_id, without ever mutating the system tier or leaking to other +// sandbox sessions. +func TestSessionCredentialsShadowSystemWithoutMutation(t *testing.T) { + s := newTestService(t) + + s.egressResolver.SetSystemCredentials(credentialsToStoredMap([]Credential{ + { + Provider: "custom_saas", + Name: "api_key", + Value: jsonStr("sk-system-default"), + Inject: &InjectPolicy{ + Type: "http-header", + Config: json.RawMessage(`{"name":"Authorization","expr":"Bearer {{.Value}}","domains":["api.custom-saas.example"]}`), + }, + }, + })) + + // No session_id yet: only the system default is visible. + if cred := s.egressResolver.ResolveFor("sandbox-a", "custom_saas/api_key"); cred == nil || rawStr(cred.Value) != "sk-system-default" { + t.Fatalf("expected system credential, got %v", cred) + } + + // sandbox-a registers its own override via PUT /v1/prepare. + if err := s.PrepareCredentials("sandbox-a", []Credential{ + {Provider: "custom_saas", Name: "api_key", Value: jsonStr("sk-sandbox-a-override")}, + }); err != nil { + t.Fatalf("PrepareCredentials: %v", err) + } + + if cred := s.egressResolver.ResolveFor("sandbox-a", "custom_saas/api_key"); cred == nil || rawStr(cred.Value) != "sk-sandbox-a-override" { + t.Fatalf("expected sandbox-a override, got %v", cred) + } + + // A different sandbox session must still see only the system default: + // sandbox-a's registration must not leak across sessions. + if cred := s.egressResolver.ResolveFor("sandbox-b", "custom_saas/api_key"); cred == nil || rawStr(cred.Value) != "sk-system-default" { + t.Fatalf("expected sandbox-b to see system default, got %v", cred) + } + + // The persisted manifest file for sandbox-a must exist on disk. + if _, err := os.Stat(s.sessionCredentialsPath("sandbox-a")); err != nil { + t.Fatalf("expected persisted session credentials file: %v", err) + } +} + +func TestPrepareCredentialsRejectsInvalidSessionID(t *testing.T) { + s := newTestService(t) + err := s.PrepareCredentials("../escape", []Credential{{Provider: "p", Name: "n", Value: jsonStr("v")}}) + if err == nil { + t.Fatal("expected error for invalid session_id") + } +} + +func TestPrepareCredentialsRequiresEgressProxyEnabled(t *testing.T) { + s := &Service{config: &Config{RuntimeDir: t.TempDir()}} + err := s.PrepareCredentials("sandbox-a", []Credential{{Provider: "p", Name: "n", Value: jsonStr("v")}}) + if err == nil { + t.Fatal("expected error when egress proxy is disabled") + } +} + +func TestDefaultCredentialEnvName(t *testing.T) { + cases := []struct { + provider, name, want string + }{ + {"github", "token", "GITHUB_TOKEN"}, + {"custom-saas", "api.key", "CUSTOM_SAAS_API_KEY"}, + {"dify_agent_stub", "auth_jwe", "DIFY_AGENT_STUB_AUTH_JWE"}, + {"123provider", "name", "_123PROVIDER_NAME"}, + {"", "", ""}, + } + for _, c := range cases { + if got := defaultCredentialEnvName(c.provider, c.name); got != c.want { + t.Errorf("defaultCredentialEnvName(%q, %q) = %q, want %q", c.provider, c.name, got, c.want) + } + } +} + +// TestSystemCredentialPlaceholderEnvInjectedIntoJob verifies that system-tier +// credentials are exposed to every job as __secret:provider/name__ +// placeholder env vars by default, so a caller doesn't need to know or +// reproduce a credential's ref manually to make use of it. +func TestSystemCredentialPlaceholderEnvInjectedIntoJob(t *testing.T) { + s := newTestService(t) + s.systemCredentials = []Credential{ + {Provider: "custom_saas", Name: "api_key", Value: jsonStr("sk-system-default")}, + {Provider: "explicit", Name: "ref", Value: jsonStr("sk-explicit"), EnvName: "MY_CUSTOM_ENV"}, + } + + env := s.systemCredentialPlaceholderEnv() + if got, want := env["CUSTOM_SAAS_API_KEY"], "__secret:custom_saas/api_key__"; got != want { + t.Errorf("derived env name: got %q, want %q", got, want) + } + if got, want := env["MY_CUSTOM_ENV"], "__secret:explicit/ref__"; got != want { + t.Errorf("explicit EnvName: got %q, want %q", got, want) + } +} + +// TestSessionCredentialPlaceholderEnvScopedToSandbox verifies that a +// sandbox's own registered credentials (via PrepareCredentials) are exposed +// as placeholder env vars only for that session_id, never for others. +func TestSessionCredentialPlaceholderEnvScopedToSandbox(t *testing.T) { + s := newTestService(t) + if err := s.PrepareCredentials("sandbox-a", []Credential{ + {Provider: "myprovider", Name: "mysecret", Value: jsonStr("sk-sandbox-a")}, + }); err != nil { + t.Fatalf("PrepareCredentials: %v", err) + } + + env := s.sessionCredentialPlaceholderEnv("sandbox-a") + if got, want := env["MYPROVIDER_MYSECRET"], "__secret:myprovider/mysecret__"; got != want { + t.Errorf("sandbox-a env: got %q, want %q", got, want) + } + + if env := s.sessionCredentialPlaceholderEnv("sandbox-b"); env != nil { + t.Errorf("expected no placeholder env for a different sandbox, got %v", env) + } + if env := s.sessionCredentialPlaceholderEnv(""); env != nil { + t.Errorf("expected no placeholder env for an empty session_id, got %v", env) + } +} diff --git a/dify-agent-runtime/tests/egress_proxy_test.go b/dify-agent-runtime/tests/egress_proxy_test.go new file mode 100644 index 00000000000000..773df21f5340c8 --- /dev/null +++ b/dify-agent-runtime/tests/egress_proxy_test.go @@ -0,0 +1,244 @@ +//go:build integration + +// This file verifies the egress proxy's credential injection against a real +// container: a dedicated dify-agent-runtime container (SHELLCTL_EGRESSPROXY_ENABLED=true, +// HTTP(S)_PROXY pointed at the in-process MITM proxy) plus an echo backend +// reachable only from inside that container's docker network as "echo-backend". +// A job script issues a real outbound curl request; the echo backend reflects +// back the headers it received, which we assert against to prove the proxy +// actually injected credentials over the wire. +// +// Provisioned by `make integration-up` (see Makefile) and exercised via +// `make integration-test` / `make integration`. +package tests + +import ( + "bytes" + "encoding/json" + "net/http" + "os" + "testing" +) + +var ( + egressGoURL = os.Getenv("SHELLCTL_EGRESS_GO_URL") + egressAuthToken = os.Getenv("SHELLCTL_EGRESS_TEST_TOKEN") + + // egressUpstreamGoURL/egressUpstreamAuthToken target a second runtime + // container whose SHELLCTL_EGRESSPROXY_UPSTREAM points at a squid + // container ("squid-upstream"), verifying that the credproxy correctly + // chains through an upstream forward proxy to reach the real destination. + egressUpstreamGoURL = os.Getenv("SHELLCTL_EGRESS_UPSTREAM_GO_URL") + egressUpstreamAuthToken = os.Getenv("SHELLCTL_EGRESS_UPSTREAM_TEST_TOKEN") +) + +func egressTarget() (target, bool) { + if egressGoURL == "" { + return target{}, false + } + return target{name: "go-egress", baseURL: egressGoURL}, true +} + +func egressUpstreamTarget() (target, bool) { + if egressUpstreamGoURL == "" { + return target{}, false + } + return target{name: "go-egress-upstream", baseURL: egressUpstreamGoURL}, true +} + +// doPutWithToken issues a PUT request with a bearer token, mirroring doPost's +// shape but for the PUT /v1/prepare endpoint. +func doPutWithToken(t *testing.T, tgt target, token, path string, payload map[string]any) *http.Response { + t.Helper() + body, _ := json.Marshal(payload) + req, _ := http.NewRequest(http.MethodPut, tgt.baseURL+path, bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+token) + resp, err := httpClient.Do(req) + if err != nil { + t.Fatalf("[%s] PUT %s failed: %v", tgt.name, path, err) + } + return resp +} + +// TestEgressProxyCredentialInjection verifies that a credential registered via +// PUT /v1/prepare with a domain-scoped http-header injection rule is actually +// injected into an outbound request made by a job script, by round-tripping +// through the real in-process MITM egress proxy to an echo backend. +func TestEgressProxyCredentialInjection(t *testing.T) { + tgt, ok := egressTarget() + if !ok { + t.Skip("SHELLCTL_EGRESS_GO_URL not set; egress proxy container not available") + } + + const sessionID = "sandbox-credential-injection" + prepareResp := doPutWithToken(t, tgt, egressAuthToken, "/v1/prepare", map[string]any{ + "session_id": sessionID, + "credentials": []map[string]any{ + { + "provider": "testprovider", + "name": "apikey", + "value": "sk-integration-test-secret", + "inject": map[string]any{ + "type": "http-header", + "config": map[string]any{ + "name": "Authorization", + "expr": "Bearer {{.Value}}", + "domains": []string{"echo-backend"}, + }, + }, + }, + }, + }) + assertStatus(t, prepareResp, 200) + readBody(t, prepareResp) + + result := runJobWithToken(t, tgt, egressAuthToken, map[string]any{ + "script": "curl -s http://echo-backend:8080/", + "timeout": 15, + "session_id": sessionID, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + + output := result["output"].(string) + var echoed map[string]any + if err := json.Unmarshal([]byte(output), &echoed); err != nil { + t.Fatalf("failed to parse echo backend response: %v\noutput: %s", err, output) + } + headers, ok := echoed["headers"].(map[string]any) + if !ok { + t.Fatalf("echo response missing headers: %s", output) + } + auth, _ := headers["authorization"].(string) + if auth != "Bearer sk-integration-test-secret" { + t.Errorf("expected injected Authorization header, got %q (full output: %s)", auth, output) + } +} + +// TestEgressProxyCredentialNotInjectedForNonMatchingDomain verifies that +// injection rules are scoped to their configured domains and are not applied +// to unrelated destinations. +// +// NOTE: credentials registered via /v1/prepare persist for the lifetime of +// the container and are never cleared between tests (this mirrors production +// behavior: a job's credentials remain registered for the sandbox's life). +// To keep this test order-independent from other tests in this file that +// also register "Authorization" injection rules for "echo-backend", this +// test uses a header name unique to itself ("X-Scoped-Test") rather than +// asserting on "Authorization", which other tests may have already caused to +// be injected for echo-backend by the time this test runs. +func TestEgressProxyCredentialNotInjectedForNonMatchingDomain(t *testing.T) { + tgt, ok := egressTarget() + if !ok { + t.Skip("SHELLCTL_EGRESS_GO_URL not set; egress proxy container not available") + } + + const sessionID = "sandbox-non-matching-domain" + prepareResp := doPutWithToken(t, tgt, egressAuthToken, "/v1/prepare", map[string]any{ + "session_id": sessionID, + "credentials": []map[string]any{ + { + "provider": "testprovider", + "name": "scoped", + "value": "sk-should-not-leak", + "inject": map[string]any{ + "type": "http-header", + "config": map[string]any{ + "name": "X-Scoped-Test", + "expr": "Bearer {{.Value}}", + "domains": []string{"some-other-host.internal"}, + }, + }, + }, + }, + }) + assertStatus(t, prepareResp, 200) + readBody(t, prepareResp) + + result := runJobWithToken(t, tgt, egressAuthToken, map[string]any{ + "script": "curl -s http://echo-backend:8080/", + "timeout": 15, + "session_id": sessionID, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + + output := result["output"].(string) + var echoed map[string]any + if err := json.Unmarshal([]byte(output), &echoed); err != nil { + t.Fatalf("failed to parse echo backend response: %v\noutput: %s", err, output) + } + headers, _ := echoed["headers"].(map[string]any) + if v, ok := headers["x-scoped-test"].(string); ok && v != "" { + t.Errorf("X-Scoped-Test header should not be injected for non-matching domain, got %q", v) + } +} + +// TestEgressProxyUpstreamChaining verifies that when SHELLCTL_EGRESSPROXY_UPSTREAM +// is configured, the credproxy correctly tunnels the outbound connection +// through the upstream forward proxy (squid) to reach the real destination, +// and that credential injection still happens (it occurs in the credproxy's +// own HTTP interceptor before the request is handed to the upstream dialer, +// so it must be unaffected by upstream chaining). +// +// NOTE: this container shares a network with echo-backend directly, so this +// test alone does NOT prove hostname passthrough through the upstream chain +// (an attempt to test that at the Docker level, by isolating this container +// from echo-backend's network, hit an unrelated Docker networking pitfall — +// see the comment in the Makefile's integration-up target). Hostname +// passthrough is instead covered reliably, without any Docker networking +// involved, by TestProxyUpstreamChainingPreservesHostname in +// internal/egressproxy/proxy_test.go. +func TestEgressProxyUpstreamChaining(t *testing.T) { + tgt, ok := egressUpstreamTarget() + if !ok { + t.Skip("SHELLCTL_EGRESS_UPSTREAM_GO_URL not set; upstream-chained egress proxy container not available") + } + + const sessionID = "sandbox-upstream-chaining" + prepareResp := doPutWithToken(t, tgt, egressUpstreamAuthToken, "/v1/prepare", map[string]any{ + "session_id": sessionID, + "credentials": []map[string]any{ + { + "provider": "testprovider", + "name": "upstreamkey", + "value": "sk-upstream-chained-secret", + "inject": map[string]any{ + "type": "http-header", + "config": map[string]any{ + "name": "Authorization", + "expr": "Bearer {{.Value}}", + "domains": []string{"echo-backend"}, + }, + }, + }, + }, + }) + assertStatus(t, prepareResp, 200) + readBody(t, prepareResp) + + // If the upstream chaining is broken (e.g. squid unreachable, CONNECT + // rejected), this curl will fail and the job's exit code will be non-zero. + result := runJobWithToken(t, tgt, egressUpstreamAuthToken, map[string]any{ + "script": "curl -sf http://echo-backend:8080/", + "timeout": 15, + "session_id": sessionID, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + + output := result["output"].(string) + var echoed map[string]any + if err := json.Unmarshal([]byte(output), &echoed); err != nil { + t.Fatalf("failed to parse echo backend response (upstream chaining likely broken): %v\noutput: %s", err, output) + } + headers, ok := echoed["headers"].(map[string]any) + if !ok { + t.Fatalf("echo response missing headers: %s", output) + } + auth, _ := headers["authorization"].(string) + if auth != "Bearer sk-upstream-chained-secret" { + t.Errorf("expected credential injection to survive upstream chaining, got %q (full output: %s)", auth, output) + } +} diff --git a/dify-agent-runtime/tests/squid-test.conf b/dify-agent-runtime/tests/squid-test.conf new file mode 100644 index 00000000000000..d8660066d14236 --- /dev/null +++ b/dify-agent-runtime/tests/squid-test.conf @@ -0,0 +1,18 @@ +# Minimal permissive Squid config used only by the egress-proxy integration +# test to verify that SHELLCTL_EGRESSPROXY_UPSTREAM chaining works end-to-end +# (the credproxy tunnels through this squid to reach the echo backend). +# +# This is NOT representative of the production squid-agent.conf.template +# ACLs; it intentionally allows everything since the goal here is to prove +# the upstream-chaining *mechanism* works, not to test ACL policy. +acl SSL_ports port 443 +acl Safe_ports port 80 +acl Safe_ports port 443 +acl Safe_ports port 8080 +acl CONNECT method CONNECT + +http_access allow CONNECT +http_access allow all + +http_port 3128 +coredump_dir /var/spool/squid diff --git a/dify-agent/pyproject.toml b/dify-agent/pyproject.toml index 36757323c6b678..2074b39e889c27 100644 --- a/dify-agent/pyproject.toml +++ b/dify-agent/pyproject.toml @@ -48,6 +48,7 @@ extraPaths = ["src", "examples/agenton", "examples/dify_agent"] [tool.pytest.ini_options] testpaths = ["tests"] python_files = ["test_*.py", "*_test.py"] +consider_namespace_packages = true markers = ["integration: requires a real external service or exercises multiple concrete adapters"] [tool.ruff] diff --git a/dify-agent/src/dify_agent/adapters/shell/protocols.py b/dify-agent/src/dify_agent/adapters/shell/protocols.py index f5f4e827cee3ec..8ee74c2c65f1df 100644 --- a/dify-agent/src/dify_agent/adapters/shell/protocols.py +++ b/dify-agent/src/dify_agent/adapters/shell/protocols.py @@ -1,7 +1,12 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Literal, Protocol +from typing import TYPE_CHECKING, Literal, Protocol + +if TYPE_CHECKING: + from collections.abc import Sequence + + from shellctl.shared.schemas import Credential @dataclass(frozen=True, slots=True) @@ -65,6 +70,8 @@ async def run( timeout: float, ) -> ShellCommandResult: ... + async def prepare(self, credentials: Sequence[Credential]) -> None: ... + async def wait( self, job_id: str, diff --git a/dify-agent/src/dify_agent/adapters/shell/shellctl.py b/dify-agent/src/dify_agent/adapters/shell/shellctl.py index 24cc9a576f8c4a..9873c63c7f2817 100644 --- a/dify-agent/src/dify_agent/adapters/shell/shellctl.py +++ b/dify-agent/src/dify_agent/adapters/shell/shellctl.py @@ -20,7 +20,12 @@ from collections.abc import Awaitable from collections.abc import Callable from dataclasses import dataclass -from typing import Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Protocol, TypeVar, cast + +if TYPE_CHECKING: + from collections.abc import Sequence + + from shellctl.shared.schemas import Credential import httpx2 as httpx from shellctl.client import ShellctlClientError @@ -58,6 +63,7 @@ _DOWNLOAD_MISSING_EXIT_CODE = 66 _WORKSPACE_PAYLOAD_BEGIN = "<<>>" _WORKSPACE_PAYLOAD_END = "<<>>" +_SESSION_ID_SANITIZER = re.compile(r"[^A-Za-z0-9_-]+") _LIST_WORKSPACE_SCRIPT = r""" import base64 @@ -220,9 +226,14 @@ async def run( *, cwd: str | None = None, env: dict[str, str] | None = None, + session_id: str | None = None, timeout: float = _DEFAULT_TIMEOUT_SECONDS, ) -> ShellctlJobResult: ... + async def prepare(self, session_id: str, credentials: list[Credential]) -> object: + """prepare the sandbox post creation. called once after the sandbox is created.""" + ... + async def wait( self, job_id: str, @@ -262,9 +273,24 @@ async def close(self) -> None: ... type ShellctlClientFactory = Callable[[], ShellctlClientProtocol] +class ShellctlSessionID(str): + """Shellctl session id identifies a logical session within the sandbox. This + is useful when multiple sessions share one shellctl container in local mode. + + When shellctl runs in isolated sandboxes, each sandbox serves only one session + so this becomes trivial. + """ + + @classmethod + def from_handle(cls, handle: str) -> ShellctlSessionID: + sanitized = _SESSION_ID_SANITIZER.sub("_", handle) + return cls(sanitized or "_") + + @dataclass(slots=True) class ShellctlCommands(ShellCommandProtocol): client: ShellctlClientProtocol + session_id: ShellctlSessionID home_dir: str | None = None workspace_dir: str | None = None @@ -283,9 +309,20 @@ async def run( ) resolved_env = _lease_env(env, home_dir=self.home_dir) return _from_job_result( - await _run_client_call(self.client.run(script, cwd=resolved_cwd, env=resolved_env, timeout=timeout)) + await _run_client_call( + self.client.run( + script, + cwd=resolved_cwd, + env=resolved_env, + session_id=self.session_id, + timeout=timeout, + ) + ) ) + async def prepare(self, credentials: Sequence[Credential]) -> None: + await _run_client_call(self.client.prepare(self.session_id, list(credentials))) + async def wait( self, job_id: str, @@ -706,5 +743,6 @@ def _shquote(value: str) -> str: "ShellctlClientProtocol", "ShellctlCommands", "ShellctlFileTransfer", + "ShellctlSessionID", "create_default_shellctl_client_factory", ] diff --git a/dify-agent/src/dify_agent/agent_stub/shell_env.py b/dify-agent/src/dify_agent/agent_stub/shell_env.py index dd81ee4e75aa80..37fa1e6d033d98 100644 --- a/dify-agent/src/dify_agent/agent_stub/shell_env.py +++ b/dify-agent/src/dify_agent/agent_stub/shell_env.py @@ -9,7 +9,9 @@ from __future__ import annotations +from dataclasses import dataclass from typing import Protocol +from urllib.parse import urlsplit from dify_agent.agent_stub._constants import AGENT_STUB_DRIVE_BASE_ENV_VAR from dify_agent.agent_stub.protocol.agent_stub import ( @@ -19,6 +21,11 @@ normalize_agent_stub_api_base_url, ) from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig +from shellctl.shared.schemas import Credential, HTTPHeaderInject, InjectPolicy + +# Placeholder pattern used to reference the JWE credential in env vars. +_JWE_CREDENTIAL_REF = "dify_agent_stub/auth_jwe" +_JWE_PLACEHOLDER = f"__secret:{_JWE_CREDENTIAL_REF}__" class ShellAgentStubTokenFactory(Protocol): @@ -27,13 +34,23 @@ class ShellAgentStubTokenFactory(Protocol): def __call__(self, execution_context: DifyExecutionContextLayerConfig, *, session_id: str | None) -> str: ... +@dataclass(frozen=True, slots=True) +class ShellAgentStubEnvResult: + """Result of building the agent stub shell environment. + + ``env`` holds the environment variables (with placeholder for JWE). + ``credentials`` holds the structured credentials for the credential proxy. + """ + + env: dict[str, str] + credentials: list[Credential] + + def build_shell_agent_stub_env( *, agent_stub_api_base_url: str | None, agent_stub_drive_ref: str | None = None, execution_context: DifyExecutionContextLayerConfig | None, - token_factory: ShellAgentStubTokenFactory | None, - session_id: str | None, ) -> dict[str, str] | None: """Build the shell-visible Agent Stub environment for one user command. @@ -41,19 +58,56 @@ def build_shell_agent_stub_env( ``dify.drive`` layer. The sandbox-local base is fixed by the Agent Stub contract and derived here at shell-run injection time. """ - if agent_stub_api_base_url is None or execution_context is None or token_factory is None: + if agent_stub_api_base_url is None or execution_context is None: return None - return { + env: dict[str, str] = { AGENT_STUB_API_BASE_URL_ENV_VAR: normalize_agent_stub_api_base_url(agent_stub_api_base_url), - AGENT_STUB_AUTH_JWE_ENV_VAR: token_factory(execution_context, session_id=session_id), + AGENT_STUB_AUTH_JWE_ENV_VAR: _JWE_PLACEHOLDER, AGENT_STUB_DRIVE_BASE_ENV_VAR: agent_stub_drive_base_for_ref(agent_stub_drive_ref), } + return env + + +def build_shell_agent_stub_credentials( + *, + agent_stub_api_base_url: str, + execution_context: DifyExecutionContextLayerConfig, + token_factory: ShellAgentStubTokenFactory, + session_id: str | None, +) -> list[Credential]: + """Build structured credentials for the JWE token with header injection. + + The returned credential instructs the sandbox credential proxy to inject + an ``Authorization: Bearer `` header on outbound HTTP requests + matching the agent stub domain. + """ + jwe = token_factory(execution_context, session_id=session_id) + parsed = urlsplit(normalize_agent_stub_api_base_url(agent_stub_api_base_url)) + domain = parsed.hostname or "" + + return [ + Credential( + provider="dify_agent_stub", + name="auth_jwe", + value=jwe, + inject=InjectPolicy( + type="http-header", + config=HTTPHeaderInject( + name="Authorization", + expr="Bearer {{.Value}}", + domains=[domain] if domain else [], + ), + ), + ), + ] __all__ = [ "AGENT_STUB_AUTH_JWE_ENV_VAR", "AGENT_STUB_DRIVE_BASE_ENV_VAR", "AGENT_STUB_API_BASE_URL_ENV_VAR", + "ShellAgentStubEnvResult", "ShellAgentStubTokenFactory", + "build_shell_agent_stub_credentials", "build_shell_agent_stub_env", ] diff --git a/dify-agent/src/dify_agent/layers/shell/layer.py b/dify-agent/src/dify_agent/layers/shell/layer.py index f987ad07371663..aa0cc4affbbc72 100644 --- a/dify-agent/src/dify_agent/layers/shell/layer.py +++ b/dify-agent/src/dify_agent/layers/shell/layer.py @@ -29,8 +29,11 @@ ShellCommandResult, ShellPromptObservation, ) -from dify_agent.agent_stub.protocol import AGENT_STUB_AUTH_JWE_ENV_VAR -from dify_agent.agent_stub.shell_env import ShellAgentStubTokenFactory, build_shell_agent_stub_env +from dify_agent.agent_stub.shell_env import ( + ShellAgentStubTokenFactory, + build_shell_agent_stub_credentials, + build_shell_agent_stub_env, +) from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig from dify_agent.layers.runtime.layer import DifyRuntimeLayer from dify_agent.layers.shell.configs import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig @@ -258,6 +261,7 @@ def tools(self) -> Sequence[PydanticAITool[object]]: @override async def on_context_create(self) -> None: + await self._prepare_credentials() bootstrap_script = _workspace_bootstrap_script(self.config) if not bootstrap_script: return @@ -510,8 +514,6 @@ def _build_shell_command_env( agent_stub_api_base_url=self.agent_stub_api_base_url, agent_stub_drive_ref=self.config.agent_stub_drive_ref, execution_context=execution_context, - token_factory=self.agent_stub_token_factory, - session_id=None, ) if agent_stub_env is None: if not require_agent_stub_env: @@ -520,25 +522,29 @@ def _build_shell_command_env( env.update(agent_stub_env) return env + async def _prepare_credentials(self) -> None: + """Register credentials with the sandbox egress proxy (once per session).""" + execution_context_layer = self.deps.execution_context + execution_context = execution_context_layer.config if execution_context_layer is not None else None + if self.agent_stub_api_base_url is None or execution_context is None or self.agent_stub_token_factory is None: + return + credentials = build_shell_agent_stub_credentials( + agent_stub_api_base_url=self.agent_stub_api_base_url, + execution_context=execution_context, + token_factory=self.agent_stub_token_factory, + session_id=None, + ) + await self._require_resource().commands.prepare(credentials) + def _redact_output(self, text: str) -> str: """Redact sensitive content from shell output before the model sees it. - Two layers of redaction are applied: - - 1. **Built-in token redaction** — the actual Agent Stub JWE token value - is always replaced with ``***``. This is unconditional and cannot be - disabled. - 2. **Pattern redaction** — regex patterns from both server-level - ``shell_redact_patterns`` and per-agent ``config.redact_patterns`` - are applied via ``re.sub`` to mask additional secrets. + Regex patterns from both server-level ``shell_redact_patterns`` and + per-agent ``config.redact_patterns`` are applied via ``re.sub`` to + mask additional secrets. """ if not text: return text - # Built-in: always redact the JWE token value. - env = self._build_shell_command_env(include_agent_stub_env=True) - jwe_value = env.get(AGENT_STUB_AUTH_JWE_ENV_VAR) - if jwe_value and len(jwe_value) > 8: - text = text.replace(jwe_value, "***") # Server-level + per-agent regex patterns. for pattern in (*self.shell_redact_patterns, *self.config.redact_patterns): text = re.sub(pattern, "***", text) diff --git a/dify-agent/src/dify_agent/runtime_backend/enterprise.py b/dify-agent/src/dify_agent/runtime_backend/enterprise.py index a87e879c8265ca..847ec73ffc1035 100644 --- a/dify-agent/src/dify_agent/runtime_backend/enterprise.py +++ b/dify-agent/src/dify_agent/runtime_backend/enterprise.py @@ -9,6 +9,7 @@ from __future__ import annotations +from dify_agent.adapters.shell.shellctl import ShellctlSessionID from dataclasses import dataclass, field import logging import shlex @@ -100,7 +101,7 @@ async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBin data_plane = await self._create_data_plane(sandbox_id) result = await run_shellctl_control_command( - ShellctlCommands(client=data_plane.client), + ShellctlCommands(client=data_plane.client, session_id=ShellctlSessionID.from_handle(sandbox_id)), "\n".join( [ "set -eu", @@ -131,7 +132,9 @@ async def acquire(self, binding_ref: str) -> RuntimeLease: data_plane: ShellctlRuntimeLease | None = None try: data_plane = await self._create_data_plane(binding_ref) - validation_commands = ShellctlCommands(client=data_plane.client) + validation_commands = ShellctlCommands( + client=data_plane.client, session_id=ShellctlSessionID.from_handle(binding_ref) + ) result = await run_shellctl_control_command( validation_commands, "\n".join( diff --git a/dify-agent/src/dify_agent/runtime_backend/shellctl.py b/dify-agent/src/dify_agent/runtime_backend/shellctl.py index 257157dfc82491..1018afa46c3e8c 100644 --- a/dify-agent/src/dify_agent/runtime_backend/shellctl.py +++ b/dify-agent/src/dify_agent/runtime_backend/shellctl.py @@ -12,6 +12,7 @@ ShellctlClientProtocol, ShellctlCommands, ShellctlFileTransfer, + ShellctlSessionID, create_default_shellctl_client_factory, ) from dify_agent.runtime_backend.protocols import FileSystem, RuntimeLayout @@ -74,6 +75,7 @@ def create_shellctl_lease( client=client, commands=ShellctlCommands( client=client, + session_id=ShellctlSessionID.from_handle(handle), home_dir=layout.home_dir, workspace_dir=layout.workspace_dir, ), diff --git a/dify-agent/src/shellctl/client/sdk.py b/dify-agent/src/shellctl/client/sdk.py index d007fff551d833..3278077c203db2 100644 --- a/dify-agent/src/shellctl/client/sdk.py +++ b/dify-agent/src/shellctl/client/sdk.py @@ -25,12 +25,14 @@ DEFAULT_TIMEOUT_SECONDS, ) from shellctl.shared.schemas import ( + Credential, DeleteJobResponse, HealthResponse, JobInfo, JobResult, JobStatusView, ListJobsResponse, + PrepareRequest, RunJobRequest, TerminalSize, ) @@ -141,19 +143,23 @@ async def run( *, cwd: str | None = None, env: dict[str, str] | None = None, + session_id: str | None = None, timeout: float = DEFAULT_TIMEOUT_SECONDS, terminal: TerminalSize | None = None, ) -> JobResult: """Create a new job and wait for initial output or completion. `cwd` and `env` preset the script's working directory and environment - overlay on the server side. + overlay on the server side. `session_id` identifies which sandbox + session's credentials (registered via `prepare()`) apply to this job's + egress traffic; it is required when the egress proxy is enabled. """ payload = RunJobRequest( script=script, cwd=cwd, env=env, + session_id=session_id, terminal=terminal, timeout=timeout, output_limit=self.output_limit, @@ -266,6 +272,23 @@ async def terminate( ) return JobStatusView.model_validate(self._decode_response(response)) + async def prepare(self, session_id: str, credentials: list[Credential]) -> dict[str, Any]: + """Register structured credentials with the sandbox credential proxy. + + Credentials are scoped strictly to `session_id`: they are persisted to + a session-specific manifest and never affect the system tier or any + other sandbox session's credentials. Pass the same `session_id` to + `run()` so the egress proxy can resolve them for that job's traffic. + """ + + payload = PrepareRequest(session_id=session_id, credentials=credentials) + response = await self._client.put( + "/v1/prepare", + json=payload.model_dump(mode="json", exclude_none=True), + headers=self._auth_headers(), + ) + return self._decode_response(response) + async def delete( self, job_id: str, diff --git a/dify-agent/src/shellctl/shared/schemas.py b/dify-agent/src/shellctl/shared/schemas.py index 0120b8b8a5f758..a618ae186c766d 100644 --- a/dify-agent/src/shellctl/shared/schemas.py +++ b/dify-agent/src/shellctl/shared/schemas.py @@ -127,17 +127,51 @@ class ErrorResponse(ShellctlModel): error: ErrorDetail +class HTTPHeaderInject(ShellctlModel): + """Inject a credential value as an HTTP request header. + + `expr` is a Go text/template string evaluated by the sandbox credential + proxy with the resolved credential value available as `{{.Value}}`, e.g. + `"Bearer {{.Value}}"` or `"{{.Value}}"`. This mirrors the wire contract of + `egressproxy.SimpleHeaderPolicy` in dify-agent-runtime. + """ + + name: str + expr: str = "" + domains: list[str] = Field(default_factory=list) + + +class InjectPolicy(ShellctlModel): + """Credential injection strategy (discriminated by type).""" + + type: str # e.g. "http-header" + config: HTTPHeaderInject | None = None + + +class Credential(ShellctlModel): + """A secret with its identity and optional injection policy.""" + + provider: str + name: str + value: str + inject: InjectPolicy | None = None + + class RunJobRequest(ShellctlModel): """HTTP request body for `POST /v1/jobs/run`. `env` augments the runner's inherited process environment instead of replacing it, so callers can preset script-local variables without losing - ambient values such as `PATH`. + ambient values such as `PATH`. Credentials are never passed here; callers + must first register them for a `session_id` via `PUT /v1/prepare`; the + egress proxy then proactively injects them into outbound HTTP requests + based on each credential's inject policy. """ script: str cwd: str | None = None env: dict[str, str] | None = None + session_id: str | None = None terminal: TerminalSize | None = None timeout: float = Field(default=DEFAULT_TIMEOUT_SECONDS, gt=0, le=MAX_WAIT_TIMEOUT_SECONDS) output_limit: int = Field(default=DEFAULT_OUTPUT_LIMIT_BYTES, ge=1, le=MAX_OUTPUT_LIMIT_BYTES) @@ -193,18 +227,42 @@ class TerminateJobRequest(ShellctlModel): grace_seconds: float = Field(default=DEFAULT_TERMINATE_GRACE_SECONDS, ge=0, le=300) +class PrepareRequest(ShellctlModel): + """HTTP request body for `PUT /v1/prepare`. + + `session_id` scopes these credentials to one sandbox session: they are + persisted server-side to a session-specific file and made visible only to + egress traffic from jobs run with the same `session_id` (see + `RunJobRequest`). They never affect the system tier or any other session. + """ + + session_id: str + credentials: list[Credential] + + +class PrepareResponse(ShellctlModel): + """Response body for `PUT /v1/prepare`.""" + + registered: int + + __all__ = [ "TERMINAL_JOB_STATUSES", + "Credential", "DeleteJobResponse", "ErrorDetail", "ErrorResponse", + "HTTPHeaderInject", "HealthResponse", + "InjectPolicy", "InputJobRequest", "JobInfo", "JobResult", "JobStatusName", "JobStatusView", "ListJobsResponse", + "PrepareRequest", + "PrepareResponse", "RunJobRequest", "ShellctlModel", "TerminalSize", diff --git a/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py b/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py index 44f897a17c0938..d9f557deeeb7a9 100644 --- a/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py +++ b/dify-agent/tests/local/dify_agent/adapters/shell/test_shellctl.py @@ -23,6 +23,7 @@ from dify_agent.adapters.shell.shellctl import ( ShellctlClientProtocol, ShellctlCommands, + ShellctlSessionID, ShellFileTransferError, ShellctlFileTransfer, ) @@ -58,6 +59,7 @@ class _RunCall: cwd: str | None env: dict[str, str] | None timeout: float + session_id: str | None = None type _RunHandler = Callable[[str, str | None, dict[str, str] | None, float], _Job] @@ -86,13 +88,17 @@ async def run( *, cwd: str | None = None, env: dict[str, str] | None = None, + session_id: str, timeout: float = 30.0, ) -> _Job: - self.run_calls.append(_RunCall(script=script, cwd=cwd, env=env, timeout=timeout)) + self.run_calls.append(_RunCall(script=script, cwd=cwd, env=env, timeout=timeout, session_id=session_id)) if self.run_handler is not None: return self.run_handler(script, cwd, env, timeout) return _Job(job_id="job", status="exited", done=True, exit_code=0) + async def prepare(self, session_id: str, credentials: object) -> object: + return {} + async def wait(self, job_id: str, *, offset: int, timeout: float = 30.0) -> _Job: self.wait_calls.append((job_id, offset, timeout)) if self.wait_handler is not None: @@ -384,7 +390,7 @@ def test_commands_forward_parameters_and_map_metadata() -> None: ) async def scenario() -> None: - commands = ShellctlCommands(_client_protocol(client)) + commands = ShellctlCommands(_client_protocol(client), session_id=ShellctlSessionID.from_handle("test-session")) run_result = await commands.run("pwd", cwd="~/workspace/abc12ff", env={"FOO": "bar"}, timeout=2.5) wait_result = await commands.wait("run-job", offset=3, timeout=4.0) read_result = await commands.read_output("run-job", offset=6) @@ -411,7 +417,9 @@ async def scenario() -> None: asyncio.run(scenario()) - assert client.run_calls == [_RunCall(script="pwd", cwd="~/workspace/abc12ff", env={"FOO": "bar"}, timeout=2.5)] + assert client.run_calls == [ + _RunCall(script="pwd", cwd="~/workspace/abc12ff", env={"FOO": "bar"}, timeout=2.5, session_id="test-session") + ] assert client.wait_calls == [ ("run-job", 3, 4.0), ("run-job", 6, 0.0), @@ -427,6 +435,7 @@ def test_commands_enforce_runtime_lease_home_and_cwd_namespace() -> None: async def scenario() -> None: commands = ShellctlCommands( _client_protocol(client), + session_id=ShellctlSessionID.from_handle("test-session"), home_dir="/homes/binding-b", workspace_dir="/workspaces/shared", ) @@ -443,12 +452,14 @@ async def scenario() -> None: cwd="/workspaces/shared", env={"HOME": "/homes/binding-b", "FOO": "bar"}, timeout=2.5, + session_id="test-session", ), _RunCall( script="pwd", cwd="/homes/binding-b/project", env={"HOME": "/homes/binding-b"}, timeout=2.5, + session_id="test-session", ), ] @@ -462,7 +473,7 @@ def test_commands_map_http_timeout_to_shell_provider_error() -> None: ) async def scenario() -> None: - commands = ShellctlCommands(_client_protocol(client)) + commands = ShellctlCommands(_client_protocol(client), session_id=ShellctlSessionID.from_handle("test-session")) with pytest.raises(ShellProviderError, match="timed out") as exc_info: await commands.run("pwd", timeout=2.5) assert exc_info.value.code == "timeout" @@ -479,7 +490,7 @@ def test_commands_map_http_request_error_to_shell_provider_error() -> None: ) async def scenario() -> None: - commands = ShellctlCommands(_client_protocol(client)) + commands = ShellctlCommands(_client_protocol(client), session_id=ShellctlSessionID.from_handle("test-session")) with pytest.raises(ShellProviderError, match="connection failed") as exc_info: await commands.wait("run-job", offset=3, timeout=4.0) assert exc_info.value.code == "request_error" @@ -495,7 +506,7 @@ def test_commands_preserve_shellctl_structured_error_fields() -> None: ) async def scenario() -> None: - commands = ShellctlCommands(_client_protocol(client)) + commands = ShellctlCommands(_client_protocol(client), session_id=ShellctlSessionID.from_handle("test-session")) with pytest.raises(ShellProviderError, match="sandbox expired") as exc_info: await commands.run("pwd", timeout=2.5) assert exc_info.value.status_code == 404 @@ -569,7 +580,7 @@ async def delete(self, job_id, *, force=False, grace_seconds=None): client = DeleteTimeoutClient() async def scenario() -> None: - commands = ShellctlCommands(_client_protocol(client)) + commands = ShellctlCommands(_client_protocol(client), session_id=ShellctlSessionID.from_handle("test-session")) with pytest.raises(ShellProviderError, match="delete timed out") as exc_info: await commands.delete("run-job", force=True, grace_seconds=2.0) assert exc_info.value.code == "timeout" @@ -590,7 +601,7 @@ async def delete(self, job_id, *, force=False, grace_seconds=None): client = DeleteRequestErrorClient() async def scenario() -> None: - commands = ShellctlCommands(_client_protocol(client)) + commands = ShellctlCommands(_client_protocol(client), session_id=ShellctlSessionID.from_handle("test-session")) with pytest.raises(ShellProviderError, match="delete connection failed") as exc_info: await commands.delete("run-job", force=True, grace_seconds=2.0) assert exc_info.value.code == "request_error" diff --git a/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py b/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py index a358d0725ce5dc..0c792538eb21a4 100644 --- a/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py @@ -194,6 +194,9 @@ async def run(self, script: str, *, cwd: str | None = None, env: dict[str, str] raise AssertionError("Unexpected run() call") return self.run_handler(script, cwd, env, timeout) + async def prepare(self, credentials: object) -> None: + pass + async def wait(self, job_id: str, *, offset: int, timeout: float): self.wait_calls.append(WaitCall(job_id=job_id, offset=offset, timeout=timeout)) if self.wait_handler is None: @@ -883,7 +886,9 @@ def run_handler( _timeout: float, ) -> ShellCommandResult: assert env is not None - assert env["DIFY_AGENT_STUB_AUTH_JWE"] == "stub-token" + # The env carries a placeholder, not the real token; the egress proxy + # resolves it at request time. + assert env["DIFY_AGENT_STUB_AUTH_JWE"] == "__secret:dify_agent_stub/auth_jwe__" return _command_result("remote-job", status="exited", done=True, exit_code=0) layer, _provider = _layer(commands=FakeCommands(run_handler=run_handler)) @@ -893,10 +898,15 @@ def run_handler( async def scenario() -> None: async with layer.resource_context(): + # on_context_create triggers _prepare_credentials which calls + # token_factory via build_shell_agent_stub_credentials. + await layer.on_context_create() _ = await layer.run_remote_script_complete("true", inject_agent_stub_env=True) asyncio.run(scenario()) + # token_factory is invoked once via the credentials path (build_shell_agent_stub_credentials) + # with session_id=None so the workspace session identity is not embedded in the token. assert seen_session_ids == [None] @@ -1214,39 +1224,6 @@ def _layer_with_redaction( return layer, provider -def test_redact_output_replaces_jwe_token_value() -> None: - """The JWE token value should always be redacted from shell output.""" - token = "eyJhbGciOiJkaXIiLCJlbmMiOiJBMjU2R0NNIn0.super-secret-token-12345" - - def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: - return _command_result( - "job-1", - status="exited", - done=True, - exit_code=0, - output=f"DIFY_AGENT_STUB_AUTH_JWE={token}\n", - offset=100, - ) - - commands = FakeCommands( - run_handler=run_handler, - tail_handler=lambda _: _command_result("job-1", done=True, exit_code=0, status="exited", offset=100), - ) - layer, _provider = _layer_with_redaction(commands=commands, token_value=token) - _bind_execution_context(layer) - layer.runtime_state = _runtime_state() - tools = {tool.name: tool for tool in layer.tools} - - async def scenario() -> None: - async with layer.resource_context(): - result = await tools["shell_run"].function_schema.call({"script": "env"}, None) # pyright: ignore[reportArgumentType] - _, output = _parse_tagged_observation(result) - assert token not in output - assert "***" in output - - asyncio.run(scenario()) - - def test_redact_output_applies_server_level_patterns() -> None: """Server-level regex patterns from env var should redact matching content.""" @@ -1313,35 +1290,3 @@ async def scenario() -> None: assert "token: ***" in output asyncio.run(scenario()) - - -def test_redact_output_skips_short_jwe_values() -> None: - """JWE values ≤8 chars should not be redacted to avoid false positives.""" - - def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: - return _command_result( - "job-1", - status="exited", - done=True, - exit_code=0, - output="short\n", - offset=6, - ) - - commands = FakeCommands( - run_handler=run_handler, - tail_handler=lambda _: _command_result("job-1", done=True, exit_code=0, status="exited", offset=6), - ) - # Token value is short — should NOT be redacted even if it appears in output. - layer, _provider = _layer_with_redaction(commands=commands, token_value="short") - _bind_execution_context(layer) - layer.runtime_state = _runtime_state() - tools = {tool.name: tool for tool in layer.tools} - - async def scenario() -> None: - async with layer.resource_context(): - result = await tools["shell_run"].function_schema.call({"script": "echo hi"}, None) # pyright: ignore[reportArgumentType] - _, output = _parse_tagged_observation(result) - assert "short" in output - - asyncio.run(scenario()) diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py index a2f8190e32fc11..4a28d7355d166e 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py @@ -2,11 +2,14 @@ from dataclasses import dataclass, field import shlex -from typing import Mapping +from typing import TYPE_CHECKING, Mapping import pytest from shellctl.shared import DeleteJobResponse, JobResult, JobStatusName, JobStatusView +if TYPE_CHECKING: + from shellctl.shared.schemas import Credential + from dify_agent.runtime_backend import ( BindingCreateError, BindingDestroyError, @@ -22,6 +25,7 @@ class _RunCall: commands: tuple[tuple[str, ...], ...] cwd: str | None env: Mapping[str, str] | None + session_id: str | None = None @dataclass(slots=True) @@ -38,14 +42,15 @@ async def run( script: str, *, cwd: str | None = None, - env: Mapping[str, str] | None = None, + env: dict[str, str] | None = None, + session_id: str | None = None, timeout: float = 10.0, ) -> JobResult: del timeout commands = tuple( tuple(shlex.split(line)) for line in script.splitlines() if line.strip() and line.strip() != "set -eu" ) - self.runs.append(_RunCall(commands=commands, cwd=cwd, env=env)) + self.runs.append(_RunCall(commands=commands, cwd=cwd, env=env, session_id=session_id)) return JobResult( job_id=f"job-{len(self.runs)}", status=JobStatusName.EXITED, @@ -57,6 +62,9 @@ async def run( truncated=False, ) + async def prepare(self, session_id: str, credentials: list[Credential]) -> object: + return {} + async def wait(self, job_id: str, *, offset: int, timeout: float = 10.0) -> JobResult: raise AssertionError((job_id, offset, timeout)) diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py index 9d07ccfb027f74..2855df22a37e85 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_shellctl_backend.py @@ -6,7 +6,7 @@ import pytest from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellCommandStatus -from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol +from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol, ShellctlSessionID from dify_agent.runtime_backend.protocols import RuntimeLayout from dify_agent.runtime_backend.shellctl import ( create_owned_shellctl_lease, @@ -55,6 +55,9 @@ async def run( del script, cwd, env, timeout return self.initial + async def prepare(self, credentials: object) -> None: + pass + async def wait(self, job_id: str, *, offset: int, timeout: float) -> ShellCommandResult: del job_id, offset, timeout if self.wait_error is not None: @@ -98,6 +101,25 @@ def _result(*, done: bool = True) -> ShellCommandResult: ) +@pytest.mark.parametrize( + ("handle", "want"), + [ + ("sandbox-1", "sandbox-1"), + ("binding-id:workspace-id", "binding-id_workspace-id"), + ("agent.stub-run:workspace.a", "agent_stub-run_workspace_a"), + ("", "_"), + ], +) +def test_session_id_from_handle_sanitizes_disallowed_characters(handle: str, want: str) -> None: + # The shellctl runtime restricts session_id to [A-Za-z0-9_-]{1,128} and + # additionally treats ':' as the Basic-Auth user:password separator when + # the id is embedded in the egress proxy's HTTP_PROXY/HTTPS_PROXY URL, so + # handles like local binding refs ("binding_id:workspace_id") must be + # sanitized before use as a session_id. + assert ShellctlSessionID.from_handle(handle) == want + assert str(ShellctlSessionID.from_handle(handle)) == want + + @pytest.mark.anyio async def test_owned_transport_is_closed_exactly_once() -> None: client = _FakeClient() diff --git a/docker/docker-compose-template.yaml b/docker/docker-compose-template.yaml index bb4a3c76efd790..fcc837cecc67fa 100644 --- a/docker/docker-compose-template.yaml +++ b/docker/docker-compose-template.yaml @@ -543,9 +543,9 @@ services: # for shellctl, and it can reach agent_backend directly) and # `local_sandbox_proxy_network` (so its egress is forced through # agent_ssrf_proxy). - # All non-agent_backend/localhost traffic goes through the Squid forward proxy - # on port 3128, which only allows agent_backend /agent-stub/ and the Dify API - # /files/* endpoints (see ssrf_proxy/squid-agent.conf.template). + # The in-process egress proxy (127.0.0.1:18080) only handles credential + # injection; network-level egress restriction (allowlist/SSRF) is enforced + # by agent_ssrf_proxy (Squid), which the egress proxy chains to as upstream. local_sandbox: image: langgenius/dify-agent-local-sandbox:1.16.1 restart: always @@ -554,9 +554,14 @@ services: required: false environment: - SHELLCTL_AUTH_TOKEN=${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}} - - HTTP_PROXY=http://agent_ssrf_proxy:3128 - - HTTPS_PROXY=http://agent_ssrf_proxy:3128 + - SHELLCTL_EGRESSPROXY_ENABLED=${DIFY_AGENT_EGRESSPROXY_ENABLED:-true} + - SHELLCTL_EGRESSPROXY_UPSTREAM=http://agent_ssrf_proxy:${SSRF_HTTP_PORT:-3128} + - SHELLCTL_EGRESSPROXY_SYSTEM_CREDENTIALS_DIR=/etc/shellctl/credentials + - HTTP_PROXY=http://127.0.0.1:18080 + - HTTPS_PROXY=http://127.0.0.1:18080 - NO_PROXY=localhost,127.0.0.1 + volumes: + - ./volumes/local_sandbox/credentials:/etc/shellctl/credentials:ro healthcheck: test: ["CMD", "curl", "-f", "http://localhost:5004/healthz"] interval: 30s diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 301c03c45acf83..cbd34fc05758f1 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -549,9 +549,9 @@ services: # for shellctl, and it can reach agent_backend directly) and # `local_sandbox_proxy_network` (so its egress is forced through # agent_ssrf_proxy). - # All non-agent_backend/localhost traffic goes through the Squid forward proxy - # on port 3128, which only allows agent_backend /agent-stub/ and the Dify API - # /files/* endpoints (see ssrf_proxy/squid-agent.conf.template). + # The in-process egress proxy (127.0.0.1:18080) only handles credential + # injection; network-level egress restriction (allowlist/SSRF) is enforced + # by agent_ssrf_proxy (Squid), which the egress proxy chains to as upstream. local_sandbox: image: langgenius/dify-agent-local-sandbox:1.16.1 restart: always @@ -560,9 +560,14 @@ services: required: false environment: - SHELLCTL_AUTH_TOKEN=${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}} - - HTTP_PROXY=http://agent_ssrf_proxy:3128 - - HTTPS_PROXY=http://agent_ssrf_proxy:3128 + - SHELLCTL_EGRESSPROXY_ENABLED=${DIFY_AGENT_EGRESSPROXY_ENABLED:-true} + - SHELLCTL_EGRESSPROXY_UPSTREAM=http://agent_ssrf_proxy:${SSRF_HTTP_PORT:-3128} + - SHELLCTL_EGRESSPROXY_SYSTEM_CREDENTIALS_DIR=/etc/shellctl/credentials + - HTTP_PROXY=http://127.0.0.1:18080 + - HTTPS_PROXY=http://127.0.0.1:18080 - NO_PROXY=localhost,127.0.0.1 + volumes: + - ./volumes/local_sandbox/credentials:/etc/shellctl/credentials:ro healthcheck: test: ["CMD", "curl", "-f", "http://localhost:5004/healthz"] interval: 30s diff --git a/docker/volumes/local_sandbox/credentials/.gitignore b/docker/volumes/local_sandbox/credentials/.gitignore new file mode 100644 index 00000000000000..d8916012c56aaa --- /dev/null +++ b/docker/volumes/local_sandbox/credentials/.gitignore @@ -0,0 +1,4 @@ +# Ignore actual credential files (real secrets) +*.cred.yaml +*.cred.yml +*.cred.json diff --git a/docker/volumes/local_sandbox/credentials/README.md b/docker/volumes/local_sandbox/credentials/README.md new file mode 100644 index 00000000000000..a712be5c22ece7 --- /dev/null +++ b/docker/volumes/local_sandbox/credentials/README.md @@ -0,0 +1,88 @@ +# System Credentials Directory + +All `.yaml`, `.yml`, and `.json` files in this directory are loaded at startup +and merged into the egress proxy's **system credential tier**. Files are loaded +in alphabetical order; later files override earlier ones on `provider/name` +conflicts. + +Files matching `*.cred.yaml`, `*.cred.yml`, and `*.cred.json` are gitignored +(see `.gitignore`) to prevent accidental commits of real secrets. + +## Example: simple header injection (API key) + +Create a file like `tavily.cred.yaml` (gitignored): + +```yaml +credentials: + - provider: tavily + name: api_key + value: tvly-dev-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx + env_name: TAVILY_API_KEY + inject: + type: http-header + config: + name: Authorization + expr: "Bearer {{.Value}}" + domains: + - api.tavily.com +``` + +## Example: AWS S3 with SigV4 re-signing + +```yaml +credentials: + - provider: aws + name: s3_prod + value: # structured value (object, not string) + access_key_id: AKIAIOSFODNN7EXAMPLE + secret_access_key: wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY + session_token: "" # optional, only for temporary credentials + env_names: # one credential → multiple env vars + - AWS_ACCESS_KEY_ID + - AWS_SECRET_ACCESS_KEY + - AWS_SESSION_TOKEN + inject: + type: aws-sigv4 + config: + service: s3 # defaults to "s3" if omitted + # region: us-east-1 # omit to auto-extract from hostname + domains: + - "*.amazonaws.com" +``` + +The proxy strips any client-supplied AWS auth headers and re-signs with the +real credentials, so both `curl` (no signature) and `aws cli` (placeholder-based +fake signature from env vars) work transparently. + +## Example: Cloudflare R2 (S3-compatible) + +```yaml +credentials: + - provider: cloudflare + name: r2_prod + value: + access_key_id: + secret_access_key: + env_names: + - AWS_ACCESS_KEY_ID + - AWS_SECRET_ACCESS_KEY + inject: + type: aws-sigv4 + config: + region: auto # R2 is region-less; must set explicitly + service: s3 + domains: + - "*.r2.cloudflarestorage.com" +``` + +## Field reference + +| Field | Description | +|---|---| +| `provider` | Credential provider namespace (e.g. `tavily`, `aws`) | +| `name` | Credential name within the provider (e.g. `api_key`, `s3_prod`) | +| `value` | The secret value — a string for simple credentials, or an object for structured credentials (e.g. AWS) | +| `env_name` | Single env var name exposed to jobs as a `__secret:provider/name__` placeholder (optional; auto-derived as `PROVIDER_NAME` uppercased if omitted) | +| `env_names` | Multiple env var names, all pointing to the same `__secret:provider/name__` placeholder (for structured credentials like AWS that need several standard env vars) | +| `inject.type` | Injection policy: `http-header` or `aws-sigv4` | +| `inject.config` | Type-specific config payload (see examples above) |