From 56fde5b20bc9baf8f234bb4cd03c796e4ff35e24 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 14 Aug 2026 13:34:04 +0800 Subject: [PATCH] feat(zookeeper): replace Java agent with native Go implementation --- .github/scripts/bump-agent-versions.mjs | 7 +- .github/scripts/bump-agent-versions.test.mjs | 19 + .../scripts/reuse-agent-release-assets.mjs | 2 +- .github/workflows/agents-release.yml | 51 +- .github/workflows/ci.yml | 88 ++ agents/README.md | 11 +- agents/README.zh-CN.md | 11 +- agents/drivers/zookeeper/agent_test.go | 336 ++++++ agents/drivers/zookeeper/build.gradle | 13 - agents/drivers/zookeeper/connection.go | 563 ++++++++++ agents/drivers/zookeeper/go.mod | 26 + agents/drivers/zookeeper/go.sum | 50 + agents/drivers/zookeeper/integration_test.go | 105 ++ agents/drivers/zookeeper/main.go | 128 +++ agents/drivers/zookeeper/operations.go | 515 ++++++++++ agents/drivers/zookeeper/sasl.go | 194 ++++ agents/drivers/zookeeper/sasl_test.go | 156 +++ .../dbx/agent/zookeeper/ZooKeeperAgent.java | 971 ------------------ .../agent/zookeeper/ZooKeeperAgentTest.java | 924 ----------------- agents/go-common/gosasl/sasl.go | 10 +- agents/go-common/gosasl/sasl_test.go | 30 + .../scripts/driver_release_packages_test.py | 16 + agents/scripts/validate_agents.py | 1 + agents/scripts/validate_agents_test.py | 3 +- agents/scripts/version_agent_artifacts.py | 2 +- agents/settings.gradle | 2 +- 26 files changed, 2306 insertions(+), 1928 deletions(-) create mode 100644 agents/drivers/zookeeper/agent_test.go delete mode 100644 agents/drivers/zookeeper/build.gradle create mode 100644 agents/drivers/zookeeper/connection.go create mode 100644 agents/drivers/zookeeper/go.mod create mode 100644 agents/drivers/zookeeper/go.sum create mode 100644 agents/drivers/zookeeper/integration_test.go create mode 100644 agents/drivers/zookeeper/main.go create mode 100644 agents/drivers/zookeeper/operations.go create mode 100644 agents/drivers/zookeeper/sasl.go create mode 100644 agents/drivers/zookeeper/sasl_test.go delete mode 100644 agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java delete mode 100644 agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java diff --git a/.github/scripts/bump-agent-versions.mjs b/.github/scripts/bump-agent-versions.mjs index 44358a6967..7fb5fef8f9 100644 --- a/.github/scripts/bump-agent-versions.mjs +++ b/.github/scripts/bump-agent-versions.mjs @@ -59,15 +59,20 @@ const nativeDriverDirectories = { vastbase: "vastbase-go", rabbitmq: "rabbitmq", rocketmq: "rocketmq", + zookeeper: "zookeeper", tdengine: "tdengine", }; -const nativeDriverModules = new Set(["cassandra", "duckdb", "hive", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "rabbitmq", "rocketmq", "tdengine"]); +const nativeDriverModules = new Set(["cassandra", "duckdb", "hive", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "rabbitmq", "rocketmq", "zookeeper", "tdengine"]); const nativeDriverSharedPaths = { hive: [ "agents/go-common/go-gssapi", "agents/go-common/gohive", "agents/go-common/gosasl", ], + zookeeper: [ + "agents/go-common/go-gssapi", + "agents/go-common/gosasl", + ], }; function resolveAgentModule(moduleName, { legacyStandaloneModules, moduleExists, readModuleFile }) { diff --git a/.github/scripts/bump-agent-versions.test.mjs b/.github/scripts/bump-agent-versions.test.mjs index 99ea4c20f4..ea6e6e06cc 100644 --- a/.github/scripts/bump-agent-versions.test.mjs +++ b/.github/scripts/bump-agent-versions.test.mjs @@ -71,6 +71,25 @@ test("bumps the native RocketMQ agent from its Go directory", () => { assert.deepEqual(result.nativeModules, ["rocketmq"]); }); +test("bumps ZooKeeper from native and shared SASL source directories", () => { + for (const changedFile of [ + "agents/drivers/zookeeper/main.go", + "agents/go-common/gosasl/sasl.go", + "agents/go-common/go-gssapi/krb5/krb5.go", + ]) { + const result = evaluateAgentVersionBump({ + versions: { zookeeper: "0.1.0" }, + changedFiles: [changedFile], + moduleExists: (path) => path === "agents/drivers/zookeeper", + readModuleFile: () => "", + }); + + assert.equal(result.versions.zookeeper, "0.1.1"); + assert.deepEqual(result.javaModules, []); + assert.deepEqual(result.nativeModules, ["zookeeper"]); + } +}); + test("bumps the native Vastbase agent from its independent Go directory", () => { const result = evaluateAgentVersionBump({ versions: { vastbase: "0.1.37" }, diff --git a/.github/scripts/reuse-agent-release-assets.mjs b/.github/scripts/reuse-agent-release-assets.mjs index 21f72bde29..61f9159e24 100644 --- a/.github/scripts/reuse-agent-release-assets.mjs +++ b/.github/scripts/reuse-agent-release-assets.mjs @@ -15,7 +15,7 @@ import { basename, join } from "node:path"; import { tmpdir } from "node:os"; const REGISTRY_ASSET = "agent-registry.json"; -const NATIVE_MODULES = new Set(["duckdb", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "rabbitmq", "rocketmq", "tdengine"]); +const NATIVE_MODULES = new Set(["duckdb", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "rabbitmq", "rocketmq", "zookeeper", "tdengine"]); const PLATFORMS = [ "macos-aarch64", "macos-x64", diff --git a/.github/workflows/agents-release.yml b/.github/workflows/agents-release.yml index 86d94837b4..b655d8e7a2 100644 --- a/.github/workflows/agents-release.yml +++ b/.github/workflows/agents-release.yml @@ -329,6 +329,46 @@ jobs: name: rocketmq-native path: "release-native/dbx-agent-rocketmq-*" + build-zookeeper-native: + needs: [bump-versions] + if: ${{ contains(fromJSON(needs.bump-versions.outputs.native_modules), 'zookeeper') }} + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version: "1.25.x" + - name: Test ZooKeeper native agent + working-directory: agents/drivers/zookeeper + run: go test -race ./... + - name: Cross-compile ZooKeeper native agent + shell: bash + run: | + mkdir -p release-native + cd agents/drivers/zookeeper + declare -A TARGETS=( + ["macos-aarch64"]="darwin/arm64" + ["macos-x64"]="darwin/amd64" + ["linux-aarch64"]="linux/arm64" + ["linux-x64"]="linux/amd64" + ["windows-aarch64"]="windows/arm64" + ["windows-x64"]="windows/amd64" + ) + for platform in "${!TARGETS[@]}"; do + IFS=/ read -r goos goarch <<< "${TARGETS[$platform]}" + output="../../../release-native/dbx-agent-zookeeper-${platform}" + if [[ "$goos" == "windows" ]]; then + output="${output}.exe" + fi + echo "Building $platform ($goos/$goarch)" + CGO_ENABLED=0 GOOS="$goos" GOARCH="$goarch" go build -trimpath -ldflags="-s -w" -o "$output" . + done + ls -lh ../../../release-native + - uses: actions/upload-artifact@v4 + with: + name: zookeeper-native + path: "release-native/dbx-agent-zookeeper-*" + build-cassandra-native: needs: [bump-versions] if: ${{ contains(fromJSON(needs.bump-versions.outputs.native_modules), 'cassandra') }} @@ -908,7 +948,7 @@ jobs: retention-days: 1 release: - needs: [bump-versions, commit-versions, build-agents, build-oracle-native, build-xugu-native, build-rabbitmq-native, build-rocketmq-native, build-cassandra-native, build-hive-native, build-kingbase-native, build-vastbase-native, build-neo4j-native, build-iotdb-native, build-duckdb-native, build-tdengine-native, build-jre, reuse-previous-assets] + needs: [bump-versions, commit-versions, build-agents, build-oracle-native, build-xugu-native, build-rabbitmq-native, build-rocketmq-native, build-zookeeper-native, build-cassandra-native, build-hive-native, build-kingbase-native, build-vastbase-native, build-neo4j-native, build-iotdb-native, build-duckdb-native, build-tdengine-native, build-jre, reuse-previous-assets] if: ${{ always() && !contains(needs.*.result, 'failure') && !contains(needs.*.result, 'cancelled') }} runs-on: ubuntu-latest steps: @@ -1031,6 +1071,7 @@ jobs: xugu) echo "虚谷 XuguDB" ;; rabbitmq) echo "RabbitMQ" ;; rocketmq) echo "Apache RocketMQ" ;; + zookeeper) echo "Apache ZooKeeper" ;; cassandra) echo "Apache Cassandra" ;; hive) echo "Apache Hive" ;; neo4j) echo "Neo4j" ;; @@ -1099,7 +1140,7 @@ jobs: [ -n "$DRIVERS" ] && DRIVERS="${DRIVERS},"$'\n' DRIVERS="${DRIVERS}$(generate_jar_entry "$name" "$label" "$f" "$jre_key" "$version" "$external_driver" "$native_json")" done - for name in oracle xugu kingbase vastbase neo4j iotdb duckdb rabbitmq rocketmq cassandra hive tdengine; do + for name in oracle xugu kingbase vastbase neo4j iotdb duckdb rabbitmq rocketmq zookeeper cassandra hive tdengine; do version=$(get_module_version "$name") [ -f "release/dbx-agent-${name}-${version}.jar" ] && continue native_json=$(generate_native_platforms "$name" "$version") @@ -1170,6 +1211,7 @@ jobs: xugu) echo "虚谷 XuguDB" ;; rabbitmq) echo "RabbitMQ" ;; rocketmq) echo "Apache RocketMQ" ;; + zookeeper) echo "Apache ZooKeeper" ;; cassandra) echo "Apache Cassandra" ;; hive) echo "Apache Hive" ;; neo4j) echo "Neo4j" ;; @@ -1219,6 +1261,11 @@ jobs: "agents/go-common/gohive/" "agents/go-common/gosasl/" ) + elif [ "$name" = "zookeeper" ]; then + LOG_PATHS+=( + "agents/go-common/go-gssapi/" + "agents/go-common/gosasl/" + ) fi while IFS= read -r line; do [ -n "$line" ] && NOTES="${NOTES}- ${line}"$'\n' diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0b4bb879b6..fc57b200e8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -744,6 +744,10 @@ jobs: run: go test -race ./... working-directory: agents/drivers/rocketmq + - name: ZooKeeper native agent tests + run: go test -race ./... + working-directory: agents/drivers/zookeeper + - name: Cassandra native agent tests run: go test ./... working-directory: agents/drivers/cassandra-go @@ -780,6 +784,90 @@ jobs: run: CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-rocketmq-linux-x64 . working-directory: agents/drivers/rocketmq + - name: ZooKeeper native agent cross-builds + shell: bash + working-directory: agents/drivers/zookeeper + run: | + set -euo pipefail + for target in darwin/arm64 darwin/amd64 linux/arm64 linux/amd64 windows/arm64 windows/amd64; do + IFS=/ read -r goos goarch <<< "$target" + output="/tmp/dbx-agent-zookeeper-${goos}-${goarch}" + [ "$goos" = windows ] && output="${output}.exe" + CGO_ENABLED=0 GOOS="$goos" GOARCH="$goarch" go build -trimpath -ldflags="-s -w" -o "$output" . + done + + - name: ZooKeeper native agent integration tests + shell: bash + working-directory: agents/drivers/zookeeper + run: | + set -euo pipefail + for version in 3.4.14 3.5.5 3.7.0 3.9.5; do + name="dbx-zookeeper-${version//./-}" + docker rm -fv "$name" >/dev/null 2>&1 || true + docker run -d --name "$name" -p 2181:2181 "zookeeper:$version" + cleanup() { + docker rm -fv "$name" >/dev/null 2>&1 || true + } + trap cleanup EXIT + ready=false + for _ in $(seq 1 90); do + if timeout 1 bash -c '/dev/null 2>&1; then + ready=true + break + fi + sleep 2 + done + if [ "$ready" != true ]; then + docker logs "$name" + exit 1 + fi + DBX_ZOOKEEPER_TEST_CONNECT_STRING=127.0.0.1:2181 \ + go test -run '^TestZooKeeperIntegration$' -count=1 ./... + cleanup + trap - EXIT + done + + - name: ZooKeeper SASL DIGEST-MD5 integration test + shell: bash + working-directory: agents/drivers/zookeeper + run: | + set -euo pipefail + name="dbx-zookeeper-sasl-3-7-0" + jaas_file=$(mktemp) + cat > "$jaas_file" <<'EOF' + Server { + org.apache.zookeeper.server.auth.DigestLoginModule required + user_dbx="dbx-secret"; + }; + EOF + docker rm -fv "$name" >/dev/null 2>&1 || true + docker run -d --name "$name" -p 2181:2181 \ + -v "$jaas_file:/conf/server-jaas.conf:ro" \ + -e ZOO_CFG_EXTRA="authProvider.1=org.apache.zookeeper.server.auth.SASLAuthenticationProvider" \ + -e JVMFLAGS="-Djava.security.auth.login.config=/conf/server-jaas.conf -Dzookeeper.sasl.serverconfig=Server -Dzookeeper.sessionRequireClientSASLAuth=true" \ + zookeeper:3.7.0 + cleanup() { + docker rm -fv "$name" >/dev/null 2>&1 || true + } + trap cleanup EXIT + ready=false + for _ in $(seq 1 90); do + if timeout 1 bash -c '/dev/null 2>&1; then + ready=true + break + fi + sleep 2 + done + if [ "$ready" != true ]; then + docker logs "$name" + exit 1 + fi + DBX_ZOOKEEPER_TEST_CONNECT_STRING=127.0.0.1:2181 \ + DBX_ZOOKEEPER_TEST_AUTH_SCHEME=sasl_digest \ + DBX_ZOOKEEPER_TEST_USERNAME=dbx \ + DBX_ZOOKEEPER_TEST_PASSWORD=dbx-secret \ + go test -run '^TestZooKeeperIntegration$' -count=1 ./... + - name: RocketMQ native agent integration tests shell: bash working-directory: agents/drivers/rocketmq diff --git a/agents/README.md b/agents/README.md index fbe698a3ee..82b5b7b7e2 100644 --- a/agents/README.md +++ b/agents/README.md @@ -43,14 +43,14 @@ Each agent runs as a standalone process and communicates with DBX via stdin/stdo | xugu | 虚谷 XuguDB | XuguDB Go native agent | | iotdb | Apache IoTDB | Apache IoTDB Go Client native agent | | etcd | etcd | jetcd | -| zookeeper | Apache ZooKeeper | Apache Curator | +| zookeeper | Apache ZooKeeper | go-zookeeper native agent | | rabbitmq | RabbitMQ | amqp091-go native agent | | rocketmq | Apache RocketMQ | rocketmq-admin-go native agent | ## Multi-JRE Support -Most Java agents target JRE 21. Native agents, such as `cassandra`, `duckdb`, `hive`, `iotdb`, `oracle`, `kingbase`, `tdengine`, `xugu`, `rabbitmq`, and `rocketmq`, do not require a JRE. DBX downloads and manages the JRE 21 installation automatically for Java agents. +Most Java agents target JRE 21. Native agents, such as `cassandra`, `duckdb`, `hive`, `iotdb`, `oracle`, `kingbase`, `tdengine`, `xugu`, `rabbitmq`, `rocketmq`, and `zookeeper`, do not require a JRE. DBX downloads and manages the JRE 21 installation automatically for Java agents. ## JDBC Connection Pooling @@ -76,7 +76,7 @@ Set `DBX_AGENT_JDBC_POOL_ENABLED=false` for a runtime-level compatibility fallba For new agents, prefer a **native (Go or Rust) driver** over a Java/JDBC agent whenever a mature, license-compatible native driver is available. Native agents ship as a single self-contained executable with no JRE, which significantly reduces memory footprint and startup time — the JVM baseline that every Java agent pays even when idle is avoided entirely. -- **Native (C++/Go/Rust)** — preferred when a usable native driver exists. See `drivers/cassandra-go` (Apache cassandra-gocql-driver), `drivers/duckdb`, `drivers/hive-go` (native HS2), `drivers/iotdb` (Apache IoTDB Go Client), `drivers/oracle-go` (go-ora), `drivers/kingbase-go` (gokb), `drivers/vastbase-go` (openGauss connector), `drivers/tdengine` (taos-connector-rust), `drivers/xugu`, `drivers/rabbitmq` (amqp091-go), and `drivers/rocketmq` (rocketmq-admin-go) as reference implementations. No JRE download or management is needed. +- **Native (C++/Go/Rust)** — preferred when a usable native driver exists. See `drivers/cassandra-go` (Apache cassandra-gocql-driver), `drivers/duckdb`, `drivers/hive-go` (native HS2), `drivers/iotdb` (Apache IoTDB Go Client), `drivers/oracle-go` (go-ora), `drivers/kingbase-go` (gokb), `drivers/vastbase-go` (openGauss connector), `drivers/tdengine` (taos-connector-rust), `drivers/xugu`, `drivers/rabbitmq` (amqp091-go), `drivers/rocketmq` (rocketmq-admin-go), and `drivers/zookeeper` (go-zookeeper) as reference implementations. No JRE download or management is needed. - **Java/JDBC** — the default fallback when only a JDBC driver exists for the database, or when the native driver is immature or unmaintained. Most agents still fall in this category. Native agents implement the same JSON-RPC contract and `versions.json` registration as Java agents; they ship an `agent` executable instead of `agent.jar`. If both native and Java source implementations exist for the same database, publish only the native artifact unless the Java variant has a separately registered compatibility profile, such as `oracle-legacy` / `oracle-10g`. @@ -97,9 +97,10 @@ Requires JDK 21 (Gradle toolchain auto-downloads if needed). (cd drivers/xugu && go build -o agent .) (cd drivers/rabbitmq && go build -o agent .) (cd drivers/rocketmq && go build -o agent .) +(cd drivers/zookeeper && go build -o agent .) ``` -Output JARs are in `drivers/{module}/build/libs/`. Native agents build from `drivers/cassandra-go`, `drivers/duckdb`, `drivers/hive-go`, `drivers/iotdb`, `drivers/oracle-go`, `drivers/kingbase-go`, `drivers/vastbase-go`, `drivers/tdengine`, `drivers/xugu`, `drivers/rabbitmq`, and `drivers/rocketmq`. +Output JARs are in `drivers/{module}/build/libs/`. Native agents build from `drivers/cassandra-go`, `drivers/duckdb`, `drivers/hive-go`, `drivers/iotdb`, `drivers/oracle-go`, `drivers/kingbase-go`, `drivers/vastbase-go`, `drivers/tdengine`, `drivers/xugu`, `drivers/rabbitmq`, `drivers/rocketmq`, and `drivers/zookeeper`. ### Local DBX Runtime Test @@ -113,7 +114,7 @@ cp agents/drivers//build/libs/*-all.jar ~/.dbx/agents/drivers/ Restart DBX or disconnect and reconnect the database so the new agent process loads the replacement JAR. -Native agents such as `cassandra`, `hive`, `iotdb`, `oracle`, `kingbase`, `tdengine`, `xugu`, `rabbitmq`, and `rocketmq` use an `agent` executable instead of `agent.jar`. TDengine builds `target/release/dbx-tdengine-driver` from `drivers/tdengine/Cargo.toml`. +Native agents such as `cassandra`, `hive`, `iotdb`, `oracle`, `kingbase`, `tdengine`, `xugu`, `rabbitmq`, `rocketmq`, and `zookeeper` use an `agent` executable instead of `agent.jar`. TDengine builds `target/release/dbx-tdengine-driver` from `drivers/tdengine/Cargo.toml`. ## Versioning diff --git a/agents/README.zh-CN.md b/agents/README.zh-CN.md index 1321edf436..da2f15944a 100644 --- a/agents/README.zh-CN.md +++ b/agents/README.zh-CN.md @@ -43,14 +43,14 @@ DBX 的 Agent 驱动 —— 通过 JDBC 和原生数据库驱动支持各种数 | xugu | 虚谷 XuguDB | XuguDB Go 原生 agent | | iotdb | Apache IoTDB | Apache IoTDB Go Client 原生 Agent | | etcd | etcd | jetcd | -| zookeeper | Apache ZooKeeper | Apache Curator | +| zookeeper | Apache ZooKeeper | go-zookeeper 原生 Agent | | rabbitmq | RabbitMQ | amqp091-go 原生 agent | | rocketmq | Apache RocketMQ | rocketmq-admin-go 原生 agent | ## 多 JRE 支持 -多数 Java agent 以 JRE 21 为目标。原生 agent（如 `cassandra`、`duckdb`、`hive`、`iotdb`、`oracle`、`kingbase`、`tdengine`、`xugu`、`rabbitmq` 和 `rocketmq`）不需要 JRE。对 Java agent，DBX 会自动下载并管理 JRE 21 安装。 +多数 Java agent 以 JRE 21 为目标。原生 agent（如 `cassandra`、`duckdb`、`hive`、`iotdb`、`oracle`、`kingbase`、`tdengine`、`xugu`、`rabbitmq`、`rocketmq` 和 `zookeeper`）不需要 JRE。对 Java agent，DBX 会自动下载并管理 JRE 21 安装。 ## JDBC 连接池 @@ -76,7 +76,7 @@ HikariCP 会直接打进启用连接池的 Agent JAR。已经使用 DBX 托管 J 对于新 agent，只要存在成熟、许可证兼容的原生驱动，优先选择**原生（Go 或 Rust）驱动**而非 Java/JDBC agent。原生 agent 以单一自包含可执行文件发布，无需 JRE，可显著降低内存占用和启动时间 —— 完全避开 Java agent 即便空闲也要付出的 JVM 基线开销。 -- **原生（Go/Rust）** —— 存在可用原生驱动时首选。参考 `drivers/cassandra-go`（Apache cassandra-gocql-driver）、`drivers/duckdb`、`drivers/hive-go`（原生 HS2）、`drivers/iotdb`（Apache IoTDB Go Client）、`drivers/oracle-go`（go-ora）、`drivers/kingbase-go`（gokb）、`drivers/vastbase-go`（openGauss connector）、`drivers/tdengine`（taos-connector-rust）、`drivers/xugu`、`drivers/rabbitmq`（amqp091-go）和 `drivers/rocketmq`（rocketmq-admin-go）。无需 JRE 下载与管理。 +- **原生（Go/Rust）** —— 存在可用原生驱动时首选。参考 `drivers/cassandra-go`（Apache cassandra-gocql-driver）、`drivers/duckdb`、`drivers/hive-go`（原生 HS2）、`drivers/iotdb`（Apache IoTDB Go Client）、`drivers/oracle-go`（go-ora）、`drivers/kingbase-go`（gokb）、`drivers/vastbase-go`（openGauss connector）、`drivers/tdengine`（taos-connector-rust）、`drivers/xugu`、`drivers/rabbitmq`（amqp091-go）、`drivers/rocketmq`（rocketmq-admin-go）和 `drivers/zookeeper`（go-zookeeper）。无需 JRE 下载与管理。 - **Java/JDBC** —— 当某数据库只有 JDBC 驱动，或原生驱动不成熟、缺乏维护时的默认兜底方案。多数 agent 仍属此类。 原生 agent 实现与 Java agent 相同的 JSON-RPC 契约和 `versions.json` 登记；它发布的是 `agent` 可执行文件而非 `agent.jar`。若同一数据库同时保留原生和 Java 源码实现，默认只发布原生产物；只有 Java 变体以独立兼容配置登记时才同时发布，例如 `oracle-legacy` / `oracle-10g`。 @@ -97,9 +97,10 @@ HikariCP 会直接打进启用连接池的 Agent JAR。已经使用 DBX 托管 J (cd drivers/xugu && go build -o agent .) (cd drivers/rabbitmq && go build -o agent .) (cd drivers/rocketmq && go build -o agent .) +(cd drivers/zookeeper && go build -o agent .) ``` -产物 JAR 在 `drivers/{module}/build/libs/`。原生 agent 从 `drivers/cassandra-go`、`drivers/duckdb`、`drivers/hive-go`、`drivers/iotdb`、`drivers/oracle-go`、`drivers/kingbase-go`、`drivers/vastbase-go`、`drivers/tdengine`、`drivers/xugu`、`drivers/rabbitmq` 和 `drivers/rocketmq` 构建。 +产物 JAR 在 `drivers/{module}/build/libs/`。原生 agent 从 `drivers/cassandra-go`、`drivers/duckdb`、`drivers/hive-go`、`drivers/iotdb`、`drivers/oracle-go`、`drivers/kingbase-go`、`drivers/vastbase-go`、`drivers/tdengine`、`drivers/xugu`、`drivers/rabbitmq`、`drivers/rocketmq` 和 `drivers/zookeeper` 构建。 ### 本地 DBX 运行时测试 @@ -113,7 +114,7 @@ cp agents/drivers//build/libs/*-all.jar ~/.dbx/agents/drivers/ 重启 DBX 或断开重连数据库，使新 agent 进程加载替换后的 JAR。 -`cassandra`、`hive`、`iotdb`、`oracle`、`kingbase`、`tdengine`、`xugu`、`rabbitmq` 和 `rocketmq` 等原生 agent 使用可执行文件而非 `agent.jar`。TDengine 从 `drivers/tdengine/Cargo.toml` 构建 `target/release/dbx-tdengine-driver`。 +`cassandra`、`hive`、`iotdb`、`oracle`、`kingbase`、`tdengine`、`xugu`、`rabbitmq`、`rocketmq` 和 `zookeeper` 等原生 agent 使用可执行文件而非 `agent.jar`。TDengine 从 `drivers/tdengine/Cargo.toml` 构建 `target/release/dbx-tdengine-driver`。 ## 版本管理 diff --git a/agents/drivers/zookeeper/agent_test.go b/agents/drivers/zookeeper/agent_test.go new file mode 100644 index 0000000000..58508ef638 --- /dev/null +++ b/agents/drivers/zookeeper/agent_test.go @@ -0,0 +1,336 @@ +package main + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "sort" + "strings" + "sync" + "testing" + + "github.com/go-zookeeper/zk" +) + +func TestHandshakeAndRPCFailures(t *testing.T) { + service := newServer() + response, shutdown := service.handleRequest([]byte(`{"jsonrpc":"2.0","id":1,"method":"handshake","params":{}}`)) + if shutdown || response.Error != nil { + t.Fatalf("handshake response = %#v, shutdown=%v", response, shutdown) + } + handshake, ok := response.Result.(handshakeResult) + if !ok || handshake.ProtocolVersion != 1 || handshake.AgentProtocolVersion != 1 || strings.Join(handshake.Capabilities, ",") != "connect,test_connection,kv" { + t.Fatalf("unexpected handshake: %#v", response.Result) + } + + response, _ = service.handleRequest([]byte(`{"jsonrpc":"2.0","id":"get","method":"kv_get","params":{"key":"/a"}}`)) + if response.Error == nil || response.Error.Message != "Not connected" { + t.Fatalf("unexpected disconnected response: %#v", response) + } + response, _ = service.handleRequest([]byte(`{"jsonrpc":"2.0","id":2,"method":"missing","params":{}}`)) + if response.Error == nil || response.Error.Message != "Unknown method: missing" { + t.Fatalf("unexpected unknown-method response: %#v", response) + } + response, _ = service.handleRequest([]byte(`not-json`)) + if response.Error == nil || response.Error.Code != -1 || response.ID != nil { + t.Fatalf("unexpected parse response: %#v", response) + } +} + +func TestConnectionConfiguration(t *testing.T) { + if got := connectionString(connectionConfig{ZooKeeperConnectString: " zk-main:2181 "}); got != "zk-main:2181" { + t.Fatalf("zookeeper connect string = %q", got) + } + if got := connectionString(connectionConfig{ConnectString: "zk-alt:2181"}); got != "zk-alt:2181" { + t.Fatalf("connect string = %q", got) + } + if got := connectionString(connectionConfig{Host: "::1", Port: 2281}); got != "[::1]:2281" { + t.Fatalf("host fallback = %q", got) + } + target, err := parseConnectTarget("zookeeper://zk-a:2181, zk-b:2182/app/root/") + if err != nil || strings.Join(target.Servers, ",") != "zk-a:2181,zk-b:2182" || target.Chroot != "/app/root" { + t.Fatalf("target=%#v err=%v", target, err) + } + if got := joinPrefix("/app/", "/tenant/"); got != "/app/tenant" { + t.Fatalf("prefix = %q", got) + } + for input, expected := range map[string]string{ + "host": "host:2181", "host:2281": "host:2281", "[::1]": "[::1]:2181", "[::1]:2281": "[::1]:2281", "::1": "[::1]:2181", + } { + got, err := endpointAddress(input) + if err != nil || got != expected { + t.Errorf("endpointAddress(%q)=%q,%v want %q", input, got, err, expected) + } + } + if resolveAuthScheme(connectionConfig{URLParams: "?foo=1;auth_scheme=SASL_DIGEST"}) != saslDigestAuthScheme { + t.Fatal("URL auth_scheme was not resolved") + } + if !hasTLSOptions(connectionConfig{CACertPath: "/tmp/ca.pem"}) { + t.Fatal("TLS path was not detected") + } + for input, expected := range map[string]int{"": 16, "bad": 16, "0": 1, "32": 32, "100": 64} { + if got := configuredStatLookupConcurrency(input); got != expected { + t.Errorf("configuredStatLookupConcurrency(%q)=%d want %d", input, got, expected) + } + } +} + +func TestConnectionValidationHappensBeforeNetwork(t *testing.T) { + tests := []struct { + name string + config connectionConfig + message string + }{ + {"tls", connectionConfig{SSL: true}, "ZooKeeper TLS is not supported"}, + {"auth", connectionConfig{AuthScheme: "sasl"}, `Unsupported auth_scheme "sasl"`}, + {"sasl username", connectionConfig{AuthScheme: saslDigestAuthScheme, Password: "secret"}, `username is required when auth_scheme = "sasl_digest"`}, + {"sasl password", connectionConfig{AuthScheme: saslDigestAuthScheme, Username: "user"}, `password is required when auth_scheme = "sasl_digest"`}, + {"negative base sleep", connectionConfig{BaseSleepTimeMS: intPointer(-1)}, "base_sleep_time_ms must be non-negative"}, + {"negative retries", connectionConfig{MaxRetries: intPointer(-1)}, "max_retries must be non-negative"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := openClient(test.config) + if err == nil || !strings.Contains(err.Error(), test.message) { + t.Fatalf("error=%v want substring %q", err, test.message) + } + }) + } +} + +func intPointer(value int) *int { return &value } + +func TestPathValueCursorAndCreateModes(t *testing.T) { + for input, expected := range map[string]string{"": "/", "/": "/", "app": "/app", "/app/": "/app", "//app//": "//app"} { + if got := normalizePath(input); got != expected { + t.Errorf("normalizePath(%q)=%q want %q", input, got, expected) + } + } + if got := encodeValue([]byte("hello")); got.Encoding != "utf8" || got.Data != "hello" { + t.Fatalf("utf8 value = %#v", got) + } + binaryValue := []byte{0xff, 0x00, 0x81} + encoded := encodeValue(binaryValue) + if encoded.Encoding != "base64" || encoded.Data != base64.StdEncoding.EncodeToString(binaryValue) { + t.Fatalf("binary value = %#v", encoded) + } + decoded, err := decodeValue(encoded) + if err != nil || !bytes.Equal(decoded, binaryValue) { + t.Fatalf("decoded=%v err=%v", decoded, err) + } + if _, err := decodeValue(valueObject{Encoding: "hex"}); err == nil { + t.Fatal("unsupported encoding was accepted") + } + cursor := listCursor{Root: "/app", Recursive: true, Offset: 12} + continuation, err := encodeCursor(cursor) + if err != nil { + t.Fatal(err) + } + decodedCursor, err := decodeCursor(continuation) + if err != nil || decodedCursor != cursor { + t.Fatalf("cursor=%#v err=%v", decodedCursor, err) + } + for mode, expected := range map[string]int32{"": 0, "persistent": 0, "ephemeral": zk.FlagEphemeral, "persistent_sequential": zk.FlagSequence, "ephemeral_sequential": zk.FlagEphemeral | zk.FlagSequence} { + flags, err := createFlags(mode) + if err != nil || flags != expected { + t.Errorf("createFlags(%q)=%d,%v want %d", mode, flags, err, expected) + } + } +} + +func TestZooKeeperRetryClassification(t *testing.T) { + for _, err := range []error{zk.ErrConnectionClosed, zk.ErrClosing, zk.ErrSessionMoved} { + if !isRetryableZooKeeperError(err) { + t.Fatalf("expected retryable error: %v", err) + } + } + for _, err := range []error{zk.ErrNoNode, zk.ErrAuthFailed, zk.ErrSessionExpired} { + if isRetryableZooKeeperError(err) { + t.Fatalf("unexpected retryable error: %v", err) + } + } +} + +func TestKVOperationsAndPagination(t *testing.T) { + client := newMemoryClient() + service := &server{activeClient: client, statLookupConcurrency: 4} + put := func(payload string) map[string]any { + result, err := service.put(json.RawMessage(payload)) + if err != nil { + t.Fatal(err) + } + return result + } + put(`{"key":"/app/name","value":{"encoding":"utf8","data":"dbx"}}`) + put(`{"key":"/app/a","value":{"encoding":"base64","data":"/wA="},"writeMode":"create"}`) + sequential := put(`{"key":"/app/seq-","value":{"data":"v"},"writeMode":"create","createMode":"persistent_sequential"}`) + if !strings.HasPrefix(sequential["createdKey"].(string), "/app/seq-") { + t.Fatalf("sequential result = %#v", sequential) + } + got, err := service.get(json.RawMessage(`{"key":"app/name/"}`)) + if err != nil || got["found"] != true || got["value"].(valueObject).Data != "dbx" { + t.Fatalf("get=%#v err=%v", got, err) + } + put(`{"key":"/app/name","value":{"data":"dbx2"},"writeMode":"update"}`) + + recursive := true + first, err := service.listPrefix(mustJSON(listRequest{Prefix: "/", Recursive: &recursive, Limit: 2})) + if err != nil || len(first.Keys) != 2 || first.Continuation == nil { + t.Fatalf("first page=%#v err=%v", first, err) + } + second, err := service.listPrefix(mustJSON(listRequest{Prefix: "/", Recursive: &recursive, Limit: 100, Continuation: *first.Continuation})) + if err != nil || len(second.Keys) < 2 || second.Continuation != nil { + t.Fatalf("second page=%#v err=%v", second, err) + } + wrongRecursive := false + if _, err := service.listPrefix(mustJSON(listRequest{Prefix: "/", Recursive: &wrongRecursive, Continuation: *first.Continuation})); err == nil || err.Error() != "Continuation does not match request" { + t.Fatalf("continuation mismatch error=%v", err) + } + + deleted, err := service.delete(json.RawMessage(`{"key":"/app","recursive":true}`)) + if err != nil || deleted["deleted"].(int) < 4 { + t.Fatalf("delete=%#v err=%v", deleted, err) + } + missing, err := service.delete(json.RawMessage(`{"key":"/app","recursive":true}`)) + if err != nil || missing["deleted"] != 0 { + t.Fatalf("missing delete=%#v err=%v", missing, err) + } + if _, err := service.put(json.RawMessage(`{"key":"/","value":{"data":"x"}}`)); err == nil || err.Error() != "Root znode cannot be modified" { + t.Fatalf("root put error=%v", err) + } + if _, err := service.delete(json.RawMessage(`{"key":"/","recursive":true}`)); err == nil || err.Error() != "Root znode cannot be deleted" { + t.Fatalf("root delete error=%v", err) + } +} + +func mustJSON(value any) json.RawMessage { + payload, err := json.Marshal(value) + if err != nil { + panic(err) + } + return payload +} + +type memoryNode struct { + data []byte + stat zk.Stat +} + +type memoryClient struct { + mutex sync.Mutex + nodes map[string]*memoryNode + sequence int +} + +func newMemoryClient() *memoryClient { + return &memoryClient{nodes: map[string]*memoryNode{"/": {stat: zk.Stat{Ctime: 1, Mtime: 1}}}} +} + +func (client *memoryClient) Close() {} + +func (client *memoryClient) Exists(path string) (bool, *zk.Stat, error) { + client.mutex.Lock() + defer client.mutex.Unlock() + node, ok := client.nodes[path] + if !ok { + return false, nil, nil + } + stat := node.stat + stat.NumChildren = int32(len(client.childrenLocked(path))) + stat.DataLength = int32(len(node.data)) + return true, &stat, nil +} + +func (client *memoryClient) Get(path string) ([]byte, *zk.Stat, error) { + exists, stat, err := client.Exists(path) + if err != nil || !exists { + return nil, nil, zk.ErrNoNode + } + client.mutex.Lock() + data := append([]byte(nil), client.nodes[path].data...) + client.mutex.Unlock() + return data, stat, nil +} + +func (client *memoryClient) Children(path string) ([]string, *zk.Stat, error) { + exists, stat, err := client.Exists(path) + if err != nil || !exists { + return nil, nil, zk.ErrNoNode + } + client.mutex.Lock() + children := client.childrenLocked(path) + client.mutex.Unlock() + return children, stat, nil +} + +func (client *memoryClient) Create(path string, data []byte, flags int32) (string, error) { + client.mutex.Lock() + defer client.mutex.Unlock() + created := path + if flags&zk.FlagSequence != 0 { + created = fmt.Sprintf("%s%010d", path, client.sequence) + client.sequence++ + } + if _, exists := client.nodes[created]; exists { + return "", zk.ErrNodeExists + } + if _, exists := client.nodes[parentPath(created)]; !exists { + return "", zk.ErrNoNode + } + owner := int64(0) + if flags&zk.FlagEphemeral != 0 { + owner = 1 + } + client.nodes[created] = &memoryNode{data: append([]byte(nil), data...), stat: zk.Stat{Ctime: 1, Mtime: 1, EphemeralOwner: owner}} + return created, nil +} + +func (client *memoryClient) Set(path string, data []byte) (*zk.Stat, error) { + client.mutex.Lock() + defer client.mutex.Unlock() + node, exists := client.nodes[path] + if !exists { + return nil, zk.ErrNoNode + } + node.data = append([]byte(nil), data...) + node.stat.Version++ + node.stat.Mtime++ + stat := node.stat + stat.DataLength = int32(len(node.data)) + return &stat, nil +} + +func (client *memoryClient) Delete(path string) error { + client.mutex.Lock() + defer client.mutex.Unlock() + if _, exists := client.nodes[path]; !exists { + return zk.ErrNoNode + } + if len(client.childrenLocked(path)) != 0 { + return zk.ErrNotEmpty + } + delete(client.nodes, path) + return nil +} + +func (client *memoryClient) childrenLocked(path string) []string { + prefix := path + if prefix != "/" { + prefix += "/" + } + children := make([]string, 0) + for candidate := range client.nodes { + if candidate == path || !strings.HasPrefix(candidate, prefix) { + continue + } + remainder := strings.TrimPrefix(candidate, prefix) + if remainder != "" && !strings.Contains(remainder, "/") { + children = append(children, remainder) + } + } + sort.Strings(children) + return children +} + +var _ znodeClient = (*memoryClient)(nil) diff --git a/agents/drivers/zookeeper/build.gradle b/agents/drivers/zookeeper/build.gradle deleted file mode 100644 index f679c11a84..0000000000 --- a/agents/drivers/zookeeper/build.gradle +++ /dev/null @@ -1,13 +0,0 @@ -dependencies { - implementation 'com.google.code.gson:gson:2.12.1' - implementation 'org.apache.curator:curator-framework:5.9.0' - runtimeOnly 'org.slf4j:slf4j-simple:1.7.36' - testImplementation 'org.apache.curator:curator-test:5.9.0' -} - -tasks.named('shadowJar') { - mergeServiceFiles() - manifest { - attributes('Agent-Label': 'Apache ZooKeeper', 'Main-Class': 'com.dbx.agent.zookeeper.ZooKeeperAgent') - } -} diff --git a/agents/drivers/zookeeper/connection.go b/agents/drivers/zookeeper/connection.go new file mode 100644 index 0000000000..a525438df3 --- /dev/null +++ b/agents/drivers/zookeeper/connection.go @@ -0,0 +1,563 @@ +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "net/url" + "strconv" + "strings" + "sync" + "time" + + "github.com/go-zookeeper/zk" +) + +const ( + defaultSessionTimeout = 30 * time.Second + defaultConnectionTimeout = 15 * time.Second + defaultProbeTimeout = 2 * time.Second + defaultBaseSleepTime = 250 * time.Millisecond + defaultMaxRetries = 2 + defaultPort = 2181 + defaultAuthScheme = "digest" + saslDigestAuthScheme = "sasl_digest" + statLookupConcurrencyEnv = "DBX_ZOOKEEPER_STAT_LOOKUP_CONCURRENCY" + defaultStatLookupWorkers = 16 + minimumStatLookupWorkers = 1 + maximumStatLookupWorkers = 64 + maximumReachabilityWorkers = 8 +) + +type connectionConfig struct { + ZooKeeperConnectString string `json:"zookeeper_connect_string"` + ConnectString string `json:"connect_string"` + ConnectionString string `json:"connection_string"` + Host string `json:"host"` + Port int `json:"port"` + Namespace string `json:"namespace"` + Username string `json:"username"` + Password string `json:"password"` + AuthScheme string `json:"auth_scheme"` + URLParams string `json:"url_params"` + SessionTimeoutMS int `json:"session_timeout_ms"` + ConnectionTimeoutMS int `json:"connection_timeout_ms"` + BaseSleepTimeMS *int `json:"base_sleep_time_ms"` + MaxRetries *int `json:"max_retries"` + SSL bool `json:"ssl"` + CACertPath string `json:"ca_cert_path"` + ClientCertPath string `json:"client_cert_path"` + ClientKeyPath string `json:"client_key_path"` + CertPath string `json:"cert_path"` + KeyPath string `json:"key_path"` +} + +type connectionParams struct { + Connection json.RawMessage `json:"connection"` +} + +type connectTarget struct { + Servers []string + Chroot string +} + +type clientSession struct { + connection *zk.Conn + prefix string + retryBase time.Duration + maxRetries int +} + +type znodeClient interface { + Close() + Exists(path string) (bool, *zk.Stat, error) + Get(path string) ([]byte, *zk.Stat, error) + Children(path string) ([]string, *zk.Stat, error) + Create(path string, data []byte, flags int32) (string, error) + Set(path string, data []byte) (*zk.Stat, error) + Delete(path string) error +} + +func decodeConnectionConfig(params json.RawMessage) (connectionConfig, error) { + var wrapper connectionParams + if err := json.Unmarshal(params, &wrapper); err != nil { + return connectionConfig{}, err + } + payload := params + if len(wrapper.Connection) > 0 && string(wrapper.Connection) != "null" { + payload = wrapper.Connection + } + var config connectionConfig + if err := json.Unmarshal(payload, &config); err != nil { + return connectionConfig{}, err + } + return config, nil +} + +func (service *server) connect(params json.RawMessage) (map[string]bool, error) { + config, err := decodeConnectionConfig(params) + if err != nil { + return nil, err + } + nextClient, err := openClient(config) + if err != nil { + return nil, err + } + previousClient := service.activeClient + service.activeClient = nextClient + if previousClient != nil { + previousClient.Close() + } + return map[string]bool{"ok": true}, nil +} + +func (service *server) testConnection(params json.RawMessage) (map[string]bool, error) { + config, err := decodeConnectionConfig(params) + if err != nil { + return nil, err + } + probe, err := openClient(config) + if err != nil { + return nil, err + } + probe.Close() + return map[string]bool{"ok": true}, nil +} + +func openClient(config connectionConfig) (*clientSession, error) { + if hasTLSOptions(config) { + return nil, errors.New("ZooKeeper TLS is not supported") + } + authScheme := resolveAuthScheme(config) + if authScheme != defaultAuthScheme && authScheme != saslDigestAuthScheme { + return nil, fmt.Errorf("Unsupported auth_scheme %q; expected %q or %q", authScheme, defaultAuthScheme, saslDigestAuthScheme) + } + if authScheme == saslDigestAuthScheme { + if strings.TrimSpace(config.Username) == "" { + return nil, errors.New(`username is required when auth_scheme = "sasl_digest"`) + } + if config.Password == "" { + return nil, errors.New(`password is required when auth_scheme = "sasl_digest"`) + } + } + if config.BaseSleepTimeMS != nil && *config.BaseSleepTimeMS < 0 { + return nil, errors.New("base_sleep_time_ms must be non-negative") + } + if config.MaxRetries != nil && *config.MaxRetries < 0 { + return nil, errors.New("max_retries must be non-negative") + } + + target, err := parseConnectTarget(connectionString(config)) + if err != nil { + return nil, err + } + connectionTimeout := millisecondsOrDefault(config.ConnectionTimeoutMS, defaultConnectionTimeout) + probeTimeout := minDuration(defaultProbeTimeout, connectionTimeout) + if err := requireReachableServer(target.Servers, probeTimeout); err != nil { + return nil, err + } + + dialer := newZooKeeperDialer(connectionTimeout, nil) + if authScheme == saslDigestAuthScheme { + dialer = newZooKeeperDialer(connectionTimeout, &saslDigestCredentials{ + Username: strings.TrimSpace(config.Username), + Password: config.Password, + }) + } + + sessionTimeout := millisecondsOrDefault(config.SessionTimeoutMS, defaultSessionTimeout) + connection, events, err := zk.Connect( + target.Servers, + sessionTimeout, + zk.WithDialer(dialer), + zk.WithLogInfo(false), + ) + if err != nil { + return nil, err + } + connected := false + timer := time.NewTimer(connectionTimeout) + defer timer.Stop() + for !connected { + select { + case event, open := <-events: + if !open { + connection.Close() + return nil, errors.New("Connection timed out") + } + if event.State == zk.StateHasSession { + connected = true + } + if event.State == zk.StateAuthFailed { + connection.Close() + return nil, errors.New("ZooKeeper authentication failed") + } + case <-timer.C: + connection.Close() + return nil, errors.New("Connection timed out") + } + } + + if authScheme == defaultAuthScheme && strings.TrimSpace(config.Username) != "" { + credentials := []byte(strings.TrimSpace(config.Username) + ":" + config.Password) + if err := connection.AddAuth(defaultAuthScheme, credentials); err != nil { + connection.Close() + return nil, err + } + } + + prefix := joinPrefix(target.Chroot, config.Namespace) + retryBase := defaultBaseSleepTime + if config.BaseSleepTimeMS != nil { + retryBase = time.Duration(*config.BaseSleepTimeMS) * time.Millisecond + } + maxRetries := defaultMaxRetries + if config.MaxRetries != nil { + maxRetries = *config.MaxRetries + } + session := &clientSession{connection: connection, prefix: prefix, retryBase: retryBase, maxRetries: maxRetries} + exists, _, err := session.Exists("/") + if err != nil || !exists { + connection.Close() + if err != nil { + return nil, err + } + return nil, errors.New("Root znode is not readable") + } + return session, nil +} + +func newZooKeeperDialer(connectionTimeout time.Duration, credentials *saslDigestCredentials) zk.Dialer { + return func(network, address string, libraryTimeout time.Duration) (net.Conn, error) { + timeout := libraryTimeout + if timeout <= 0 || connectionTimeout < timeout { + timeout = connectionTimeout + } + connection, err := net.DialTimeout(network, address, timeout) + if err != nil { + return nil, err + } + if credentials == nil { + return connection, nil + } + return newSASLHandshakeConn(connection, timeout, *credentials), nil + } +} + +func connectionString(config connectionConfig) string { + for _, candidate := range []string{config.ZooKeeperConnectString, config.ConnectString, config.ConnectionString} { + if strings.TrimSpace(candidate) != "" { + return strings.TrimSpace(candidate) + } + } + host := strings.TrimSpace(config.Host) + if host == "" { + host = "127.0.0.1" + } + port := config.Port + if port <= 0 { + port = defaultPort + } + return net.JoinHostPort(strings.Trim(host, "[]"), strconv.Itoa(port)) +} + +func parseConnectTarget(value string) (connectTarget, error) { + connectString := strings.TrimSpace(strings.TrimPrefix(value, "zookeeper://")) + slash := strings.Index(connectString, "/") + hostsPart := connectString + chroot := "" + if slash >= 0 { + hostsPart = connectString[:slash] + chroot = normalizePrefix(connectString[slash:]) + } + servers := make([]string, 0) + for _, item := range strings.Split(hostsPart, ",") { + server := strings.TrimSpace(item) + if server != "" { + servers = append(servers, server) + } + } + if len(servers) == 0 { + return connectTarget{}, errors.New("ZooKeeper connect string contains no servers") + } + return connectTarget{Servers: servers, Chroot: chroot}, nil +} + +func requireReachableServer(servers []string, timeout time.Duration) error { + ctx, cancel := context.WithTimeout(context.Background(), timeout+500*time.Millisecond) + defer cancel() + workers := minInt(len(servers), maximumReachabilityWorkers) + jobs := make(chan string) + reachable := make(chan struct{}, 1) + var waitGroup sync.WaitGroup + for worker := 0; worker < workers; worker++ { + waitGroup.Add(1) + go func() { + defer waitGroup.Done() + for server := range jobs { + address, err := endpointAddress(server) + if err != nil { + continue + } + connection, err := net.DialTimeout("tcp", address, timeout) + if err == nil { + connection.Close() + select { + case reachable <- struct{}{}: + default: + } + } + } + }() + } + go func() { + defer close(jobs) + for _, server := range servers { + select { + case jobs <- server: + case <-ctx.Done(): + return + } + } + }() + done := make(chan struct{}) + go func() { + waitGroup.Wait() + close(done) + }() + select { + case <-reachable: + return nil + case <-done: + case <-ctx.Done(): + } + return fmt.Errorf("No reachable ZooKeeper server within %dms: %s", timeout.Milliseconds(), strings.Join(servers, ",")) +} + +func endpointAddress(endpoint string) (string, error) { + value := strings.TrimSpace(endpoint) + if strings.HasPrefix(value, "[") { + if _, _, err := net.SplitHostPort(value); err == nil { + return value, nil + } + return value + ":" + strconv.Itoa(defaultPort), nil + } + if strings.Count(value, ":") == 0 { + return net.JoinHostPort(value, strconv.Itoa(defaultPort)), nil + } + if strings.Count(value, ":") == 1 { + if _, _, err := net.SplitHostPort(value); err != nil { + return "", err + } + return value, nil + } + return net.JoinHostPort(value, strconv.Itoa(defaultPort)), nil +} + +func resolveAuthScheme(config connectionConfig) string { + if strings.TrimSpace(config.AuthScheme) != "" { + return strings.ToLower(strings.TrimSpace(config.AuthScheme)) + } + params := strings.TrimPrefix(strings.TrimSpace(config.URLParams), "?") + params = strings.ReplaceAll(params, ";", "&") + if parsed, err := url.ParseQuery(params); err == nil { + if configured := strings.TrimSpace(parsed.Get("auth_scheme")); configured != "" { + return strings.ToLower(configured) + } + } + return defaultAuthScheme +} + +func hasTLSOptions(config connectionConfig) bool { + return config.SSL || firstNonBlank( + config.CACertPath, + config.ClientCertPath, + config.ClientKeyPath, + config.CertPath, + config.KeyPath, + ) != "" +} + +func joinPrefix(chroot, namespace string) string { + parts := make([]string, 0, 2) + if normalized := normalizePrefix(chroot); normalized != "" { + parts = append(parts, strings.Trim(normalized, "/")) + } + if normalized := normalizePrefix(namespace); normalized != "" { + parts = append(parts, strings.Trim(normalized, "/")) + } + if len(parts) == 0 { + return "" + } + return "/" + strings.Join(parts, "/") +} + +func normalizePrefix(value string) string { + trimmed := strings.Trim(strings.TrimSpace(value), "/") + if trimmed == "" { + return "" + } + return "/" + trimmed +} + +func millisecondsOrDefault(value int, fallback time.Duration) time.Duration { + if value <= 0 { + return fallback + } + return time.Duration(value) * time.Millisecond +} + +func configuredStatLookupConcurrency(value string) int { + parsed, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil { + return defaultStatLookupWorkers + } + return maxInt(minimumStatLookupWorkers, minInt(maximumStatLookupWorkers, parsed)) +} + +func (service *server) requireClient() (znodeClient, error) { + if service.activeClient == nil { + return nil, errors.New("Not connected") + } + return service.activeClient, nil +} + +func (service *server) closeClient() { + if service.activeClient != nil { + service.activeClient.Close() + service.activeClient = nil + } +} + +func (session *clientSession) Close() { + if session != nil && session.connection != nil { + session.connection.Close() + } +} + +func (session *clientSession) physicalPath(logicalPath string) string { + logical := normalizePath(logicalPath) + if session.prefix == "" { + return logical + } + if logical == "/" { + return session.prefix + } + return session.prefix + logical +} + +func (session *clientSession) logicalPath(physicalPath string) string { + if session.prefix == "" { + return normalizePath(physicalPath) + } + trimmed := strings.TrimPrefix(physicalPath, session.prefix) + return normalizePath(trimmed) +} + +func (session *clientSession) Exists(path string) (bool, *zk.Stat, error) { + type result struct { + exists bool + stat *zk.Stat + } + value, err := retryZooKeeper(session, func() (result, error) { + exists, stat, err := session.connection.Exists(session.physicalPath(path)) + return result{exists: exists, stat: stat}, err + }) + return value.exists, value.stat, err +} + +func (session *clientSession) Get(path string) ([]byte, *zk.Stat, error) { + type result struct { + data []byte + stat *zk.Stat + } + value, err := retryZooKeeper(session, func() (result, error) { + data, stat, err := session.connection.Get(session.physicalPath(path)) + return result{data: data, stat: stat}, err + }) + return value.data, value.stat, err +} + +func (session *clientSession) Children(path string) ([]string, *zk.Stat, error) { + type result struct { + children []string + stat *zk.Stat + } + value, err := retryZooKeeper(session, func() (result, error) { + children, stat, err := session.connection.Children(session.physicalPath(path)) + return result{children: children, stat: stat}, err + }) + return value.children, value.stat, err +} + +func (session *clientSession) Create(path string, data []byte, flags int32) (string, error) { + createdPath, err := retryZooKeeper(session, func() (string, error) { + return session.connection.Create(session.physicalPath(path), data, flags, zk.WorldACL(zk.PermAll)) + }) + if err != nil { + return "", err + } + return session.logicalPath(createdPath), nil +} + +func (session *clientSession) Set(path string, data []byte) (*zk.Stat, error) { + return retryZooKeeper(session, func() (*zk.Stat, error) { + return session.connection.Set(session.physicalPath(path), data, -1) + }) +} + +func (session *clientSession) Delete(path string) error { + _, err := retryZooKeeper(session, func() (struct{}, error) { + return struct{}{}, session.connection.Delete(session.physicalPath(path), -1) + }) + return err +} + +func retryZooKeeper[T any](session *clientSession, operation func() (T, error)) (T, error) { + for attempt := 0; ; attempt++ { + value, err := operation() + if err == nil || attempt >= session.maxRetries || !isRetryableZooKeeperError(err) { + return value, err + } + delay := session.retryBase * time.Duration(1< 0 { + time.Sleep(delay) + } + } +} + +func isRetryableZooKeeperError(err error) bool { + return errors.Is(err, zk.ErrConnectionClosed) || errors.Is(err, zk.ErrClosing) || errors.Is(err, zk.ErrSessionMoved) +} + +func minDuration(first, second time.Duration) time.Duration { + if first < second { + return first + } + return second +} + +func firstNonBlank(values ...string) string { + for _, value := range values { + if strings.TrimSpace(value) != "" { + return value + } + } + return "" +} + +func minInt(first, second int) int { + if first < second { + return first + } + return second +} + +func maxInt(first, second int) int { + if first > second { + return first + } + return second +} diff --git a/agents/drivers/zookeeper/go.mod b/agents/drivers/zookeeper/go.mod new file mode 100644 index 0000000000..d8d2908a3c --- /dev/null +++ b/agents/drivers/zookeeper/go.mod @@ -0,0 +1,26 @@ +module github.com/t8y2/dbx/agents/drivers/zookeeper + +go 1.23.0 + +require ( + github.com/beltran/gosasl v1.0.0 + github.com/go-zookeeper/zk v1.0.4 +) + +require ( + github.com/alexbrainman/sspi v0.0.0-20250919150558-7d374ff0d59e // indirect + github.com/golang-auth/go-gssapi/v2 v2.0.0 // indirect + github.com/hashicorp/go-uuid v1.0.2 // indirect + github.com/jcmturner/aescts/v2 v2.0.0 // indirect + github.com/jcmturner/dnsutils/v2 v2.0.0 // indirect + github.com/jcmturner/gofork v1.0.0 // indirect + github.com/jcmturner/gokrb5 v8.4.2+incompatible // indirect + github.com/jcmturner/gokrb5/v8 v8.4.2 // indirect + github.com/jcmturner/rpc/v2 v2.0.3 // indirect + golang.org/x/crypto v0.0.0-20201112155050-0c6587e931a9 // indirect + golang.org/x/net v0.0.0-20210326220855-61e056675ecf // indirect +) + +replace github.com/beltran/gosasl => ../../go-common/gosasl + +replace github.com/golang-auth/go-gssapi/v2 => ../../go-common/go-gssapi diff --git a/agents/drivers/zookeeper/go.sum b/agents/drivers/zookeeper/go.sum new file mode 100644 index 0000000000..e054ddb80b --- /dev/null +++ b/agents/drivers/zookeeper/go.sum @@ -0,0 +1,50 @@ +github.com/alexbrainman/sspi v0.0.0-20250919150558-7d374ff0d59e h1:4dAU9FXIyQktpoUAgOJK3OTFc/xug0PCXYCqU0FgDKI= +github.com/alexbrainman/sspi v0.0.0-20250919150558-7d374ff0d59e/go.mod h1:cEWa1LVoE5KvSD9ONXsZrj0z6KqySlCCNKHlLzbqAt4= +github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/go-zookeeper/zk v1.0.4 h1:DPzxraQx7OrPyXq2phlGlNSIyWEsAox0RJmjTseMV6I= +github.com/go-zookeeper/zk v1.0.4/go.mod h1:nOB03cncLtlp4t+UAkGSV+9beXP/akpekBwL+UX1Qcw= +github.com/gorilla/securecookie v1.1.1/go.mod h1:ra0sb63/xPlUeL+yeDciTfxMRAA+MP+HVt/4epWDjd4= +github.com/gorilla/sessions v1.2.1/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM= +github.com/hashicorp/go-uuid v1.0.2 h1:cfejS+Tpcp13yd5nYHWDI6qVCny6wyX2Mt5SGur2IGE= +github.com/hashicorp/go-uuid v1.0.2/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro= +github.com/jcmturner/aescts/v2 v2.0.0 h1:9YKLH6ey7H4eDBXW8khjYslgyqG2xZikXP0EQFKrle8= +github.com/jcmturner/aescts/v2 v2.0.0/go.mod h1:AiaICIRyfYg35RUkr8yESTqvSy7csK90qZ5xfvvsoNs= +github.com/jcmturner/dnsutils/v2 v2.0.0 h1:lltnkeZGL0wILNvrNiVCR6Ro5PGU/SeBvVO/8c/iPbo= +github.com/jcmturner/dnsutils/v2 v2.0.0/go.mod h1:b0TnjGOvI/n42bZa+hmXL+kFJZsFT7G4t3HTlQ184QM= +github.com/jcmturner/gofork v1.0.0 h1:J7uCkflzTEhUZ64xqKnkDxq3kzc96ajM1Gli5ktUem8= +github.com/jcmturner/gofork v1.0.0/go.mod h1:MK8+TM0La+2rjBD4jE12Kj1pCCxK7d2LK/UM3ncEo0o= +github.com/jcmturner/goidentity/v6 v6.0.1 h1:VKnZd2oEIMorCTsFBnJWbExfNN7yZr3EhJAxwOkZg6o= +github.com/jcmturner/goidentity/v6 v6.0.1/go.mod h1:X1YW3bgtvwAXju7V3LCIMpY0Gbxyjn/mY9zx4tFonSg= +github.com/jcmturner/gokrb5 v8.4.2+incompatible h1:MQW70Fbazv31g6URAXCjO2bGenIL0wVt3wqcpc0EjHI= +github.com/jcmturner/gokrb5 v8.4.2+incompatible/go.mod h1:0Q5eFyVvYsEsZ8xl1A/jUqhXvxUp/X9ELrJm+zieq5E= +github.com/jcmturner/gokrb5/v8 v8.4.2 h1:6ZIM6b/JJN0X8UM43ZOM6Z4SJzla+a/u7scXFJzodkA= +github.com/jcmturner/gokrb5/v8 v8.4.2/go.mod h1:sb+Xq/fTY5yktf/VxLsE3wlfPqQjp0aWNYyvBVK62bc= +github.com/jcmturner/rpc/v2 v2.0.3 h1:7FXXj8Ti1IaVFpSAziCZWNzbNuZmnvw/i6CqLNdWfZY= +github.com/jcmturner/rpc/v2 v2.0.3/go.mod h1:VUJYCIDm3PVOEHw8sgt091/20OJjskO/YJki3ELg/Hc= +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/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= +github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= +golang.org/x/crypto v0.0.0-20201112155050-0c6587e931a9 h1:umElSU9WZirRdgu2yFHY0ayQkEnKiOC1TtM3fWXFnoU= +golang.org/x/crypto v0.0.0-20201112155050-0c6587e931a9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= +golang.org/x/net v0.0.0-20200114155413-6afb5195e5aa/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20210326220855-61e056675ecf h1:WUcCxqQqDT0aXO4VnQbfMvp4zh7m1Gb2clVuHUAGGRE= +golang.org/x/net v0.0.0-20210326220855-61e056675ecf/go.mod h1:uSPa2vr4CLtc/ILN5odXGNXS6mhrKVzTaCXzk9m6W3k= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210324051608-47abb6519492/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/agents/drivers/zookeeper/integration_test.go b/agents/drivers/zookeeper/integration_test.go new file mode 100644 index 0000000000..6dd9d9b13d --- /dev/null +++ b/agents/drivers/zookeeper/integration_test.go @@ -0,0 +1,105 @@ +package main + +import ( + "encoding/json" + "fmt" + "os" + "testing" + "time" +) + +const integrationConnectStringEnv = "DBX_ZOOKEEPER_TEST_CONNECT_STRING" + +func TestZooKeeperIntegration(t *testing.T) { + connectString := os.Getenv(integrationConnectStringEnv) + if connectString == "" { + t.Skip(integrationConnectStringEnv + " is not set") + } + service := &server{statLookupConcurrency: 4} + connection := map[string]any{"zookeeper_connect_string": connectString} + if authScheme := os.Getenv("DBX_ZOOKEEPER_TEST_AUTH_SCHEME"); authScheme != "" { + connection["auth_scheme"] = authScheme + connection["username"] = os.Getenv("DBX_ZOOKEEPER_TEST_USERNAME") + connection["password"] = os.Getenv("DBX_ZOOKEEPER_TEST_PASSWORD") + } + params := mustJSON(map[string]any{"connection": connection}) + if _, err := service.connect(params); err != nil { + t.Fatal(err) + } + t.Cleanup(service.closeClient) + root := fmt.Sprintf("/dbx-go-integration-%d", time.Now().UnixNano()) + if _, err := service.put(mustJSON(map[string]any{"key": root + "/parent/value", "value": map[string]string{"encoding": "utf8", "data": "first"}})); err != nil { + t.Fatal(err) + } + result, err := service.get(mustJSON(map[string]any{"key": root + "/parent/value"})) + if err != nil || result["found"] != true { + t.Fatalf("get=%#v err=%v", result, err) + } + badConnection := mustJSON(map[string]any{"connection": map[string]any{"zookeeper_connect_string": "127.0.0.1:1", "connection_timeout_ms": 50}}) + if _, err := service.testConnection(badConnection); err == nil { + t.Fatal("unreachable test connection unexpectedly succeeded") + } + if _, err := service.connect(badConnection); err == nil { + t.Fatal("unreachable replacement connection unexpectedly succeeded") + } + result, err = service.get(mustJSON(map[string]any{"key": root + "/parent/value"})) + if err != nil || result["found"] != true { + t.Fatalf("active connection was replaced after failed probe: get=%#v err=%v", result, err) + } + if _, err := service.put(mustJSON(map[string]any{"key": root + "/ephemeral-", "value": map[string]string{"data": "e"}, "writeMode": "create", "createMode": "ephemeral_sequential"})); err != nil { + t.Fatal(err) + } + listed, err := service.listPrefix(mustJSON(map[string]any{"prefix": root, "recursive": true, "limit": 100})) + if err != nil || len(listed.Keys) < 3 { + t.Fatalf("list=%#v err=%v", listed, err) + } + deleted, err := service.delete(mustJSON(map[string]any{"key": root, "recursive": true})) + if err != nil || deleted["deleted"].(int) < 4 { + t.Fatalf("delete=%#v err=%v", deleted, err) + } +} + +func BenchmarkZooKeeperOperations(b *testing.B) { + connectString := os.Getenv(integrationConnectStringEnv) + if connectString == "" { + b.Skip(integrationConnectStringEnv + " is not set") + } + service := &server{statLookupConcurrency: 16} + if _, err := service.connect(mustJSON(map[string]any{"connection": map[string]any{"zookeeper_connect_string": connectString}})); err != nil { + b.Fatal(err) + } + b.Cleanup(service.closeClient) + root := fmt.Sprintf("/dbx-go-benchmark-%d", time.Now().UnixNano()) + if _, err := service.put(mustJSON(map[string]any{"key": root + "/value", "value": map[string]string{"data": "warmup"}})); err != nil { + b.Fatal(err) + } + b.Cleanup(func() { _, _ = service.delete(mustJSON(map[string]any{"key": root, "recursive": true})) }) + + b.Run("get", func(b *testing.B) { + params := json.RawMessage(fmt.Sprintf(`{"key":%q}`, root+"/value")) + b.ResetTimer() + for index := 0; index < b.N; index++ { + if _, err := service.get(params); err != nil { + b.Fatal(err) + } + } + }) + b.Run("put", func(b *testing.B) { + params := json.RawMessage(fmt.Sprintf(`{"key":%q,"value":{"data":"updated"}}`, root+"/value")) + b.ResetTimer() + for index := 0; index < b.N; index++ { + if _, err := service.put(params); err != nil { + b.Fatal(err) + } + } + }) + b.Run("list", func(b *testing.B) { + params := json.RawMessage(fmt.Sprintf(`{"prefix":%q,"recursive":true,"limit":100}`, root)) + b.ResetTimer() + for index := 0; index < b.N; index++ { + if _, err := service.listPrefix(params); err != nil { + b.Fatal(err) + } + } + }) +} diff --git a/agents/drivers/zookeeper/main.go b/agents/drivers/zookeeper/main.go new file mode 100644 index 0000000000..95ae868d89 --- /dev/null +++ b/agents/drivers/zookeeper/main.go @@ -0,0 +1,128 @@ +package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" +) + +const ( + protocolVersion = 1 + agentProtocolVersion = 1 + maxRPCMessageBytes = 32 * 1024 * 1024 +) + +var capabilities = []string{"connect", "test_connection", "kv"} + +type rpcRequest struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + Params json.RawMessage `json:"params"` +} + +type rpcError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +type rpcResponse struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result any `json:"result,omitempty"` + Error *rpcError `json:"error,omitempty"` +} + +type handshakeResult struct { + ProtocolVersion int `json:"protocolVersion"` + AgentProtocolVersion int `json:"agentProtocolVersion"` + Capabilities []string `json:"capabilities"` +} + +type server struct { + activeClient znodeClient + statLookupConcurrency int +} + +func main() { + service := newServer() + defer service.closeClient() + + encoder := json.NewEncoder(os.Stdout) + encoder.SetEscapeHTML(false) + fmt.Fprintln(os.Stdout, `{"ready":true}`) + + scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 64*1024), maxRPCMessageBytes) + for scanner.Scan() { + response, shutdown := service.handleRequest(scanner.Bytes()) + if err := encoder.Encode(response); err != nil { + fmt.Fprintln(os.Stderr, err) + return + } + if shutdown { + return + } + } + if err := scanner.Err(); err != nil { + fmt.Fprintln(os.Stderr, err) + } +} + +func newServer() *server { + return &server{statLookupConcurrency: configuredStatLookupConcurrency(os.Getenv(statLookupConcurrencyEnv))} +} + +func (service *server) handleRequest(payload []byte) (rpcResponse, bool) { + var request rpcRequest + if err := json.Unmarshal(payload, &request); err != nil { + return rpcResponse{JSONRPC: "2.0", Error: &rpcError{Code: -1, Message: err.Error()}}, false + } + + result, shutdown, err := service.dispatch(request.Method, request.Params) + if err != nil { + return rpcResponse{ + JSONRPC: "2.0", + ID: request.ID, + Error: &rpcError{Code: -1, Message: err.Error()}, + }, false + } + return rpcResponse{JSONRPC: "2.0", ID: request.ID, Result: result}, shutdown +} + +func (service *server) dispatch(method string, params json.RawMessage) (any, bool, error) { + switch method { + case "handshake": + return handshakeResult{ + ProtocolVersion: protocolVersion, + AgentProtocolVersion: agentProtocolVersion, + Capabilities: capabilities, + }, false, nil + case "connect": + result, err := service.connect(params) + return result, false, err + case "test_connection": + result, err := service.testConnection(params) + return result, false, err + case "kv_list_prefix": + result, err := service.listPrefix(params) + return result, false, err + case "kv_get": + result, err := service.get(params) + return result, false, err + case "kv_put": + result, err := service.put(params) + return result, false, err + case "kv_delete": + result, err := service.delete(params) + return result, false, err + case "disconnect": + service.closeClient() + return map[string]bool{"ok": true}, false, nil + case "shutdown": + service.closeClient() + return map[string]bool{"ok": true}, true, nil + default: + return nil, false, fmt.Errorf("Unknown method: %s", method) + } +} diff --git a/agents/drivers/zookeeper/operations.go b/agents/drivers/zookeeper/operations.go new file mode 100644 index 0000000000..d446d4e9a6 --- /dev/null +++ b/agents/drivers/zookeeper/operations.go @@ -0,0 +1,515 @@ +package main + +import ( + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "sort" + "strings" + "sync" + "unicode/utf8" + + "github.com/go-zookeeper/zk" +) + +const defaultListLimit = 100 + +type valueObject struct { + Encoding string `json:"encoding"` + Data string `json:"data"` +} + +type getRequest struct { + Key string `json:"key"` +} + +type putRequest struct { + Key string `json:"key"` + Value valueObject `json:"value"` + WriteMode string `json:"writeMode"` + CreateMode string `json:"createMode"` +} + +type deleteRequest struct { + Key string `json:"key"` + Recursive bool `json:"recursive"` +} + +type listRequest struct { + Prefix string `json:"prefix"` + Recursive *bool `json:"recursive"` + Limit int `json:"limit"` + Continuation string `json:"continuation"` +} + +type listCursor struct { + Root string `json:"root"` + Recursive bool `json:"recursive"` + Offset int `json:"offset"` +} + +type listResponse struct { + Keys []map[string]any `json:"keys"` + Continuation *string `json:"continuation"` +} + +func (service *server) get(params json.RawMessage) (map[string]any, error) { + client, err := service.requireClient() + if err != nil { + return nil, err + } + var request getRequest + if err := json.Unmarshal(params, &request); err != nil { + return nil, err + } + path := normalizePath(request.Key) + exists, _, err := client.Exists(path) + if err != nil { + return nil, err + } + if !exists { + return map[string]any{"found": false, "key": path, "value": nil, "metadata": nil}, nil + } + data, stat, err := client.Get(path) + if err == zk.ErrNoNode { + return map[string]any{"found": false, "key": path, "value": nil, "metadata": nil}, nil + } + if err != nil { + return nil, err + } + return map[string]any{ + "found": true, + "key": path, + "value": encodeValue(data), + "metadata": statMetadata(stat), + }, nil +} + +func (service *server) put(params json.RawMessage) (map[string]any, error) { + client, err := service.requireClient() + if err != nil { + return nil, err + } + var request putRequest + if err := json.Unmarshal(params, &request); err != nil { + return nil, err + } + path := normalizePath(request.Key) + if path == "/" { + return nil, errors.New("Root znode cannot be modified") + } + data, err := decodeValue(request.Value) + if err != nil { + return nil, err + } + writeMode := request.WriteMode + if writeMode == "" { + writeMode = "upsert" + } + switch writeMode { + case "create": + return createNode(client, path, data, request.CreateMode) + case "update": + stat, err := client.Set(path, data) + if err != nil { + return nil, err + } + return putResult(stat), nil + case "upsert": + exists, _, err := client.Exists(path) + if err != nil { + return nil, err + } + if exists { + stat, err := client.Set(path, data) + if err != nil { + return nil, err + } + return putResult(stat), nil + } + if err := createParents(client, parentPath(path)); err != nil { + return nil, err + } + createdPath, err := client.Create(path, data, 0) + if err != nil { + return nil, err + } + _, stat, err := client.Get(createdPath) + if err != nil { + return nil, err + } + return putResult(stat), nil + default: + return nil, fmt.Errorf("Unsupported writeMode: %s", writeMode) + } +} + +func createNode(client znodeClient, path string, data []byte, mode string) (map[string]any, error) { + flags, err := createFlags(mode) + if err != nil { + return nil, err + } + if err := createParents(client, parentPath(path)); err != nil { + return nil, err + } + createdPath, err := client.Create(path, data, flags) + if err != nil { + return nil, err + } + _, stat, err := client.Get(createdPath) + if err != nil { + return nil, err + } + result := putResult(stat) + result["key"] = createdPath + result["createdKey"] = createdPath + return result, nil +} + +func createFlags(mode string) (int32, error) { + switch mode { + case "", "persistent": + return 0, nil + case "ephemeral": + return zk.FlagEphemeral, nil + case "persistent_sequential": + return zk.FlagSequence, nil + case "ephemeral_sequential": + return zk.FlagEphemeral | zk.FlagSequence, nil + default: + return 0, fmt.Errorf("Unsupported createMode: %s", mode) + } +} + +func createParents(client znodeClient, parent string) error { + if parent == "/" { + return nil + } + current := "" + for _, segment := range strings.Split(strings.Trim(parent, "/"), "/") { + current = childPath(current, segment) + exists, _, err := client.Exists(current) + if err != nil { + return err + } + if exists { + continue + } + if _, err := client.Create(current, nil, 0); err != nil && err != zk.ErrNodeExists { + return err + } + } + return nil +} + +func (service *server) delete(params json.RawMessage) (map[string]any, error) { + client, err := service.requireClient() + if err != nil { + return nil, err + } + var request deleteRequest + if err := json.Unmarshal(params, &request); err != nil { + return nil, err + } + path := normalizePath(request.Key) + if path == "/" { + return nil, errors.New("Root znode cannot be deleted") + } + exists, _, err := client.Exists(path) + if err != nil { + return nil, err + } + if !exists { + return map[string]any{"deleted": 0}, nil + } + if !request.Recursive { + if err := client.Delete(path); err != nil { + if err == zk.ErrNoNode { + return map[string]any{"deleted": 0}, nil + } + return nil, err + } + return map[string]any{"deleted": 1}, nil + } + deleted, err := deleteSubtree(client, path) + if err != nil { + return nil, err + } + return map[string]any{"deleted": deleted}, nil +} + +func deleteSubtree(client znodeClient, path string) (int, error) { + children, _, err := client.Children(path) + if err == zk.ErrNoNode { + return 0, nil + } + if err != nil { + return 0, err + } + deleted := 0 + for _, child := range children { + count, err := deleteSubtree(client, childPath(path, child)) + if err != nil { + return deleted, err + } + deleted += count + } + if err := client.Delete(path); err != nil { + if err == zk.ErrNoNode { + return deleted, nil + } + return deleted, err + } + return deleted + 1, nil +} + +func (service *server) listPrefix(params json.RawMessage) (listResponse, error) { + client, err := service.requireClient() + if err != nil { + return listResponse{}, err + } + var request listRequest + if err := json.Unmarshal(params, &request); err != nil { + return listResponse{}, err + } + root := normalizePath(request.Prefix) + recursive := true + if request.Recursive != nil { + recursive = *request.Recursive + } + limit := request.Limit + if limit < 1 { + limit = defaultListLimit + } + cursor := listCursor{Root: root, Recursive: recursive} + if strings.TrimSpace(request.Continuation) != "" { + decoded, err := decodeCursor(request.Continuation) + if err != nil { + return listResponse{}, err + } + if decoded.Root != root || decoded.Recursive != recursive { + return listResponse{}, errors.New("Continuation does not match request") + } + cursor = decoded + } + exists, _, err := client.Exists(root) + if err != nil { + return listResponse{}, err + } + if !exists { + return listResponse{Keys: []map[string]any{}, Continuation: nil}, nil + } + var paths []string + if recursive { + paths, err = listRecursive(client, root) + } else { + paths, err = listDirectChildren(client, root) + } + if err != nil { + return listResponse{}, err + } + sort.Strings(paths) + offset := maxInt(0, cursor.Offset) + if offset > len(paths) { + offset = len(paths) + } + end := minInt(len(paths), offset+limit) + rows, err := service.rowsWithMetadata(client, paths[offset:end]) + if err != nil { + return listResponse{}, err + } + var continuation *string + if end < len(paths) { + encoded, err := encodeCursor(listCursor{Root: root, Recursive: recursive, Offset: end}) + if err != nil { + return listResponse{}, err + } + continuation = &encoded + } + return listResponse{Keys: rows, Continuation: continuation}, nil +} + +func listDirectChildren(client znodeClient, root string) ([]string, error) { + children, _, err := client.Children(root) + if err == zk.ErrNoNode { + return []string{}, nil + } + if err != nil { + return nil, err + } + sort.Strings(children) + paths := make([]string, 0, len(children)) + for _, child := range children { + paths = append(paths, childPath(root, child)) + } + return paths, nil +} + +func listRecursive(client znodeClient, root string) ([]string, error) { + result := make([]string, 0) + if err := collectRecursive(client, root, &result); err != nil { + return nil, err + } + return result, nil +} + +func collectRecursive(client znodeClient, root string, result *[]string) error { + children, _, err := client.Children(root) + if err == zk.ErrNoNode { + return nil + } + if err != nil { + return err + } + sort.Strings(children) + for _, child := range children { + path := childPath(root, child) + *result = append(*result, path) + if err := collectRecursive(client, path, result); err != nil { + return err + } + } + return nil +} + +func (service *server) rowsWithMetadata(client znodeClient, paths []string) ([]map[string]any, error) { + if len(paths) == 0 { + return []map[string]any{}, nil + } + rows := make([]map[string]any, len(paths)) + jobs := make(chan int) + workers := minInt(service.statLookupConcurrency, len(paths)) + var waitGroup sync.WaitGroup + var firstError error + var errorMutex sync.Mutex + for worker := 0; worker < workers; worker++ { + waitGroup.Add(1) + go func() { + defer waitGroup.Done() + for index := range jobs { + exists, stat, err := client.Exists(paths[index]) + if err == zk.ErrNoNode || !exists { + continue + } + if err != nil { + errorMutex.Lock() + if firstError == nil { + firstError = err + } + errorMutex.Unlock() + continue + } + row := statMetadata(stat) + row["key"] = paths[index] + rows[index] = row + } + }() + } + for index := range paths { + jobs <- index + } + close(jobs) + waitGroup.Wait() + if firstError != nil { + return nil, firstError + } + filtered := make([]map[string]any, 0, len(rows)) + for _, row := range rows { + if row != nil { + filtered = append(filtered, row) + } + } + return filtered, nil +} + +func statMetadata(stat *zk.Stat) map[string]any { + return map[string]any{ + "czxid": stat.Czxid, + "mzxid": stat.Mzxid, + "pzxid": stat.Pzxid, + "ctime": stat.Ctime, + "mtime": stat.Mtime, + "version": stat.Version, + "cversion": stat.Cversion, + "aversion": stat.Aversion, + "ephemeralOwner": stat.EphemeralOwner, + "dataLength": stat.DataLength, + "numChildren": stat.NumChildren, + "createRevision": stat.Czxid, + "modRevision": stat.Mzxid, + "valueSize": stat.DataLength, + } +} + +func putResult(stat *zk.Stat) map[string]any { + return map[string]any{"version": stat.Version, "mtime": stat.Mtime} +} + +func encodeValue(data []byte) valueObject { + if utf8.Valid(data) { + return valueObject{Encoding: "utf8", Data: string(data)} + } + return valueObject{Encoding: "base64", Data: base64.StdEncoding.EncodeToString(data)} +} + +func decodeValue(value valueObject) ([]byte, error) { + switch value.Encoding { + case "", "utf8": + return []byte(value.Data), nil + case "base64": + return base64.StdEncoding.DecodeString(value.Data) + default: + return nil, fmt.Errorf("Unsupported value encoding: %s", value.Encoding) + } +} + +func encodeCursor(cursor listCursor) (string, error) { + payload, err := json.Marshal(cursor) + if err != nil { + return "", err + } + return base64.StdEncoding.EncodeToString(payload), nil +} + +func decodeCursor(value string) (listCursor, error) { + payload, err := base64.StdEncoding.DecodeString(value) + if err != nil { + return listCursor{}, err + } + var cursor listCursor + if err := json.Unmarshal(payload, &cursor); err != nil { + return listCursor{}, err + } + return cursor, nil +} + +func normalizePath(value string) string { + if value == "" || value == "/" { + return "/" + } + normalized := value + if !strings.HasPrefix(normalized, "/") { + normalized = "/" + normalized + } + for len(normalized) > 1 && strings.HasSuffix(normalized, "/") { + normalized = strings.TrimSuffix(normalized, "/") + } + return normalized +} + +func childPath(parent, child string) string { + if parent == "" || parent == "/" { + return "/" + strings.Trim(child, "/") + } + return strings.TrimRight(parent, "/") + "/" + strings.Trim(child, "/") +} + +func parentPath(value string) string { + path := normalizePath(value) + separator := strings.LastIndex(path, "/") + if separator <= 0 { + return "/" + } + return path[:separator] +} diff --git a/agents/drivers/zookeeper/sasl.go b/agents/drivers/zookeeper/sasl.go new file mode 100644 index 0000000000..381d1ab098 --- /dev/null +++ b/agents/drivers/zookeeper/sasl.go @@ -0,0 +1,194 @@ +package main + +import ( + "encoding/binary" + "errors" + "fmt" + "io" + "net" + "sync" + "time" + + "github.com/beltran/gosasl" +) + +const ( + zooKeeperSASLOpcode = int32(102) + zooKeeperSASLXIDBase = int32(0x3fff0000) + zooKeeperSASLMaxRounds = 8 + zooKeeperMaximumFrameLen = 16 * 1024 * 1024 +) + +type saslDigestCredentials struct { + Username string + Password string +} + +type saslHandshakeConn struct { + net.Conn + timeout time.Duration + credentials saslDigestCredentials + mutex sync.Mutex + initialized bool + buffer []byte + error error +} + +var authenticateZooKeeperSASL = authenticateSASLDigest + +func newSASLHandshakeConn(connection net.Conn, timeout time.Duration, credentials saslDigestCredentials) net.Conn { + return &saslHandshakeConn{ + Conn: connection, + timeout: timeout, + credentials: credentials, + } +} + +func (connection *saslHandshakeConn) Read(payload []byte) (int, error) { + connection.mutex.Lock() + defer connection.mutex.Unlock() + if !connection.initialized { + connection.initialized = true + connection.buffer, connection.error = readZooKeeperFrame(connection.Conn) + if connection.error == nil { + connection.error = authenticateZooKeeperSASL(connection.Conn, connection.timeout, connection.credentials) + } + } + if connection.error != nil { + return 0, connection.error + } + if len(connection.buffer) > 0 { + read := copy(payload, connection.buffer) + connection.buffer = connection.buffer[read:] + return read, nil + } + return connection.Conn.Read(payload) +} + +func authenticateSASLDigest(connection net.Conn, timeout time.Duration, credentials saslDigestCredentials) error { + mechanism := gosasl.NewDigestMD5Mechanism("zookeeper", credentials.Username, credentials.Password) + saslClient := gosasl.NewSaslClient("zk-sasl-md5", mechanism) + defer saslClient.Dispose() + return negotiateSASLDigest(connection, timeout, saslClient) +} + +type saslClient interface { + Start() ([]byte, error) + Step(challenge []byte) ([]byte, error) + Complete() bool +} + +func negotiateSASLDigest(connection net.Conn, timeout time.Duration, saslClient saslClient) error { + if timeout <= 0 { + timeout = defaultConnectionTimeout + } + if err := connection.SetDeadline(time.Now().Add(timeout)); err != nil { + return err + } + defer connection.SetDeadline(time.Time{}) + + token, err := saslClient.Start() + if err != nil { + return fmt.Errorf("start ZooKeeper DIGEST-MD5 negotiation: %w", err) + } + for round := 0; round < zooKeeperSASLMaxRounds; round++ { + challenge, err := zooKeeperSASLRound(connection, zooKeeperSASLXIDBase+int32(round), token) + if err != nil { + return fmt.Errorf("ZooKeeper SASL round %d: %w", round+1, err) + } + if saslClient.Complete() { + if len(challenge) != 0 { + return errors.New("ZooKeeper sent an unexpected token after DIGEST-MD5 completion") + } + return nil + } + token, err = saslClient.Step(challenge) + if err != nil { + return fmt.Errorf("continue ZooKeeper DIGEST-MD5 negotiation at round %d: %w", round+1, err) + } + if saslClient.Complete() { + if len(token) != 0 { + return errors.New("ZooKeeper DIGEST-MD5 completed with an unexpected client token") + } + return nil + } + } + return fmt.Errorf("ZooKeeper DIGEST-MD5 negotiation exceeded %d rounds", zooKeeperSASLMaxRounds) +} + +func zooKeeperSASLRound(connection net.Conn, xid int32, token []byte) ([]byte, error) { + payload := make([]byte, 12+len(token)) + binary.BigEndian.PutUint32(payload[0:4], uint32(xid)) + binary.BigEndian.PutUint32(payload[4:8], uint32(zooKeeperSASLOpcode)) + binary.BigEndian.PutUint32(payload[8:12], uint32(len(token))) + copy(payload[12:], token) + if err := writeZooKeeperFrame(connection, payload); err != nil { + return nil, err + } + response, err := readZooKeeperFrame(connection) + if err != nil { + return nil, err + } + if len(response) < 20 { + return nil, errors.New("ZooKeeper SASL response is truncated") + } + responseXID := int32(binary.BigEndian.Uint32(response[4:8])) + if responseXID != xid { + return nil, fmt.Errorf("ZooKeeper SASL response xid %d does not match request xid %d", responseXID, xid) + } + errorCode := int32(binary.BigEndian.Uint32(response[16:20])) + if errorCode != 0 { + return nil, fmt.Errorf("ZooKeeper SASL server returned error %d", errorCode) + } + if len(response) < 24 { + return nil, errors.New("ZooKeeper SASL token is truncated") + } + tokenLength := int(int32(binary.BigEndian.Uint32(response[20:24]))) + if tokenLength < 0 || tokenLength > zooKeeperMaximumFrameLen || 24+tokenLength > len(response) { + return nil, fmt.Errorf("ZooKeeper SASL token length %d is invalid", tokenLength) + } + return append([]byte(nil), response[24:24+tokenLength]...), nil +} + +func readZooKeeperFrame(reader io.Reader) ([]byte, error) { + header := make([]byte, 4) + if _, err := io.ReadFull(reader, header); err != nil { + return nil, err + } + length := int(int32(binary.BigEndian.Uint32(header))) + if length < 0 || length > zooKeeperMaximumFrameLen { + return nil, fmt.Errorf("ZooKeeper frame length %d is invalid", length) + } + payload := make([]byte, length+4) + copy(payload, header) + if _, err := io.ReadFull(reader, payload[4:]); err != nil { + return nil, err + } + return payload, nil +} + +func writeZooKeeperFrame(writer io.Writer, payload []byte) error { + if len(payload) > zooKeeperMaximumFrameLen { + return fmt.Errorf("ZooKeeper frame length %d exceeds maximum %d", len(payload), zooKeeperMaximumFrameLen) + } + header := make([]byte, 4) + binary.BigEndian.PutUint32(header, uint32(len(payload))) + if err := writeAll(writer, header); err != nil { + return err + } + return writeAll(writer, payload) +} + +func writeAll(writer io.Writer, payload []byte) error { + for len(payload) > 0 { + written, err := writer.Write(payload) + if err != nil { + return err + } + if written <= 0 { + return io.ErrShortWrite + } + payload = payload[written:] + } + return nil +} diff --git a/agents/drivers/zookeeper/sasl_test.go b/agents/drivers/zookeeper/sasl_test.go new file mode 100644 index 0000000000..de4f065ea6 --- /dev/null +++ b/agents/drivers/zookeeper/sasl_test.go @@ -0,0 +1,156 @@ +package main + +import ( + "bytes" + "encoding/binary" + "errors" + "io" + "net" + "testing" + "time" +) + +func TestSASLHandshakeConnReplaysConnectResponse(t *testing.T) { + client, server := net.Pipe() + defer client.Close() + defer server.Close() + previous := authenticateZooKeeperSASL + t.Cleanup(func() { authenticateZooKeeperSASL = previous }) + called := make(chan saslDigestCredentials, 1) + authenticateZooKeeperSASL = func(_ net.Conn, _ time.Duration, credentials saslDigestCredentials) error { + called <- credentials + return nil + } + go func() { _ = writeZooKeeperFrame(server, []byte("connect-response")) }() + wrapper := newSASLHandshakeConn(client, time.Second, saslDigestCredentials{Username: "user", Password: "secret"}) + frame, err := readZooKeeperFrame(wrapper) + if err != nil || string(frame[4:]) != "connect-response" { + t.Fatalf("frame=%q err=%v", frame, err) + } + if credentials := <-called; credentials.Username != "user" || credentials.Password != "secret" { + t.Fatalf("credentials=%#v", credentials) + } +} + +func TestZooKeeperSASLRound(t *testing.T) { + client, server := net.Pipe() + defer client.Close() + defer server.Close() + serverError := make(chan error, 1) + go func() { + request, err := readZooKeeperFrame(server) + if err != nil { + serverError <- err + return + } + if len(request) != 4+12+len("client-token") || int32(binary.BigEndian.Uint32(request[4:8])) != 41 || int32(binary.BigEndian.Uint32(request[8:12])) != zooKeeperSASLOpcode || string(request[16:]) != "client-token" { + serverError <- errors.New("unexpected SASL request") + return + } + response := make([]byte, 20+4+len("server-token")) + binary.BigEndian.PutUint32(response[0:4], 41) + binary.BigEndian.PutUint32(response[16:20], uint32(len("server-token"))) + copy(response[20:], "server-token") + serverError <- writeZooKeeperFrame(server, response) + }() + token, err := zooKeeperSASLRound(client, 41, []byte("client-token")) + if err != nil || string(token) != "server-token" { + t.Fatalf("token=%q err=%v", token, err) + } + if err := <-serverError; err != nil { + t.Fatal(err) + } +} + +func TestZooKeeperSASLRoundRejectsWrongXID(t *testing.T) { + client, server := net.Pipe() + defer client.Close() + defer server.Close() + go func() { + _, _ = readZooKeeperFrame(server) + response := make([]byte, 24) + binary.BigEndian.PutUint32(response[0:4], 99) + _ = writeZooKeeperFrame(server, response) + }() + if _, err := zooKeeperSASLRound(client, 41, nil); err == nil || !bytes.Contains([]byte(err.Error()), []byte("does not match")) { + t.Fatalf("error=%v", err) + } +} + +func TestZooKeeperFrameValidation(t *testing.T) { + oversized := make([]byte, 4) + binary.BigEndian.PutUint32(oversized, zooKeeperMaximumFrameLen+1) + if _, err := readZooKeeperFrame(bytes.NewReader(oversized)); err == nil { + t.Fatal("oversized frame was accepted") + } + truncated := append([]byte{0, 0, 0, 4}, []byte{1, 2}...) + if _, err := readZooKeeperFrame(bytes.NewReader(truncated)); !errors.Is(err, io.ErrUnexpectedEOF) { + t.Fatalf("truncated error=%v", err) + } +} + +func TestSASLNegotiationStopsAfterRspauth(t *testing.T) { + clientConnection, serverConnection := net.Pipe() + defer clientConnection.Close() + defer serverConnection.Close() + client := &scriptedDigestClient{} + serverError := make(chan error, 1) + go func() { + for round, challenge := range [][]byte{[]byte("challenge"), []byte("rspauth")} { + request, err := readZooKeeperFrame(serverConnection) + if err != nil { + serverError <- err + return + } + expectedToken := "" + if round == 1 { + expectedToken = "response" + } + if string(request[16:]) != expectedToken { + serverError <- errors.New("unexpected client token") + return + } + response := make([]byte, 24+len(challenge)) + xid := int32(binary.BigEndian.Uint32(request[4:8])) + binary.BigEndian.PutUint32(response[0:4], uint32(xid)) + binary.BigEndian.PutUint32(response[16:20], uint32(len(challenge))) + copy(response[20:], challenge) + if err := writeZooKeeperFrame(serverConnection, response); err != nil { + serverError <- err + return + } + } + _ = serverConnection.SetReadDeadline(time.Now().Add(50 * time.Millisecond)) + buffer := make([]byte, 1) + if _, err := serverConnection.Read(buffer); err == nil { + serverError <- errors.New("client sent an unexpected third SASL round") + return + } + serverError <- nil + }() + if err := negotiateSASLDigest(clientConnection, time.Second, client); err != nil { + t.Fatal(err) + } + if err := <-serverError; err != nil { + t.Fatal(err) + } +} + +type scriptedDigestClient struct { + step int +} + +func (client *scriptedDigestClient) Start() ([]byte, error) { return nil, nil } + +func (client *scriptedDigestClient) Step(challenge []byte) ([]byte, error) { + client.step++ + if client.step == 1 && string(challenge) == "challenge" { + return []byte("response"), nil + } + if client.step == 2 && string(challenge) == "rspauth" { + return nil, nil + } + return nil, errors.New("unexpected challenge") +} + +func (client *scriptedDigestClient) Complete() bool { return client.step == 2 } diff --git a/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java b/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java deleted file mode 100644 index bb29e1a040..0000000000 --- a/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java +++ /dev/null @@ -1,971 +0,0 @@ -package com.dbx.agent.zookeeper; - -import com.dbx.agent.AgentProtocol; -import com.google.gson.*; -import org.apache.curator.framework.CuratorFramework; -import org.apache.curator.framework.CuratorFrameworkFactory; -import org.apache.curator.retry.ExponentialBackoffRetry; -import org.apache.zookeeper.CreateMode; -import org.apache.zookeeper.KeeperException; -import org.apache.zookeeper.client.ZKClientConfig; -import org.apache.zookeeper.data.Stat; - -import java.io.BufferedReader; -import java.io.InputStreamReader; -import java.net.InetSocketAddress; -import java.net.Socket; -import java.net.URLDecoder; -import java.nio.ByteBuffer; -import java.nio.charset.CharacterCodingException; -import java.nio.charset.CharsetDecoder; -import java.nio.charset.CodingErrorAction; -import java.nio.charset.StandardCharsets; -import java.util.*; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.ThreadFactory; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicInteger; -import javax.security.auth.login.AppConfigurationEntry; -import javax.security.auth.login.Configuration; - -public final class ZooKeeperAgent { - private static final Gson GSON = new GsonBuilder().serializeNulls().create(); - private static final int DEFAULT_LIMIT = 100; - private static final int DEFAULT_SESSION_TIMEOUT_MS = 30000; - private static final int DEFAULT_CONNECTION_TIMEOUT_MS = 15000; - private static final int DEFAULT_BASE_SLEEP_TIME_MS = 250; - private static final int DEFAULT_MAX_RETRIES = 2; - private static final int DEFAULT_PORT = 2181; - private static final int DEFAULT_PROBE_TIMEOUT_MS = 2000; - private static final String STAT_LOOKUP_CONCURRENCY_PROPERTY = "dbx.zookeeper.statLookupConcurrency"; - private static final String STAT_LOOKUP_CONCURRENCY_ENV = "DBX_ZOOKEEPER_STAT_LOOKUP_CONCURRENCY"; - private static final int DEFAULT_STAT_LOOKUP_CONCURRENCY = 16; - private static final int MIN_STAT_LOOKUP_CONCURRENCY = 1; - private static final int MAX_STAT_LOOKUP_CONCURRENCY = 64; - private static final String DEFAULT_AUTH_SCHEME = "digest"; - private static final String SASL_DIGEST_AUTH_SCHEME = "sasl_digest"; - private static final String SASL_LOGIN_CONTEXT_PREFIX = "DbxZooKeeperClient"; - private static final Object JAAS_LOCK = new Object(); - private static final AtomicInteger SASL_CONTEXT_SEQUENCE = new AtomicInteger(1); - private static final int STAT_LOOKUP_CONCURRENCY = configuredStatLookupConcurrency(); - private static final ExecutorService STAT_LOOKUP_EXECUTOR = Executors.newFixedThreadPool( - STAT_LOOKUP_CONCURRENCY, - daemonThreadFactory("dbx-zookeeper-stat-lookup-") - ); - private static final List CAPABILITIES = Collections.unmodifiableList(Arrays.asList( - AgentProtocol.CAPABILITY_CONNECT, - AgentProtocol.CAPABILITY_TEST_CONNECTION, - AgentProtocol.CAPABILITY_KV - )); - private static CuratorFramework client; - private static SaslLease activeSaslLease; - private static JaasRegistryConfiguration jaasRegistry; - - private ZooKeeperAgent() { - } - - private static Object handshakeResult() { - return new HandshakeResult(AgentProtocol.PROTOCOL_VERSION, AgentProtocol.PROTOCOL_VERSION, CAPABILITIES); - } - - private static Object connect(JsonObject params) throws Exception { - JsonObject connection = connectionObject(params); - SaslLease nextLease = SASL_DIGEST_AUTH_SCHEME.equals(authScheme(connection)) ? acquireSaslLease(connection) : null; - CuratorFramework nextClient = null; - try { - nextClient = buildClient(connection, nextLease == null ? null : nextLease.contextName); - startAndVerify(nextClient, connection); - closeClient(); - client = nextClient; - activeSaslLease = nextLease; - nextLease = null; - return Collections.singletonMap("ok", true); - } catch (Exception e) { - if (nextClient != null) { - nextClient.close(); - } - throw e; - } finally { - if (nextLease != null) { - nextLease.close(); - } - } - } - - private static Object testConnection(JsonObject params) throws Exception { - JsonObject connection = connectionObject(params); - SaslLease probeLease = SASL_DIGEST_AUTH_SCHEME.equals(authScheme(connection)) ? acquireSaslLease(connection) : null; - CuratorFramework probe = null; - try { - probe = buildClient(connection, probeLease == null ? null : probeLease.contextName); - startAndVerify(probe, connection); - return Collections.singletonMap("ok", true); - } finally { - try { - if (probe != null) { - probe.close(); - } - } finally { - if (probeLease != null) { - probeLease.close(); - } - } - } - } - - static CuratorFramework buildClient(JsonObject connection) { - return buildClient(connection, null); - } - - private static CuratorFramework buildClient(JsonObject connection, String saslContextName) { - if (hasTlsOptions(connection)) { - throw new IllegalArgumentException("ZooKeeper TLS is not supported"); - } - - String scheme = authScheme(connection); - if (!DEFAULT_AUTH_SCHEME.equals(scheme) && !SASL_DIGEST_AUTH_SCHEME.equals(scheme)) { - throw new IllegalArgumentException( - "Unsupported auth_scheme \"" + scheme + "\"; expected \"digest\" or \"sasl_digest\"" - ); - } - if (SASL_DIGEST_AUTH_SCHEME.equals(scheme)) { - saslCredentials(connection); - } - - int baseSleepTimeMs = intOrDefault(connection, "base_sleep_time_ms", DEFAULT_BASE_SLEEP_TIME_MS); - int maxRetries = intOrDefault(connection, "max_retries", DEFAULT_MAX_RETRIES); - int connectionTimeoutMs = intOrDefault(connection, "connection_timeout_ms", DEFAULT_CONNECTION_TIMEOUT_MS); - String connectString = reachableConnectString( - connectString(connection), - Math.min(DEFAULT_PROBE_TIMEOUT_MS, connectionTimeoutMs) - ); - CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder() - .connectString(connectString) - .sessionTimeoutMs(intOrDefault(connection, "session_timeout_ms", DEFAULT_SESSION_TIMEOUT_MS)) - .connectionTimeoutMs(connectionTimeoutMs) - .retryPolicy(new ExponentialBackoffRetry(baseSleepTimeMs, maxRetries)); - - if (SASL_DIGEST_AUTH_SCHEME.equals(scheme)) { - if (saslContextName == null) { - throw new IllegalStateException("SASL JAAS context was not initialized"); - } - ZKClientConfig clientConfig = new ZKClientConfig(); - clientConfig.setProperty(ZKClientConfig.ENABLE_CLIENT_SASL_KEY, "true"); - clientConfig.setProperty(ZKClientConfig.LOGIN_CONTEXT_NAME_KEY, saslContextName); - builder.zkClientConfig(clientConfig); - } - - String namespace = trimSlashes(stringOrEmpty(connection, "namespace")); - if (!namespace.isBlank()) { - builder.namespace(namespace); - } - - if (DEFAULT_AUTH_SCHEME.equals(scheme)) { - String username = stringOrEmpty(connection, "username"); - String password = stringOrEmpty(connection, "password"); - if (!username.isBlank()) { - builder.authorization(DEFAULT_AUTH_SCHEME, (username + ":" + password).getBytes(StandardCharsets.UTF_8)); - } - } - - return builder.build(); - } - - static String connectString(JsonObject connection) { - String configured = firstNonBlank( - stringOrNull(connection, "zookeeper_connect_string"), - stringOrNull(connection, "connect_string"), - stringOrNull(connection, "connection_string") - ); - if (configured != null) { - return configured; - } - return stringOrDefault(connection, "host", "127.0.0.1") + ":" - + intOrDefault(connection, "port", DEFAULT_PORT); - } - - static String reachableConnectString(String connectString, int probeTimeoutMs) { - int slash = connectString.indexOf('/'); - String hostsPart = slash >= 0 ? connectString.substring(0, slash) : connectString; - String chroot = slash >= 0 ? connectString.substring(slash) : ""; - - List hosts = new ArrayList<>(); - for (String item : hostsPart.split(",")) { - String trimmed = item.trim(); - if (!trimmed.isEmpty()) { - hosts.add(trimmed); - } - } - - List reachable = probeReachableHosts(hosts, probeTimeoutMs); - if (reachable.isEmpty()) { - throw new IllegalStateException( - "No reachable ZooKeeper server within " + probeTimeoutMs + "ms: " + hostsPart - ); - } - // Curator needs the complete ensemble to reconnect when nodes recover or fail over later. - return String.join(",", hosts) + chroot; - } - - private static List probeReachableHosts(List hosts, int probeTimeoutMs) { - if (hosts.isEmpty()) { - return Collections.emptyList(); - } - ExecutorService executor = Executors.newFixedThreadPool( - Math.min(hosts.size(), 8), - daemonThreadFactory("dbx-zookeeper-probe-") - ); - try { - List> probes = new ArrayList<>(); - for (String host : hosts) { - probes.add(executor.submit(() -> isEndpointReachable(host, probeTimeoutMs))); - } - // Probes run in parallel; the deadline covers socket timeout plus DNS slack. - long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(probeTimeoutMs + 500L); - List reachable = new ArrayList<>(); - for (int i = 0; i < hosts.size(); i++) { - long remainingMs = Math.max(1, TimeUnit.NANOSECONDS.toMillis(deadline - System.nanoTime())); - try { - if (probes.get(i).get(remainingMs, TimeUnit.MILLISECONDS)) { - reachable.add(hosts.get(i)); - } - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IllegalStateException("ZooKeeper reachability probe interrupted", e); - } catch (ExecutionException | TimeoutException e) { - // Treat probe failures as unreachable. - } - } - return reachable; - } finally { - executor.shutdownNow(); - } - } - - private static boolean isEndpointReachable(String endpoint, int probeTimeoutMs) { - try (Socket socket = new Socket()) { - socket.connect(parseEndpoint(endpoint), probeTimeoutMs); - return true; - } catch (Exception e) { - return false; - } - } - - static InetSocketAddress parseEndpoint(String endpoint) { - String host = endpoint; - int port = DEFAULT_PORT; - int bracketEnd = endpoint.indexOf(']'); - if (endpoint.startsWith("[") && bracketEnd > 0) { - host = endpoint.substring(1, bracketEnd); - if (bracketEnd + 1 < endpoint.length() && endpoint.charAt(bracketEnd + 1) == ':') { - port = Integer.parseInt(endpoint.substring(bracketEnd + 2)); - } - } else { - int colon = endpoint.lastIndexOf(':'); - if (colon >= 0 && endpoint.indexOf(':') == colon) { - host = endpoint.substring(0, colon); - port = Integer.parseInt(endpoint.substring(colon + 1)); - } - } - return new InetSocketAddress(host, port); - } - - private static void startAndVerify(CuratorFramework active, JsonObject connection) throws Exception { - active.start(); - int timeoutMs = intOrDefault(connection, "connection_timeout_ms", DEFAULT_CONNECTION_TIMEOUT_MS); - if (!active.blockUntilConnected(timeoutMs, TimeUnit.MILLISECONDS)) { - throw new IllegalStateException("Connection timed out"); - } - if (active.checkExists().forPath("/") == null) { - throw new IllegalStateException("Root znode is not readable"); - } - } - - private static Object listPrefix(JsonObject params) throws Exception { - CuratorFramework active = requireClient(); - String root = normalizePath(stringOrDefault(params, "prefix", "/")); - boolean recursive = boolOrDefault(params, "recursive", true); - int limit = Math.max(1, intOrDefault(params, "limit", DEFAULT_LIMIT)); - Cursor cursor = null; - String continuation = stringOrNull(params, "continuation"); - if (continuation != null && !continuation.isBlank()) { - cursor = decodeCursor(continuation); - if (!root.equals(cursor.root) || recursive != cursor.recursive) { - throw new IllegalArgumentException("Continuation does not match request"); - } - } - - Map result = new LinkedHashMap<>(); - List> keys = new ArrayList<>(); - Stat rootStat = active.checkExists().forPath(root); - if (rootStat == null) { - result.put("keys", keys); - result.put("continuation", null); - return result; - } - - List paths = recursive ? listRecursive(active, root) : listDirectChildren(active, root); - Collections.sort(paths); - int offset = cursor == null ? 0 : Math.max(0, cursor.offset); - int end = Math.min(paths.size(), offset + limit); - keys.addAll(listRowsWithMetadata(active, paths.subList(offset, end))); - - result.put("keys", keys); - result.put("continuation", end < paths.size() ? encodeCursor(new Cursor(root, recursive, end)) : null); - return result; - } - - private static Object get(JsonObject params) throws Exception { - CuratorFramework active = requireClient(); - String path = normalizePath(stringOrDefault(params, "key", "")); - Map result = new LinkedHashMap<>(); - Stat exists = active.checkExists().forPath(path); - if (exists == null) { - result.put("found", false); - result.put("key", path); - result.put("value", null); - result.put("metadata", null); - return result; - } - - Stat stat = new Stat(); - byte[] bytes = active.getData().storingStatIn(stat).forPath(path); - result.put("found", true); - result.put("key", path); - result.put("value", valueObject(bytes)); - result.put("metadata", statMetadata(stat)); - return result; - } - - private static Object put(JsonObject params) throws Exception { - CuratorFramework active = requireClient(); - String path = normalizePath(stringOrDefault(params, "key", "")); - requireMutablePath(path, "modified"); - byte[] bytes = parseValue(params.getAsJsonObject("value")); - - String writeMode = stringOrDefault(params, "writeMode", "upsert"); - return switch (writeMode) { - case "create" -> create(active, path, bytes, createMode(params)); - case "update" -> update(active, path, bytes); - case "upsert" -> upsert(active, path, bytes); - default -> throw new IllegalArgumentException("Unsupported writeMode: " + writeMode); - }; - } - - private static Object upsert(CuratorFramework active, String path, byte[] bytes) throws Exception { - Stat stat = active.checkExists().forPath(path); - if (stat == null) { - active.create().creatingParentsIfNeeded().forPath(path, bytes); - stat = active.checkExists().forPath(path); - } else { - stat = active.setData().forPath(path, bytes); - } - return putResult(stat); - } - - private static Object create(CuratorFramework active, String path, byte[] bytes, CreateMode createMode) throws Exception { - String createdPath = active.create() - .creatingParentsIfNeeded() - .withMode(createMode) - .forPath(path, bytes); - Stat stat = active.checkExists().forPath(createdPath); - Map result = putResult(stat); - result.put("key", createdPath); - result.put("createdKey", createdPath); - return result; - } - - private static Object update(CuratorFramework active, String path, byte[] bytes) throws Exception { - return putResult(active.setData().forPath(path, bytes)); - } - - private static Map putResult(Stat stat) { - Map result = new LinkedHashMap<>(); - result.put("version", stat.getVersion()); - result.put("mtime", stat.getMtime()); - return result; - } - - private static CreateMode createMode(JsonObject params) { - String createMode = stringOrDefault(params, "createMode", "persistent"); - return switch (createMode) { - case "persistent" -> CreateMode.PERSISTENT; - case "ephemeral" -> CreateMode.EPHEMERAL; - case "persistent_sequential" -> CreateMode.PERSISTENT_SEQUENTIAL; - case "ephemeral_sequential" -> CreateMode.EPHEMERAL_SEQUENTIAL; - default -> throw new IllegalArgumentException("Unsupported createMode: " + createMode); - }; - } - - private static Object delete(JsonObject params) throws Exception { - CuratorFramework active = requireClient(); - String path = normalizePath(stringOrDefault(params, "key", "")); - requireMutablePath(path, "deleted"); - boolean recursive = boolOrDefault(params, "recursive", false); - int deleted; - try { - if (active.checkExists().forPath(path) == null) { - deleted = 0; - } else if (recursive) { - deleted = countSubtree(active, path); - active.delete().deletingChildrenIfNeeded().forPath(path); - } else { - active.delete().forPath(path); - deleted = 1; - } - } catch (KeeperException.NoNodeException e) { - deleted = 0; - } - return Collections.singletonMap("deleted", deleted); - } - - private static Object dispatch(String method, JsonObject params) throws Exception { - return switch (method) { - case AgentProtocol.METHOD_HANDSHAKE -> handshakeResult(); - case AgentProtocol.METHOD_CONNECT -> connect(params); - case AgentProtocol.METHOD_TEST_CONNECTION -> testConnection(params); - case AgentProtocol.KV_METHOD_LIST_PREFIX -> listPrefix(params); - case AgentProtocol.KV_METHOD_GET -> get(params); - case AgentProtocol.KV_METHOD_PUT -> put(params); - case AgentProtocol.KV_METHOD_DELETE -> delete(params); - case AgentProtocol.METHOD_DISCONNECT -> { - closeClient(); - yield Collections.singletonMap("ok", true); - } - case AgentProtocol.METHOD_SHUTDOWN -> { - closeClient(); - STAT_LOOKUP_EXECUTOR.shutdownNow(); - System.exit(0); - yield Collections.singletonMap("ok", true); - } - default -> throw new IllegalArgumentException("Unknown method: " + method); - }; - } - - static String handleRequest(String line) { - JsonObject req = JsonParser.parseString(line).getAsJsonObject(); - JsonElement id = req.get("id"); - String method = req.get("method").getAsString(); - JsonObject params = req.has("params") && req.get("params").isJsonObject() - ? req.getAsJsonObject("params") - : new JsonObject(); - - JsonObject response = new JsonObject(); - response.addProperty("jsonrpc", "2.0"); - response.add("id", id); - - try { - Object result = dispatch(method, params); - response.add("result", GSON.toJsonTree(result)); - } catch (Exception e) { - JsonObject error = new JsonObject(); - error.addProperty("code", -1); - error.addProperty("message", e.getMessage() == null ? "Unknown error" : e.getMessage()); - response.add("error", error); - } - - return GSON.toJson(response); - } - - private static JsonObject connectionObject(JsonObject params) { - JsonElement connection = params.get("connection"); - return connection != null && connection.isJsonObject() ? connection.getAsJsonObject() : params; - } - - private static CuratorFramework requireClient() { - if (client == null) { - throw new IllegalStateException("Not connected"); - } - return client; - } - - private static void closeClient() { - try { - if (client != null) { - client.close(); - client = null; - } - } finally { - if (activeSaslLease != null) { - activeSaslLease.close(); - activeSaslLease = null; - } - } - } - - static String normalizePath(String path) { - String normalized = path == null ? "" : path; - if (normalized.isEmpty() || "/".equals(normalized)) { - return "/"; - } - if (!normalized.startsWith("/")) { - normalized = "/" + normalized; - } - while (normalized.length() > 1 && normalized.endsWith("/")) { - normalized = normalized.substring(0, normalized.length() - 1); - } - return normalized; - } - - private static String childPath(String parent, String child) { - return "/".equals(parent) ? "/" + child : parent + "/" + child; - } - - private static void requireMutablePath(String path, String action) { - if ("/".equals(path)) { - throw new IllegalArgumentException("Root znode cannot be " + action); - } - } - - private static Map statMetadata(Stat stat) { - Map metadata = new LinkedHashMap<>(); - long czxid = stat.getCzxid(); - long mzxid = stat.getMzxid(); - long pzxid = stat.getPzxid(); - int dataLength = stat.getDataLength(); - metadata.put("czxid", czxid); - metadata.put("mzxid", mzxid); - metadata.put("pzxid", pzxid); - metadata.put("ctime", stat.getCtime()); - metadata.put("mtime", stat.getMtime()); - metadata.put("version", stat.getVersion()); - metadata.put("cversion", stat.getCversion()); - metadata.put("aversion", stat.getAversion()); - metadata.put("ephemeralOwner", stat.getEphemeralOwner()); - metadata.put("dataLength", dataLength); - metadata.put("numChildren", stat.getNumChildren()); - metadata.put("createRevision", czxid); - metadata.put("modRevision", mzxid); - metadata.put("valueSize", dataLength); - return metadata; - } - - private static Map valueObject(byte[] bytes) { - Map value = new LinkedHashMap<>(); - // ZooKeeper represents znodes created without a data payload as null. - // Expose them as empty UTF-8 values instead of passing null to decoders. - if (bytes == null) { - bytes = new byte[0]; - } - String utf8 = strictUtf8(bytes); - if (utf8 != null) { - value.put("encoding", "utf8"); - value.put("data", utf8); - } else { - value.put("encoding", "base64"); - value.put("data", Base64.getEncoder().encodeToString(bytes)); - } - return value; - } - - private static byte[] parseValue(JsonObject value) { - if (value == null) { - return new byte[0]; - } - String encoding = stringOrDefault(value, "encoding", "utf8"); - String data = stringOrDefault(value, "data", ""); - if ("base64".equals(encoding)) { - return Base64.getDecoder().decode(data); - } - if (!"utf8".equals(encoding)) { - throw new IllegalArgumentException("Unsupported value encoding: " + encoding); - } - return data.getBytes(StandardCharsets.UTF_8); - } - - private static String strictUtf8(byte[] bytes) { - CharsetDecoder decoder = StandardCharsets.UTF_8.newDecoder() - .onMalformedInput(CodingErrorAction.REPORT) - .onUnmappableCharacter(CodingErrorAction.REPORT); - try { - return decoder.decode(ByteBuffer.wrap(bytes)).toString(); - } catch (CharacterCodingException e) { - return null; - } - } - - private static String encodeCursor(Cursor cursor) { - JsonObject object = new JsonObject(); - object.addProperty("root", cursor.root); - object.addProperty("recursive", cursor.recursive); - object.addProperty("offset", cursor.offset); - return Base64.getEncoder().encodeToString(GSON.toJson(object).getBytes(StandardCharsets.UTF_8)); - } - - private static Cursor decodeCursor(String continuation) { - String json = new String(Base64.getDecoder().decode(continuation), StandardCharsets.UTF_8); - JsonObject object = JsonParser.parseString(json).getAsJsonObject(); - return new Cursor( - object.get("root").getAsString(), - object.get("recursive").getAsBoolean(), - object.get("offset").getAsInt() - ); - } - - private static List listDirectChildren(CuratorFramework active, String root) throws Exception { - List paths = new ArrayList<>(); - List children = new ArrayList<>(active.getChildren().forPath(root)); - Collections.sort(children); - for (String child : children) { - paths.add(childPath(root, child)); - } - return paths; - } - - private static List listRecursive(CuratorFramework active, String root) throws Exception { - List result = new ArrayList<>(); - collectRecursive(active, root, result); - return result; - } - - private static List> listRowsWithMetadata(CuratorFramework active, List paths) throws Exception { - if (paths.isEmpty()) { - return Collections.emptyList(); - } - - try { - List>> futures = new ArrayList<>(); - for (String path : paths) { - futures.add(STAT_LOOKUP_EXECUTOR.submit(() -> rowWithMetadata(active, path))); - } - - List> rows = new ArrayList<>(); - for (Future> future : futures) { - Map row = future.get(); - if (row != null) { - rows.add(row); - } - } - return rows; - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw e; - } catch (ExecutionException e) { - Throwable cause = e.getCause(); - if (cause instanceof Exception exception) { - throw exception; - } - if (cause instanceof Error error) { - throw error; - } - throw new RuntimeException(cause); - } - } - - private static Map rowWithMetadata(CuratorFramework active, String path) throws Exception { - try { - Stat stat = active.checkExists().forPath(path); - if (stat == null) { - return null; - } - Map row = new LinkedHashMap<>(); - row.put("key", path); - row.putAll(statMetadata(stat)); - return row; - } catch (KeeperException.NoNodeException e) { - return null; - } - } - - private static void collectRecursive(CuratorFramework active, String path, List result) throws Exception { - List children; - try { - children = new ArrayList<>(active.getChildren().forPath(path)); - } catch (KeeperException.NoNodeException e) { - return; - } - Collections.sort(children); - for (String child : children) { - String childPath = childPath(path, child); - result.add(childPath); - collectRecursive(active, childPath, result); - } - } - - private static int countSubtree(CuratorFramework active, String path) throws Exception { - try { - if (active.checkExists().forPath(path) == null) { - return 0; - } - int count = 1; - for (String child : active.getChildren().forPath(path)) { - count += countSubtree(active, childPath(path, child)); - } - return count; - } catch (KeeperException.NoNodeException e) { - return 0; - } - } - - private static boolean hasTlsOptions(JsonObject connection) { - return boolOrDefault(connection, "ssl", false) - || firstNonBlank( - stringOrNull(connection, "ca_cert_path"), - stringOrNull(connection, "client_cert_path"), - stringOrNull(connection, "client_key_path"), - stringOrNull(connection, "cert_path"), - stringOrNull(connection, "key_path") - ) != null; - } - - static String authScheme(JsonObject connection) { - String configured = connectionOption(connection, "auth_scheme"); - return configured == null ? DEFAULT_AUTH_SCHEME : configured.toLowerCase(Locale.ROOT); - } - - private static String connectionOption(JsonObject connection, String key) { - String direct = stringOrNull(connection, key); - if (direct != null && !direct.isBlank()) { - return direct.trim(); - } - String params = stringOrEmpty(connection, "url_params").trim(); - if (params.startsWith("?")) { - params = params.substring(1); - } - for (String part : params.split("[&;]")) { - int separator = part.indexOf('='); - if (separator <= 0) { - continue; - } - String candidateKey = URLDecoder.decode(part.substring(0, separator), StandardCharsets.UTF_8); - if (key.equalsIgnoreCase(candidateKey.trim())) { - String value = URLDecoder.decode(part.substring(separator + 1), StandardCharsets.UTF_8).trim(); - return value.isEmpty() ? null : value; - } - } - return null; - } - - private static SaslCredentials saslCredentials(JsonObject connection) { - String username = stringOrEmpty(connection, "username").trim(); - String password = stringOrEmpty(connection, "password"); - if (username.isEmpty()) { - throw new IllegalArgumentException("username is required when auth_scheme = \"sasl_digest\""); - } - if (password.isEmpty()) { - throw new IllegalArgumentException("password is required when auth_scheme = \"sasl_digest\""); - } - return new SaslCredentials(username, password); - } - - private static SaslLease acquireSaslLease(JsonObject connection) { - SaslCredentials credentials = saslCredentials(connection); - synchronized (JAAS_LOCK) { - if (jaasRegistry == null) { - jaasRegistry = new JaasRegistryConfiguration(Configuration.getConfiguration()); - Configuration.setConfiguration(jaasRegistry); - } - String contextName = SASL_LOGIN_CONTEXT_PREFIX + SASL_CONTEXT_SEQUENCE.getAndIncrement(); - jaasRegistry.add(contextName, credentials); - return new SaslLease(contextName); - } - } - - private static void releaseSaslLease(String contextName) { - synchronized (JAAS_LOCK) { - if (jaasRegistry == null) { - return; - } - jaasRegistry.remove(contextName); - if (jaasRegistry.isEmpty()) { - Configuration base = jaasRegistry.base; - jaasRegistry = null; - Configuration.setConfiguration(base); - } - } - } - - private static String trimSlashes(String value) { - String trimmed = value.trim(); - while (trimmed.startsWith("/")) { - trimmed = trimmed.substring(1); - } - while (trimmed.endsWith("/")) { - trimmed = trimmed.substring(0, trimmed.length() - 1); - } - return trimmed; - } - - private static String stringOrNull(JsonObject object, String key) { - JsonElement element = object.get(key); - return element == null || element.isJsonNull() ? null : element.getAsString(); - } - - private static String stringOrEmpty(JsonObject object, String key) { - return stringOrDefault(object, key, ""); - } - - private static String stringOrDefault(JsonObject object, String key, String fallback) { - String value = stringOrNull(object, key); - return value == null ? fallback : value; - } - - private static int intOrDefault(JsonObject object, String key, int fallback) { - JsonElement element = object.get(key); - return element == null || element.isJsonNull() ? fallback : element.getAsInt(); - } - - private static boolean boolOrDefault(JsonObject object, String key, boolean fallback) { - JsonElement element = object.get(key); - return element == null || element.isJsonNull() ? fallback : element.getAsBoolean(); - } - - static int configuredStatLookupConcurrency(String propertyValue, String envValue) { - String configured = firstNonBlank(propertyValue, envValue); - if (configured == null) { - return DEFAULT_STAT_LOOKUP_CONCURRENCY; - } - try { - int parsed = Integer.parseInt(configured.trim()); - return Math.max(MIN_STAT_LOOKUP_CONCURRENCY, Math.min(MAX_STAT_LOOKUP_CONCURRENCY, parsed)); - } catch (NumberFormatException e) { - return DEFAULT_STAT_LOOKUP_CONCURRENCY; - } - } - - private static int configuredStatLookupConcurrency() { - return configuredStatLookupConcurrency( - System.getProperty(STAT_LOOKUP_CONCURRENCY_PROPERTY), - System.getenv(STAT_LOOKUP_CONCURRENCY_ENV) - ); - } - - private static ThreadFactory daemonThreadFactory(String prefix) { - AtomicInteger sequence = new AtomicInteger(1); - return runnable -> { - Thread thread = new Thread(runnable, prefix + sequence.getAndIncrement()); - thread.setDaemon(true); - return thread; - }; - } - - private static String firstNonBlank(String... values) { - for (String value : values) { - if (value != null && !value.isBlank()) { - return value; - } - } - return null; - } - - public static void main(String[] args) throws Exception { - // Entry point for the standalone DBX agent process. - System.setProperty("org.slf4j.simpleLogger.logFile", "System.err"); - System.out.println("{\"ready\":true}"); - System.out.flush(); - - BufferedReader reader = new BufferedReader(new InputStreamReader(System.in)); - while (true) { - String line = reader.readLine(); - if (line == null) { - break; - } - - System.out.println(handleRequest(line)); - System.out.flush(); - } - } - - private static final class SaslCredentials { - private final String username; - private final String password; - - private SaslCredentials(String username, String password) { - this.username = username; - this.password = password; - } - } - - private static final class JaasRegistryConfiguration extends Configuration { - private final Configuration base; - private final Map entries = new HashMap<>(); - - private JaasRegistryConfiguration(Configuration base) { - this.base = base; - } - - private void add(String contextName, SaslCredentials credentials) { - Map options = new HashMap<>(); - options.put("username", credentials.username); - options.put("password", credentials.password); - entries.put(contextName, new AppConfigurationEntry[] { - new AppConfigurationEntry( - "org.apache.zookeeper.server.auth.DigestLoginModule", - AppConfigurationEntry.LoginModuleControlFlag.REQUIRED, - Collections.unmodifiableMap(options) - ) - }); - } - - private void remove(String contextName) { - entries.remove(contextName); - } - - private boolean isEmpty() { - return entries.isEmpty(); - } - - @Override - public AppConfigurationEntry[] getAppConfigurationEntry(String name) { - synchronized (JAAS_LOCK) { - AppConfigurationEntry[] configured = entries.get(name); - if (configured != null) { - return configured.clone(); - } - } - return base == null ? null : base.getAppConfigurationEntry(name); - } - - @Override - public void refresh() { - if (base != null) { - base.refresh(); - } - } - } - - private static final class SaslLease implements AutoCloseable { - private final String contextName; - private boolean closed; - - private SaslLease(String contextName) { - this.contextName = contextName; - } - - @Override - public void close() { - if (!closed) { - releaseSaslLease(contextName); - closed = true; - } - } - } - - private static final class Cursor { - private final String root; - private final boolean recursive; - private final int offset; - - private Cursor(String root, boolean recursive, int offset) { - this.root = root; - this.recursive = recursive; - this.offset = offset; - } - } - - private static final class HandshakeResult { - private final int protocolVersion; - private final int agentProtocolVersion; - private final List capabilities; - - private HandshakeResult(int protocolVersion, int agentProtocolVersion, List capabilities) { - this.protocolVersion = protocolVersion; - this.agentProtocolVersion = agentProtocolVersion; - this.capabilities = capabilities; - } - } -} diff --git a/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java b/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java deleted file mode 100644 index 00ccae239a..0000000000 --- a/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java +++ /dev/null @@ -1,924 +0,0 @@ -package com.dbx.agent.zookeeper; - -import com.google.gson.JsonArray; -import com.google.gson.JsonElement; -import com.google.gson.JsonObject; -import com.google.gson.JsonParser; -import org.apache.curator.framework.CuratorFramework; -import org.apache.curator.framework.CuratorFrameworkFactory; -import org.apache.curator.retry.ExponentialBackoffRetry; -import org.apache.curator.test.InstanceSpec; -import org.apache.curator.test.TestingServer; -import org.apache.zookeeper.CreateMode; -import org.apache.zookeeper.ZooDefs; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Test; - -import java.io.IOException; -import java.net.InetSocketAddress; -import java.net.ServerSocket; -import java.net.Socket; -import java.util.ArrayList; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; - -final class ZooKeeperAgentTest { - @AfterEach - void disconnect() { - ZooKeeperAgent.handleRequest("{\"jsonrpc\":\"2.0\",\"id\":999,\"method\":\"disconnect\",\"params\":{}}"); - } - - @Test - void handshakeAdvertisesKvCapability() { - String response = ZooKeeperAgent.handleRequest( - "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"handshake\",\"params\":{}}" - ); - - JsonObject result = JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("result"); - - Assertions.assertEquals(1, result.get("protocolVersion").getAsInt()); - Assertions.assertTrue(result.getAsJsonArray("capabilities").contains(JsonParser.parseString("\"kv\""))); - Assertions.assertTrue(result.getAsJsonArray("capabilities").contains(JsonParser.parseString("\"connect\""))); - Assertions.assertTrue(result.getAsJsonArray("capabilities").contains(JsonParser.parseString("\"test_connection\""))); - } - - @Test - void kvMethodDispatchReturnsJsonRpcErrorWhenDisconnected() { - String response = ZooKeeperAgent.handleRequest( - "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"kv_get\",\"params\":{\"key\":\"/app/name\"}}" - ); - - JsonObject error = JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("error"); - - Assertions.assertEquals(-1, error.get("code").getAsInt()); - Assertions.assertEquals("Not connected", error.get("message").getAsString()); - } - - @Test - void unknownMethodReturnsJsonRpcError() { - String response = ZooKeeperAgent.handleRequest( - "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"missing_method\",\"params\":{}}" - ); - - JsonObject error = JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("error"); - - Assertions.assertEquals(-1, error.get("code").getAsInt()); - Assertions.assertEquals("Unknown method: missing_method", error.get("message").getAsString()); - } - - @Test - void connectStringUsesConfiguredValuesAndHostPortFallback() { - Assertions.assertEquals( - "zk-1:2181,zk-2:2181", - ZooKeeperAgent.connectString( - JsonParser.parseString("{\"connect_string\":\"zk-1:2181,zk-2:2181\"}").getAsJsonObject() - ) - ); - Assertions.assertEquals( - "zk-main:2181", - ZooKeeperAgent.connectString( - JsonParser.parseString("{\"zookeeper_connect_string\":\"zk-main:2181\"}").getAsJsonObject() - ) - ); - Assertions.assertEquals( - "127.0.0.1:2181", - ZooKeeperAgent.connectString(JsonParser.parseString("{}").getAsJsonObject()) - ); - Assertions.assertEquals( - "zk.local:2281", - ZooKeeperAgent.connectString(JsonParser.parseString("{\"host\":\"zk.local\",\"port\":2281}").getAsJsonObject()) - ); - } - - @Test - void buildClientRejectsTlsOptionsUntilTlsIsSupported() { - JsonObject connection = JsonParser.parseString( - "{\"ssl\":true,\"ca_cert_path\":\"ca.pem\",\"client_cert_path\":\"client.pem\",\"client_key_path\":\"client.key\"}" - ).getAsJsonObject(); - - IllegalArgumentException error = Assertions.assertThrows( - IllegalArgumentException.class, - () -> ZooKeeperAgent.buildClient(connection) - ); - - Assertions.assertEquals("ZooKeeper TLS is not supported", error.getMessage()); - } - - @Test - void buildClientRejectsUnsupportedAuthSchemeBeforeNetworkProbe() { - JsonObject connection = JsonParser.parseString( - "{\"auth_scheme\":\"world\",\"connect_string\":\"127.0.0.1:1\"}" - ).getAsJsonObject(); - - IllegalArgumentException error = Assertions.assertThrows( - IllegalArgumentException.class, - () -> ZooKeeperAgent.buildClient(connection) - ); - - Assertions.assertTrue(error.getMessage().contains("world"), error.getMessage()); - Assertions.assertFalse(error.getMessage().contains("No reachable"), error.getMessage()); - } - - @Test - void saslDigestRequiresUsernameAndPasswordBeforeNetworkProbe() { - JsonObject missingUsername = JsonParser.parseString( - "{\"auth_scheme\":\"sasl_digest\",\"password\":\"secret\",\"connect_string\":\"127.0.0.1:1\"}" - ).getAsJsonObject(); - JsonObject missingPassword = JsonParser.parseString( - "{\"auth_scheme\":\"sasl_digest\",\"username\":\"dbx\",\"connect_string\":\"127.0.0.1:1\"}" - ).getAsJsonObject(); - - IllegalArgumentException usernameError = Assertions.assertThrows( - IllegalArgumentException.class, - () -> ZooKeeperAgent.buildClient(missingUsername) - ); - IllegalArgumentException passwordError = Assertions.assertThrows( - IllegalArgumentException.class, - () -> ZooKeeperAgent.buildClient(missingPassword) - ); - - Assertions.assertTrue(usernameError.getMessage().contains("username"), usernameError.getMessage()); - Assertions.assertTrue(passwordError.getMessage().contains("password"), passwordError.getMessage()); - } - - @Test - void rejectsGenericSaslSchemeBeforeNetworkProbe() { - JsonObject connection = JsonParser.parseString( - "{\"auth_scheme\":\"sasl\",\"username\":\"dbx\",\"password\":\"secret\",\"connect_string\":\"127.0.0.1:1\"}" - ).getAsJsonObject(); - - IllegalArgumentException error = Assertions.assertThrows( - IllegalArgumentException.class, - () -> ZooKeeperAgent.buildClient(connection) - ); - - Assertions.assertTrue(error.getMessage().contains("sasl_digest"), error.getMessage()); - Assertions.assertFalse(error.getMessage().contains("No reachable"), error.getMessage()); - } - - @Test - void connectsToServerThatRequiresSaslDigestAuthenticationAndRestoresJaasConfiguration() throws Exception { - javax.security.auth.login.Configuration previousConfiguration = - javax.security.auth.login.Configuration.getConfiguration(); - String previousRequiredSasl = System.getProperty("zookeeper.sessionRequireClientSASLAuth"); - String previousServerConfig = System.getProperty("zookeeper.sasl.serverconfig"); - String previousLoginConfig = System.getProperty("java.security.auth.login.config"); - java.nio.file.Path jaas = java.nio.file.Files.createTempFile("dbx-zookeeper-server", ".conf"); - java.nio.file.Files.writeString(jaas, """ - Server { - org.apache.zookeeper.server.auth.DigestLoginModule required - user_dbx="secret"; - }; - """); - - try { - System.setProperty("java.security.auth.login.config", jaas.toString()); - System.setProperty("zookeeper.sasl.serverconfig", "Server"); - System.setProperty("zookeeper.sessionRequireClientSASLAuth", "true"); - javax.security.auth.login.Configuration.getConfiguration().refresh(); - - InstanceSpec spec = new InstanceSpec( - null, - -1, - -1, - -1, - true, - -1, - -1, - -1, - Map.of("authProvider.1", "org.apache.zookeeper.server.auth.SASLAuthenticationProvider") - ); - try (TestingServer server = new TestingServer(spec, true)) { - String connection = "{\"connection\":{\"url_params\":\"auth_scheme=sasl_digest\"," - + "\"username\":\"dbx\",\"password\":\"secret\"," - + "\"connect_string\":\"" + server.getConnectString() + "\"}}"; - - JsonObject connect = result(request(1, "connect", connection)); - Assertions.assertNotNull(connect, "SASL connect returned an RPC error"); - Assertions.assertTrue(connect.get("ok").getAsBoolean()); - - JsonObject probe = result(request(2, "test_connection", connection)); - Assertions.assertNotNull(probe, "concurrent SASL test_connection returned an RPC error"); - Assertions.assertTrue(probe.get("ok").getAsBoolean()); - - JsonObject root = result(request(3, "kv_get", "{\"key\":\"/\"}")); - Assertions.assertTrue(root.get("found").getAsBoolean(), "active SASL client must remain usable"); - - JsonObject disconnect = result(request(4, "disconnect", "{}")); - Assertions.assertTrue(disconnect.get("ok").getAsBoolean()); - - String wrongConnection = "{\"connection\":{\"url_params\":\"auth_scheme=sasl_digest\"," - + "\"username\":\"dbx\",\"password\":\"wrong\"," - + "\"connect_string\":\"" + server.getConnectString() + "\"}}"; - JsonObject authError = error(request(5, "test_connection", wrongConnection)); - Assertions.assertTrue( - authError.get("message").getAsString().toLowerCase(java.util.Locale.ROOT).contains("auth"), - authError.toString() - ); - } - - Assertions.assertSame( - previousConfiguration, - javax.security.auth.login.Configuration.getConfiguration(), - "client SASL configuration must be restored after authentication" - ); - } finally { - javax.security.auth.login.Configuration.setConfiguration(previousConfiguration); - restoreSystemProperty("zookeeper.sessionRequireClientSASLAuth", previousRequiredSasl); - restoreSystemProperty("zookeeper.sasl.serverconfig", previousServerConfig); - restoreSystemProperty("java.security.auth.login.config", previousLoginConfig); - javax.security.auth.login.Configuration.getConfiguration().refresh(); - java.nio.file.Files.deleteIfExists(jaas); - } - } - - @Test - void normalizePathStandardizesPublicPaths() { - Assertions.assertEquals("/", ZooKeeperAgent.normalizePath("")); - Assertions.assertEquals("/", ZooKeeperAgent.normalizePath("/")); - Assertions.assertEquals("/app", ZooKeeperAgent.normalizePath("app")); - Assertions.assertEquals("/app", ZooKeeperAgent.normalizePath("/app/")); - Assertions.assertEquals("/app/name", ZooKeeperAgent.normalizePath("/app/name")); - } - - @Test - void statLookupConcurrencyUsesConservativeDefaultsAndOverrides() { - Assertions.assertEquals(16, ZooKeeperAgent.configuredStatLookupConcurrency(null, null)); - Assertions.assertEquals(20, ZooKeeperAgent.configuredStatLookupConcurrency("20", "12")); - Assertions.assertEquals(12, ZooKeeperAgent.configuredStatLookupConcurrency(null, "12")); - Assertions.assertEquals(1, ZooKeeperAgent.configuredStatLookupConcurrency("0", null)); - Assertions.assertEquals(64, ZooKeeperAgent.configuredStatLookupConcurrency("128", null)); - Assertions.assertEquals(16, ZooKeeperAgent.configuredStatLookupConcurrency("not-a-number", null)); - } - - @Test - void connectAndTestConnectionWorkAgainstTestingServer() throws Exception { - try (TestingServer server = new TestingServer()) { - JsonObject connect = result(request( - 1, - "connect", - "{\"connection\":{\"connect_string\":\"" + server.getConnectString() + "\"}}" - )); - Assertions.assertTrue(connect.get("ok").getAsBoolean()); - - JsonObject test = result(request( - 2, - "test_connection", - "{\"connection\":{\"connect_string\":\"" + server.getConnectString() + "\"}}" - )); - Assertions.assertTrue(test.get("ok").getAsBoolean()); - } - } - - @Test - void defaultTimeoutAllowsSessionEstablishmentAfterFiveSeconds() throws Exception { - try (TestingServer server = new TestingServer(); - DelayedTcpProxy proxy = new DelayedTcpProxy(server.getConnectString(), 5500)) { - long startedAt = System.nanoTime(); - JsonObject connect = result(request( - 1, - "connect", - "{\"connection\":{\"connect_string\":\"" + proxy.getConnectString() + "\"}}" - )); - long elapsedMs = (System.nanoTime() - startedAt) / 1_000_000; - - Assertions.assertTrue(connect.get("ok").getAsBoolean()); - Assertions.assertTrue(elapsedMs >= 5000, "expected delayed handshake, took " + elapsedMs + "ms"); - } - } - - @Test - void connectFailsFastWhenAllServersAreUnreachable() { - long startedAt = System.nanoTime(); - JsonObject error = error(request( - 1, - "connect", - "{\"connection\":{\"connect_string\":\"127.0.0.1:1\",\"connection_timeout_ms\":15000}}" - )); - long elapsedMs = (System.nanoTime() - startedAt) / 1_000_000; - - Assertions.assertTrue( - error.get("message").getAsString().contains("No reachable ZooKeeper server"), - "expected fast unreachable error, got: " + error.get("message").getAsString() - ); - Assertions.assertTrue(elapsedMs < 4000, "expected fast failure, took " + elapsedMs + "ms"); - } - - @Test - void connectUsesReachableServerWithoutDroppingEnsembleMembers() throws Exception { - try (TestingServer server = new TestingServer()) { - JsonObject connect = result(request( - 1, - "connect", - "{\"connection\":{\"connect_string\":\"127.0.0.1:1," + server.getConnectString() + "\"}}" - )); - Assertions.assertTrue(connect.get("ok").getAsBoolean()); - - JsonObject get = result(request(2, "kv_get", "{\"key\":\"/\"}")); - Assertions.assertTrue(get.get("found").getAsBoolean()); - } - } - - @Test - void reachableConnectStringPreservesFullEnsembleAndChroot() throws Exception { - try (TestingServer server = new TestingServer()) { - String alive = server.getConnectString(); - String ensemble = "127.0.0.1:1," + alive + "/app"; - - Assertions.assertEquals( - ensemble, - ZooKeeperAgent.reachableConnectString(ensemble, 1000) - ); - Assertions.assertEquals(alive, ZooKeeperAgent.reachableConnectString(alive, 1000)); - - IllegalStateException error = Assertions.assertThrows( - IllegalStateException.class, - () -> ZooKeeperAgent.reachableConnectString("127.0.0.1:1/app", 1000) - ); - Assertions.assertTrue(error.getMessage().contains("No reachable ZooKeeper server")); - } - } - - @Test - void testConnectionDoesNotReplaceActiveClient() throws Exception { - try (TestingServer active = new TestingServer(); TestingServer probe = new TestingServer()) { - result(request(1, "connect", "{\"connection\":{\"connect_string\":\"" + active.getConnectString() + "\"}}")); - result(request( - 2, - "kv_put", - "{\"key\":\"/active\",\"value\":{\"encoding\":\"utf8\",\"data\":\"still-here\"}}" - )); - - result(request(3, "test_connection", "{\"connection\":{\"connect_string\":\"" + probe.getConnectString() + "\"}}")); - - JsonObject get = result(request(4, "kv_get", "{\"key\":\"/active\"}")); - Assertions.assertTrue(get.get("found").getAsBoolean()); - Assertions.assertEquals("still-here", get.getAsJsonObject("value").get("data").getAsString()); - } - } - - @Test - void kvPutCreatesMissingParentsAndGetReadsUtf8Value() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - JsonObject put = result(request( - 2, - "kv_put", - "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}" - )); - Assertions.assertTrue(put.has("version")); - - JsonObject get = result(request(3, "kv_get", "{\"key\":\"/app/name\"}")); - - Assertions.assertTrue(get.get("found").getAsBoolean()); - Assertions.assertEquals("utf8", get.getAsJsonObject("value").get("encoding").getAsString()); - Assertions.assertEquals("dbx", get.getAsJsonObject("value").get("data").getAsString()); - JsonObject metadata = get.getAsJsonObject("metadata"); - Assertions.assertEquals(0, metadata.get("numChildren").getAsInt()); - Assertions.assertEquals(metadata.get("czxid").getAsLong(), metadata.get("createRevision").getAsLong()); - Assertions.assertEquals(metadata.get("mzxid").getAsLong(), metadata.get("modRevision").getAsLong()); - Assertions.assertEquals(metadata.get("dataLength").getAsInt(), metadata.get("valueSize").getAsInt()); - } - } - - @Test - void kvGetReadsZnodeCreatedWithoutDataAsEmptyUtf8() throws Exception { - try (TestingServer server = new TestingServer(); - CuratorFramework curator = CuratorFrameworkFactory.newClient( - server.getConnectString(), - new ExponentialBackoffRetry(100, 3) - )) { - curator.start(); - curator.getZookeeperClient().getZooKeeper().create( - "/empty", - null, - ZooDefs.Ids.OPEN_ACL_UNSAFE, - CreateMode.PERSISTENT - ); - connect(server); - - JsonObject get = result(request(3, "kv_get", "{\"key\":\"/empty\"}")); - - Assertions.assertTrue(get.get("found").getAsBoolean()); - Assertions.assertEquals("utf8", get.getAsJsonObject("value").get("encoding").getAsString()); - Assertions.assertEquals("", get.getAsJsonObject("value").get("data").getAsString()); - Assertions.assertEquals(0, get.getAsJsonObject("metadata").get("valueSize").getAsInt()); - } - } - - @Test - void kvPutUpdatesExistingZnode() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); - result(request( - 3, - "kv_put", - "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx-updated\"}}" - )); - - JsonObject value = result(request(4, "kv_get", "{\"key\":\"/app/name\"}")).getAsJsonObject("value"); - Assertions.assertEquals("dbx-updated", value.get("data").getAsString()); - } - } - - @Test - void kvPutCreatePersistentCreatesOnlyMissingZnode() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - JsonObject created = result(request( - 2, - "kv_put", - "{\"key\":\"/app/created\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}," - + "\"writeMode\":\"create\",\"createMode\":\"persistent\"}" - )); - JsonObject get = result(request(3, "kv_get", "{\"key\":\"/app/created\"}")); - JsonObject duplicate = error(request( - 4, - "kv_put", - "{\"key\":\"/app/created\",\"value\":{\"encoding\":\"utf8\",\"data\":\"again\"}," - + "\"writeMode\":\"create\",\"createMode\":\"persistent\"}" - )); - - Assertions.assertEquals("/app/created", requiredString(created, "key")); - Assertions.assertEquals("/app/created", requiredString(created, "createdKey")); - Assertions.assertEquals("dbx", get.getAsJsonObject("value").get("data").getAsString()); - Assertions.assertTrue(duplicate.get("message").getAsString().contains("NodeExists")); - } - } - - @Test - void kvPutCreateEphemeralCreatesEphemeralZnode() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - JsonObject created = result(request( - 2, - "kv_put", - "{\"key\":\"/session/node\",\"value\":{\"encoding\":\"utf8\",\"data\":\"temporary\"}," - + "\"writeMode\":\"create\",\"createMode\":\"ephemeral\"}" - )); - JsonObject get = result(request(3, "kv_get", "{\"key\":\"/session/node\"}")); - - Assertions.assertEquals("/session/node", requiredString(created, "createdKey")); - Assertions.assertEquals("temporary", get.getAsJsonObject("value").get("data").getAsString()); - Assertions.assertTrue(get.getAsJsonObject("metadata").get("ephemeralOwner").getAsLong() > 0); - } - } - - @Test - void kvPutCreatePersistentSequentialReturnsGeneratedPath() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - JsonObject created = result(request( - 2, - "kv_put", - "{\"key\":\"/jobs/job-\",\"value\":{\"encoding\":\"utf8\",\"data\":\"one\"}," - + "\"writeMode\":\"create\",\"createMode\":\"persistent_sequential\"}" - )); - String createdKey = requiredString(created, "createdKey"); - JsonObject get = result(request(3, "kv_get", "{\"key\":\"" + createdKey + "\"}")); - - Assertions.assertEquals(createdKey, requiredString(created, "key")); - Assertions.assertTrue(createdKey.startsWith("/jobs/job-")); - Assertions.assertTrue(createdKey.length() > "/jobs/job-".length()); - Assertions.assertEquals("one", get.getAsJsonObject("value").get("data").getAsString()); - Assertions.assertEquals(0, get.getAsJsonObject("metadata").get("ephemeralOwner").getAsLong()); - } - } - - @Test - void kvPutCreateEphemeralSequentialReturnsGeneratedEphemeralPath() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - JsonObject created = result(request( - 2, - "kv_put", - "{\"key\":\"/sessions/member-\",\"value\":{\"encoding\":\"utf8\",\"data\":\"online\"}," - + "\"writeMode\":\"create\",\"createMode\":\"ephemeral_sequential\"}" - )); - String createdKey = requiredString(created, "createdKey"); - JsonObject get = result(request(3, "kv_get", "{\"key\":\"" + createdKey + "\"}")); - - Assertions.assertEquals(createdKey, requiredString(created, "key")); - Assertions.assertTrue(createdKey.startsWith("/sessions/member-")); - Assertions.assertEquals("online", get.getAsJsonObject("value").get("data").getAsString()); - Assertions.assertTrue(get.getAsJsonObject("metadata").get("ephemeralOwner").getAsLong() > 0); - } - } - - @Test - void kvPutUpdateOnlyUpdatesExistingZnode() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - JsonObject missing = error(request( - 2, - "kv_put", - "{\"key\":\"/app/missing\",\"value\":{\"encoding\":\"utf8\",\"data\":\"x\"},\"writeMode\":\"update\"}" - )); - result(request( - 3, - "kv_put", - "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"before\"},\"writeMode\":\"create\"}" - )); - JsonObject updated = result(request( - 4, - "kv_put", - "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"after\"},\"writeMode\":\"update\"}" - )); - JsonObject value = result(request(5, "kv_get", "{\"key\":\"/app/name\"}")).getAsJsonObject("value"); - - Assertions.assertTrue(missing.get("message").getAsString().contains("NoNode")); - Assertions.assertTrue(updated.get("version").getAsInt() > 0); - Assertions.assertEquals("after", value.get("data").getAsString()); - } - } - - @Test - void kvGetReturnsBase64ForBinaryData() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - result(request(2, "kv_put", "{\"key\":\"/bin\",\"value\":{\"encoding\":\"base64\",\"data\":\"/wAB\"}}")); - - JsonObject value = result(request(3, "kv_get", "{\"key\":\"/bin\"}")).getAsJsonObject("value"); - Assertions.assertEquals("base64", value.get("encoding").getAsString()); - Assertions.assertEquals("/wAB", value.get("data").getAsString()); - } - } - - @Test - void kvPutRejectsRootZnode() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - JsonObject error = error(request( - 2, - "kv_put", - "{\"key\":\"/\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}" - )); - - Assertions.assertEquals("Root znode cannot be modified", error.get("message").getAsString()); - } - } - - @Test - void kvDeleteReturnsZeroForMissingZnode() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - JsonObject delete = result(request(4, "kv_delete", "{\"key\":\"/missing\"}")); - - Assertions.assertEquals(0, delete.get("deleted").getAsInt()); - } - } - - @Test - void kvDeleteRemovesLeafZnode() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - result(request(2, "kv_put", "{\"key\":\"/leaf\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); - - JsonObject delete = result(request(3, "kv_delete", "{\"key\":\"/leaf\"}")); - JsonObject get = result(request(4, "kv_get", "{\"key\":\"/leaf\"}")); - - Assertions.assertEquals(1, delete.get("deleted").getAsInt()); - Assertions.assertFalse(get.get("found").getAsBoolean()); - } - } - - @Test - void kvDeleteRejectsNonRecursiveDeleteOfNonEmptyZnode() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); - - JsonObject error = error(request(3, "kv_delete", "{\"key\":\"/app\"}")); - - Assertions.assertTrue(error.get("message").getAsString().contains("not empty")); - } - } - - @Test - void kvDeleteRecursiveDeletesFullSubtreeAndCountsZnodes() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); - result(request(3, "kv_put", "{\"key\":\"/app/config\",\"value\":{\"encoding\":\"utf8\",\"data\":\"cfg\"}}")); - - JsonObject delete = result(request(4, "kv_delete", "{\"key\":\"/app\",\"recursive\":true}")); - JsonObject get = result(request(5, "kv_get", "{\"key\":\"/app/name\"}")); - - Assertions.assertEquals(3, delete.get("deleted").getAsInt()); - Assertions.assertFalse(get.get("found").getAsBoolean()); - } - } - - @Test - void kvDeleteRejectsRootZnode() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - JsonObject error = error(request(2, "kv_delete", "{\"key\":\"/\",\"recursive\":true}")); - - Assertions.assertEquals("Root znode cannot be deleted", error.get("message").getAsString()); - } - } - - @Test - void listPrefixListsDirectChildrenOfRootWhenRecursiveIsFalse() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); - result(request(3, "kv_put", "{\"key\":\"/config\",\"value\":{\"encoding\":\"utf8\",\"data\":\"cfg\"}}")); - - JsonObject list = result(request(5, "kv_list_prefix", "{\"prefix\":\"/\",\"recursive\":false}")); - - Assertions.assertEquals(List.of("/app", "/config", "/zookeeper"), listedKeys(list)); - Assertions.assertTrue(list.get("continuation").isJsonNull()); - } - } - - @Test - void listPrefixKeepsDirectChildMetadataForLazyExpansion() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); - result(request(3, "kv_put", "{\"key\":\"/config\",\"value\":{\"encoding\":\"utf8\",\"data\":\"cfg\"}}")); - - JsonObject list = result(request(5, "kv_list_prefix", "{\"prefix\":\"/\",\"recursive\":false}")); - JsonObject app = listedRow(list, "/app"); - JsonObject config = listedRow(list, "/config"); - - Assertions.assertEquals(1, app.get("numChildren").getAsInt()); - Assertions.assertEquals(0, config.get("numChildren").getAsInt()); - Assertions.assertEquals(3, config.get("dataLength").getAsInt()); - } - } - - @Test - void listPrefixDefaultsToRecursiveForDbxKvBrowser() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); - result(request(3, "kv_put", "{\"key\":\"/config\",\"value\":{\"encoding\":\"utf8\",\"data\":\"cfg\"}}")); - - JsonObject list = result(request(5, "kv_list_prefix", "{\"prefix\":\"/\"}")); - - Assertions.assertEquals( - List.of("/app", "/app/name", "/config", "/zookeeper", "/zookeeper/config", "/zookeeper/quota"), - listedKeys(list) - ); - Assertions.assertTrue(list.get("continuation").isJsonNull()); - } - } - - @Test - void listPrefixListsRecursiveDescendantsOfRoot() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - result(request(2, "kv_put", "{\"key\":\"/app/name\",\"value\":{\"encoding\":\"utf8\",\"data\":\"dbx\"}}")); - result(request(3, "kv_put", "{\"key\":\"/config\",\"value\":{\"encoding\":\"utf8\",\"data\":\"cfg\"}}")); - - JsonObject list = result(request(6, "kv_list_prefix", "{\"prefix\":\"/\",\"recursive\":true}")); - - Assertions.assertEquals( - List.of("/app", "/app/name", "/config", "/zookeeper", "/zookeeper/config", "/zookeeper/quota"), - listedKeys(list) - ); - } - } - - @Test - void listPrefixReturnsEmptyResultForMissingPrefix() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - - JsonObject list = result(request(7, "kv_list_prefix", "{\"prefix\":\"/missing\"}")); - - Assertions.assertEquals(List.of(), listedKeys(list)); - Assertions.assertTrue(list.get("continuation").isJsonNull()); - } - } - - @Test - void listPrefixPaginatesBySortedPath() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - result(request(2, "kv_put", "{\"key\":\"/c\",\"value\":{\"encoding\":\"utf8\",\"data\":\"c\"}}")); - result(request(3, "kv_put", "{\"key\":\"/a\",\"value\":{\"encoding\":\"utf8\",\"data\":\"a\"}}")); - result(request(4, "kv_put", "{\"key\":\"/b\",\"value\":{\"encoding\":\"utf8\",\"data\":\"b\"}}")); - - JsonObject first = result(request(5, "kv_list_prefix", "{\"prefix\":\"/\",\"limit\":2}")); - String continuation = first.get("continuation").getAsString(); - JsonObject second = result(request( - 6, - "kv_list_prefix", - "{\"prefix\":\"/\",\"limit\":2,\"continuation\":\"" + continuation + "\"}" - )); - String secondContinuation = second.get("continuation").getAsString(); - JsonObject third = result(request( - 7, - "kv_list_prefix", - "{\"prefix\":\"/\",\"limit\":2,\"continuation\":\"" + secondContinuation + "\"}" - )); - - Assertions.assertEquals(List.of("/a", "/b"), listedKeys(first)); - Assertions.assertFalse(continuation.isBlank()); - Assertions.assertEquals(List.of("/c", "/zookeeper"), listedKeys(second)); - Assertions.assertFalse(secondContinuation.isBlank()); - Assertions.assertEquals(List.of("/zookeeper/config", "/zookeeper/quota"), listedKeys(third)); - Assertions.assertTrue(third.get("continuation").isJsonNull()); - } - } - - @Test - void listPrefixPaginatesDirectChildrenWhenRecursiveIsFalse() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - result(request(2, "kv_put", "{\"key\":\"/app/c/deep\",\"value\":{\"encoding\":\"utf8\",\"data\":\"c\"}}")); - result(request(3, "kv_put", "{\"key\":\"/app/a/deep\",\"value\":{\"encoding\":\"utf8\",\"data\":\"a\"}}")); - result(request(4, "kv_put", "{\"key\":\"/app/b/deep\",\"value\":{\"encoding\":\"utf8\",\"data\":\"b\"}}")); - - JsonObject first = result(request(5, "kv_list_prefix", "{\"prefix\":\"/app\",\"recursive\":false,\"limit\":2}")); - String continuation = first.get("continuation").getAsString(); - JsonObject second = result(request( - 6, - "kv_list_prefix", - "{\"prefix\":\"/app\",\"recursive\":false,\"limit\":2,\"continuation\":\"" + continuation + "\"}" - )); - - Assertions.assertEquals(List.of("/app/a", "/app/b"), listedKeys(first)); - Assertions.assertFalse(continuation.isBlank()); - Assertions.assertEquals(List.of("/app/c"), listedKeys(second)); - Assertions.assertTrue(second.get("continuation").isJsonNull()); - } - } - - @Test - void listPrefixRejectsContinuationForDifferentRequest() throws Exception { - try (TestingServer server = new TestingServer()) { - connect(server); - result(request(2, "kv_put", "{\"key\":\"/app/a\",\"value\":{\"encoding\":\"utf8\",\"data\":\"a\"}}")); - result(request(3, "kv_put", "{\"key\":\"/app/b\",\"value\":{\"encoding\":\"utf8\",\"data\":\"b\"}}")); - - JsonObject first = result(request(4, "kv_list_prefix", "{\"prefix\":\"/app\",\"limit\":1}")); - String continuation = first.get("continuation").getAsString(); - JsonObject error = error(request( - 5, - "kv_list_prefix", - "{\"prefix\":\"/\",\"limit\":1,\"continuation\":\"" + continuation + "\"}" - )); - - Assertions.assertEquals("Continuation does not match request", error.get("message").getAsString()); - } - } - - private static JsonObject connect(TestingServer server) { - return result(request(1, "connect", "{\"connection\":{\"connect_string\":\"" + server.getConnectString() + "\"}}")); - } - - private static String request(int id, String method, String params) { - return ZooKeeperAgent.handleRequest( - "{\"jsonrpc\":\"2.0\",\"id\":" + id + ",\"method\":\"" + method + "\",\"params\":" + params + "}" - ); - } - - private static JsonObject result(String response) { - return JsonParser.parseString(response).getAsJsonObject().getAsJsonObject("result"); - } - - private static JsonObject error(String response) { - JsonObject object = JsonParser.parseString(response).getAsJsonObject(); - Assertions.assertTrue(object.has("error"), "expected JSON-RPC error but got " + response); - return object.getAsJsonObject("error"); - } - - private static String requiredString(JsonObject object, String key) { - Assertions.assertTrue(object.has(key), "expected result to contain " + key); - return object.get(key).getAsString(); - } - - private static void restoreSystemProperty(String key, String value) { - if (value == null) { - System.clearProperty(key); - } else { - System.setProperty(key, value); - } - } - - private static List listedKeys(JsonObject listResult) { - List keys = new ArrayList<>(); - JsonArray rows = listResult.getAsJsonArray("keys"); - for (JsonElement row : rows) { - keys.add(row.getAsJsonObject().get("key").getAsString()); - } - return keys; - } - - private static JsonObject listedRow(JsonObject listResult, String key) { - JsonArray rows = listResult.getAsJsonArray("keys"); - for (JsonElement row : rows) { - JsonObject object = row.getAsJsonObject(); - if (key.equals(object.get("key").getAsString())) { - return object; - } - } - Assertions.fail("expected listed row for " + key); - return new JsonObject(); - } - - private static final class DelayedTcpProxy implements AutoCloseable { - private final ServerSocket listener; - private final InetSocketAddress target; - private final long delayMs; - private final ExecutorService executor; - private final Set sockets = ConcurrentHashMap.newKeySet(); - - private DelayedTcpProxy(String targetConnectString, long delayMs) throws IOException { - this.listener = new ServerSocket(0); - this.target = ZooKeeperAgent.parseEndpoint(targetConnectString); - this.delayMs = delayMs; - this.executor = Executors.newCachedThreadPool(task -> { - Thread thread = new Thread(task, "dbx-zookeeper-delayed-proxy"); - thread.setDaemon(true); - return thread; - }); - executor.submit(this::acceptConnections); - } - - private String getConnectString() { - return "127.0.0.1:" + listener.getLocalPort(); - } - - private void acceptConnections() { - while (!listener.isClosed()) { - try { - Socket client = listener.accept(); - sockets.add(client); - executor.submit(() -> forwardAfterDelay(client)); - } catch (IOException error) { - if (!listener.isClosed()) { - throw new IllegalStateException("Delayed proxy failed to accept a connection", error); - } - } - } - } - - private void forwardAfterDelay(Socket client) { - try { - Thread.sleep(delayMs); - Socket server = new Socket(); - sockets.add(server); - server.connect(target); - executor.submit(() -> copyUntilClosed(client, server)); - copyUntilClosed(server, client); - } catch (InterruptedException error) { - Thread.currentThread().interrupt(); - closeSocket(client); - } catch (IOException error) { - closeSocket(client); - } - } - - private void copyUntilClosed(Socket source, Socket destination) { - try { - source.getInputStream().transferTo(destination.getOutputStream()); - } catch (IOException ignored) { - // Closing either side is expected when a probe or client disconnects. - } finally { - closeSocket(source); - closeSocket(destination); - } - } - - private void closeSocket(Socket socket) { - sockets.remove(socket); - try { - socket.close(); - } catch (IOException ignored) { - // Best effort cleanup for test sockets. - } - } - - @Override - public void close() throws IOException { - listener.close(); - for (Socket socket : sockets) { - closeSocket(socket); - } - executor.shutdownNow(); - } - } -} diff --git a/agents/go-common/gosasl/sasl.go b/agents/go-common/gosasl/sasl.go index 43356a3492..d88e340c03 100644 --- a/agents/go-common/gosasl/sasl.go +++ b/agents/go-common/gosasl/sasl.go @@ -374,15 +374,19 @@ func (m *DigestMD5Mechanism) step(challenge []byte) ([]byte, error) { a2String := "AUTHENTICATE:" + digestUri maxBuf := "" - if c["qop"] != AUTH { + if m.auth != AUTH { a2String += ":00000000000000000000000000000000" maxBuf = ",maxbuf=16777215" } // Set nonce count nc nc := fmt.Sprintf("%08x", m.nonceCount) + charset := "" + if strings.EqualFold(c["charset"], "utf-8") { + charset = ",charset=utf-8" + } // Create final response sent to server - resp := "qop=" + c["qop"] + ",realm=" + strconv.Quote(c["realm"]) + ",username=" + strconv.Quote(m.username) + ",nonce=" + strconv.Quote(m.nonce) + - ",cnonce=" + strconv.Quote(m.cnonce) + ",nc=" + nc + ",digest-uri=" + strconv.Quote(digestUri) + ",response=" + m.getHash(digestUri, a2String, c) + maxBuf + resp := "qop=" + m.auth + ",realm=" + strconv.Quote(c["realm"]) + ",username=" + strconv.Quote(m.username) + ",nonce=" + strconv.Quote(m.nonce) + + ",cnonce=" + strconv.Quote(m.cnonce) + ",nc=" + nc + ",digest-uri=" + strconv.Quote(digestUri) + ",response=" + m.getHash(digestUri, a2String, c) + charset + maxBuf return []byte(resp), nil } diff --git a/agents/go-common/gosasl/sasl_test.go b/agents/go-common/gosasl/sasl_test.go index 06dc7945ab..54bf0df076 100644 --- a/agents/go-common/gosasl/sasl_test.go +++ b/agents/go-common/gosasl/sasl_test.go @@ -153,3 +153,33 @@ func TestDigestMD5RejectsUnsupportedQOP(t *testing.T) { t.Fatalf("expected unsupported QOP error, got %v", err) } } + +func TestDigestMD5EchoesUTF8Charset(t *testing.T) { + mechanism := NewDigestMD5Mechanism("zookeeper", "user", "password") + client := NewSaslClient("zk-sasl-md5", mechanism) + if _, err := client.Start(); err != nil { + t.Fatal(err) + } + response, err := client.Step([]byte(`realm="zk-sasl-md5",nonce="nonce",qop="auth",charset=utf-8,algorithm=md5-sess`)) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(response), "charset=utf-8") { + t.Fatalf("response does not echo UTF-8 charset: %s", response) + } +} + +func TestDigestMD5DefaultsMissingQOPToAuth(t *testing.T) { + mechanism := NewDigestMD5Mechanism("zookeeper", "user", "password") + client := NewSaslClient("zk-sasl-md5", mechanism) + if _, err := client.Start(); err != nil { + t.Fatal(err) + } + response, err := client.Step([]byte(`realm="zk-sasl-md5",nonce="nonce",charset=utf-8,algorithm=md5-sess`)) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(response), "qop=auth") || strings.Contains(string(response), "maxbuf=") { + t.Fatalf("response does not use the DIGEST-MD5 default qop=auth: %s", response) + } +} diff --git a/agents/scripts/driver_release_packages_test.py b/agents/scripts/driver_release_packages_test.py index 7c471ddda4..79ec082621 100644 --- a/agents/scripts/driver_release_packages_test.py +++ b/agents/scripts/driver_release_packages_test.py @@ -43,6 +43,7 @@ class DriverReleasePackagesTest(unittest.TestCase): "duckdb": "0.1.0", "rabbitmq": "0.1.0", "rocketmq": "0.1.0", + "zookeeper": "0.1.0", "cassandra": "0.1.37", "hive": "0.1.43", "tdengine": "0.1.0", @@ -299,6 +300,21 @@ class DriverReleasePackagesTest(unittest.TestCase): self.assertFalse(source.exists()) self.assertEqual(versioned.read_bytes(), b"MZtest-hive-agent") + def test_versions_zookeeper_native_artifacts(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + release_dir = Path(temp_dir) + source = release_dir / "dbx-agent-zookeeper-linux-aarch64" + source.write_bytes(b"\x7fELFtest-zookeeper-agent") + versions = {driver: "0.1.0" for driver in NATIVE_DRIVERS} + versions["zookeeper"] = "0.1.8" + + renamed = version_agent_artifacts(release_dir, versions) + versioned = release_dir / "dbx-agent-zookeeper-0.1.8-linux-aarch64" + + self.assertEqual(renamed, [versioned]) + self.assertFalse(source.exists()) + self.assertEqual(versioned.read_bytes(), b"\x7fELFtest-zookeeper-agent") + def test_full_offline_bundle_includes_supported_windows_artifacts(self) -> None: with tempfile.TemporaryDirectory() as temp_dir: release_dir = Path(temp_dir) diff --git a/agents/scripts/validate_agents.py b/agents/scripts/validate_agents.py index ef8698d832..604e470333 100644 --- a/agents/scripts/validate_agents.py +++ b/agents/scripts/validate_agents.py @@ -25,6 +25,7 @@ NATIVE_ONLY_AGENT_MODULES = { "xugu": "drivers/xugu", "rabbitmq": "drivers/rabbitmq", "rocketmq": "drivers/rocketmq", + "zookeeper": "drivers/zookeeper", } AUTO_VERSIONED_NATIVE_MODULES = {"duckdb"} JDBC_ARCHITECTURE_ALLOWLIST = { diff --git a/agents/scripts/validate_agents_test.py b/agents/scripts/validate_agents_test.py index 66ebac8bbd..509992b80e 100644 --- a/agents/scripts/validate_agents_test.py +++ b/agents/scripts/validate_agents_test.py @@ -148,11 +148,12 @@ class ValidateAgentsTest(unittest.TestCase): "duckdb", "rabbitmq", "rocketmq", + "zookeeper", "tdengine", ): (root / "drivers" / driver).mkdir(parents=True) (root / "versions.json").write_text( - json.dumps({"h2": "0.1.0", "cassandra": "0.1.0", "hive": "0.1.0", "oracle": "0.1.0", "kingbase": "0.1.0", "iotdb": "0.1.0", "neo4j": "0.1.0", "vastbase": "0.1.0", "xugu": "0.1.0", "rabbitmq": "0.1.0", "rocketmq": "0.1.0", "tdengine": "0.1.0"}), + json.dumps({"h2": "0.1.0", "cassandra": "0.1.0", "hive": "0.1.0", "oracle": "0.1.0", "kingbase": "0.1.0", "iotdb": "0.1.0", "neo4j": "0.1.0", "vastbase": "0.1.0", "xugu": "0.1.0", "rabbitmq": "0.1.0", "rocketmq": "0.1.0", "zookeeper": "0.1.0", "tdengine": "0.1.0"}), encoding="utf-8", ) diff --git a/agents/scripts/version_agent_artifacts.py b/agents/scripts/version_agent_artifacts.py index 91de053acc..3c5a3a3e6c 100644 --- a/agents/scripts/version_agent_artifacts.py +++ b/agents/scripts/version_agent_artifacts.py @@ -4,7 +4,7 @@ import json from pathlib import Path -NATIVE_DRIVERS = ("cassandra", "hive", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "duckdb", "rabbitmq", "rocketmq", "tdengine") +NATIVE_DRIVERS = ("cassandra", "hive", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "duckdb", "rabbitmq", "rocketmq", "zookeeper", "tdengine") PLATFORMS = ( "macos-aarch64", "macos-x64", diff --git a/agents/settings.gradle b/agents/settings.gradle index 7b6ca2846e..a2420f51d7 100644 --- a/agents/settings.gradle +++ b/agents/settings.gradle @@ -6,7 +6,7 @@ def driverModules = [ 'teradata', 'vertica', 'firebird', 'exasol', 'oceanbase-oracle', 'gbase8a', 'gbase8s', 'bigquery', 'kylin', 'sundb', 'h2', 'h2-legacy', 'snowflake', 'trino', 'spark', 'db2', 'informix', 'mongodb', 'highgo', 'uxdb', 'yashandb', 'oscar', - 'iris', 'etcd', 'zookeeper', 'kafka', 'sqlserver-legacy' + 'iris', 'etcd', 'kafka', 'sqlserver-legacy' ] include(*(infrastructureModules + driverModules))