From c68cb88bf9b5c8159acb512005627df15485cf84 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 11 Aug 2026 23:12:28 +0800 Subject: [PATCH] feat(hive): migrate agent to native Go --- .github/scripts/bump-agent-versions.mjs | 14 +- .github/scripts/bump-agent-versions.test.mjs | 20 + .github/workflows/agents-release.yml | 55 +- .github/workflows/ci.yml | 8 + agents/README.md | 11 +- agents/README.zh-CN.md | 11 +- agents/build.gradle | 2 +- agents/drivers/hive-go/MIGRATION_PARITY.md | 154 ++ agents/drivers/hive-go/bench/README.md | 101 + agents/drivers/hive-go/bench/agent_compare.py | 664 +++++++ .../drivers/hive-go/bench/functional_probe.py | 203 ++ .../drivers/hive-go/bench/kdc_fixture/main.go | 99 + agents/drivers/hive-go/config.go | 1180 ++++++++++++ agents/drivers/hive-go/config_test.go | 474 +++++ agents/drivers/hive-go/connector.go | 194 ++ agents/drivers/hive-go/connector_test.go | 155 ++ agents/drivers/hive-go/discovery.go | 392 ++++ agents/drivers/hive-go/discovery_test.go | 221 +++ agents/drivers/hive-go/go.mod | 35 + agents/drivers/hive-go/go.sum | 68 + agents/drivers/hive-go/init_test.go | 65 + .../drivers/hive-go/kerberos_defaults_unix.go | 16 + .../hive-go/kerberos_defaults_windows.go | 19 + agents/drivers/hive-go/main.go | 611 ++++++ agents/drivers/hive-go/metadata.go | 753 ++++++++ agents/drivers/hive-go/metadata_test.go | 212 +++ agents/drivers/hive-go/protocol_error.go | 145 ++ agents/drivers/hive-go/query.go | 529 ++++++ agents/drivers/hive-go/query_test.go | 557 ++++++ agents/drivers/hive-go/zookeeper_protocol.go | 568 ++++++ .../hive-go/zookeeper_protocol_test.go | 562 ++++++ agents/drivers/hive-go/zookeeper_tls.go | 240 +++ agents/drivers/hive-go/zookeeper_tls_test.go | 162 ++ agents/go-common/go-gssapi/LICENSE | 201 ++ agents/go-common/go-gssapi/README.md | 228 +++ .../go-gssapi/common/channel_binding.go | 17 + agents/go-common/go-gssapi/flags.go | 59 + agents/go-common/go-gssapi/flags_test.go | 29 + agents/go-common/go-gssapi/go.mod | 11 + agents/go-common/go-gssapi/go.sum | 49 + agents/go-common/go-gssapi/interface.go | 95 + agents/go-common/go-gssapi/krb5/APRep.go | 138 ++ agents/go-common/go-gssapi/krb5/APRep_test.go | 118 ++ .../go-common/go-gssapi/krb5/context_token.go | 232 +++ .../go-gssapi/krb5/context_token_test.go | 162 ++ .../go-gssapi/krb5/credentials_test.go | 46 + .../go-gssapi/krb5/default_paths_unix.go | 20 + .../go-gssapi/krb5/default_paths_windows.go | 23 + agents/go-common/go-gssapi/krb5/keyinfo.go | 117 ++ .../go-common/go-gssapi/krb5/keyinfo_test.go | 55 + agents/go-common/go-gssapi/krb5/krb5.go | 1122 +++++++++++ agents/go-common/go-gssapi/krb5/krb5_test.go | 8 + .../go-common/go-gssapi/krb5/message_token.go | 506 +++++ .../go-gssapi/krb5/message_token_test.go | 208 +++ .../go-common/go-gssapi/krb5/sample_test.go | 110 ++ agents/go-common/go-gssapi/registry.go | 61 + agents/go-common/gohive/LICENSE | 20 + agents/go-common/gohive/browser_auth.go | 242 +++ agents/go-common/gohive/browser_auth_test.go | 116 ++ agents/go-common/gohive/connector.go | 153 ++ agents/go-common/gohive/driver.go | 346 ++++ agents/go-common/gohive/driver_test.go | 53 + agents/go-common/gohive/dsn.go | 107 ++ agents/go-common/gohive/dsn_test.go | 23 + agents/go-common/gohive/go.mod | 28 + agents/go-common/gohive/go.sum | 56 + agents/go-common/gohive/hive.go | 1650 +++++++++++++++++ agents/go-common/gohive/http_auth_test.go | 375 ++++ agents/go-common/gohive/metadata.go | 161 ++ agents/go-common/gohive/sasl_transport.go | 304 +++ .../go-common/gohive/sasl_transport_test.go | 44 + agents/go-common/gosasl/LICENSE | 19 + agents/go-common/gosasl/README.md | 65 + agents/go-common/gosasl/go.mod | 22 + agents/go-common/gosasl/go.sum | 48 + agents/go-common/gosasl/gssapi.go | 399 ++++ .../gosasl/gssapi_backend_nonwindows.go | 11 + .../gosasl/gssapi_backend_windows.go | 132 ++ agents/go-common/gosasl/gssapi_purego.go | 89 + agents/go-common/gosasl/http_spnego.go | 272 +++ agents/go-common/gosasl/http_spnego_test.go | 182 ++ agents/go-common/gosasl/sasl.go | 489 +++++ agents/go-common/gosasl/sasl_gssapi_test.go | 523 ++++++ agents/go-common/gosasl/sasl_test.go | 155 ++ .../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 +- crates/dbx-core/src/agent_manager.rs | 16 + crates/dbx-core/src/agent_service.rs | 10 +- 91 files changed, 18227 insertions(+), 22 deletions(-) create mode 100644 agents/drivers/hive-go/MIGRATION_PARITY.md create mode 100644 agents/drivers/hive-go/bench/README.md create mode 100755 agents/drivers/hive-go/bench/agent_compare.py create mode 100644 agents/drivers/hive-go/bench/functional_probe.py create mode 100644 agents/drivers/hive-go/bench/kdc_fixture/main.go create mode 100644 agents/drivers/hive-go/config.go create mode 100644 agents/drivers/hive-go/config_test.go create mode 100644 agents/drivers/hive-go/connector.go create mode 100644 agents/drivers/hive-go/connector_test.go create mode 100644 agents/drivers/hive-go/discovery.go create mode 100644 agents/drivers/hive-go/discovery_test.go create mode 100644 agents/drivers/hive-go/go.mod create mode 100644 agents/drivers/hive-go/go.sum create mode 100644 agents/drivers/hive-go/init_test.go create mode 100644 agents/drivers/hive-go/kerberos_defaults_unix.go create mode 100644 agents/drivers/hive-go/kerberos_defaults_windows.go create mode 100644 agents/drivers/hive-go/main.go create mode 100644 agents/drivers/hive-go/metadata.go create mode 100644 agents/drivers/hive-go/metadata_test.go create mode 100644 agents/drivers/hive-go/protocol_error.go create mode 100644 agents/drivers/hive-go/query.go create mode 100644 agents/drivers/hive-go/query_test.go create mode 100644 agents/drivers/hive-go/zookeeper_protocol.go create mode 100644 agents/drivers/hive-go/zookeeper_protocol_test.go create mode 100644 agents/drivers/hive-go/zookeeper_tls.go create mode 100644 agents/drivers/hive-go/zookeeper_tls_test.go create mode 100644 agents/go-common/go-gssapi/LICENSE create mode 100644 agents/go-common/go-gssapi/README.md create mode 100644 agents/go-common/go-gssapi/common/channel_binding.go create mode 100644 agents/go-common/go-gssapi/flags.go create mode 100644 agents/go-common/go-gssapi/flags_test.go create mode 100644 agents/go-common/go-gssapi/go.mod create mode 100644 agents/go-common/go-gssapi/go.sum create mode 100644 agents/go-common/go-gssapi/interface.go create mode 100644 agents/go-common/go-gssapi/krb5/APRep.go create mode 100644 agents/go-common/go-gssapi/krb5/APRep_test.go create mode 100644 agents/go-common/go-gssapi/krb5/context_token.go create mode 100644 agents/go-common/go-gssapi/krb5/context_token_test.go create mode 100644 agents/go-common/go-gssapi/krb5/credentials_test.go create mode 100644 agents/go-common/go-gssapi/krb5/default_paths_unix.go create mode 100644 agents/go-common/go-gssapi/krb5/default_paths_windows.go create mode 100644 agents/go-common/go-gssapi/krb5/keyinfo.go create mode 100644 agents/go-common/go-gssapi/krb5/keyinfo_test.go create mode 100644 agents/go-common/go-gssapi/krb5/krb5.go create mode 100644 agents/go-common/go-gssapi/krb5/krb5_test.go create mode 100644 agents/go-common/go-gssapi/krb5/message_token.go create mode 100644 agents/go-common/go-gssapi/krb5/message_token_test.go create mode 100644 agents/go-common/go-gssapi/krb5/sample_test.go create mode 100644 agents/go-common/go-gssapi/registry.go create mode 100644 agents/go-common/gohive/LICENSE create mode 100644 agents/go-common/gohive/browser_auth.go create mode 100644 agents/go-common/gohive/browser_auth_test.go create mode 100644 agents/go-common/gohive/connector.go create mode 100644 agents/go-common/gohive/driver.go create mode 100644 agents/go-common/gohive/driver_test.go create mode 100644 agents/go-common/gohive/dsn.go create mode 100644 agents/go-common/gohive/dsn_test.go create mode 100644 agents/go-common/gohive/go.mod create mode 100644 agents/go-common/gohive/go.sum create mode 100644 agents/go-common/gohive/hive.go create mode 100644 agents/go-common/gohive/http_auth_test.go create mode 100644 agents/go-common/gohive/metadata.go create mode 100644 agents/go-common/gohive/sasl_transport.go create mode 100644 agents/go-common/gohive/sasl_transport_test.go create mode 100644 agents/go-common/gosasl/LICENSE create mode 100644 agents/go-common/gosasl/README.md create mode 100644 agents/go-common/gosasl/go.mod create mode 100644 agents/go-common/gosasl/go.sum create mode 100644 agents/go-common/gosasl/gssapi.go create mode 100644 agents/go-common/gosasl/gssapi_backend_nonwindows.go create mode 100644 agents/go-common/gosasl/gssapi_backend_windows.go create mode 100644 agents/go-common/gosasl/gssapi_purego.go create mode 100644 agents/go-common/gosasl/http_spnego.go create mode 100644 agents/go-common/gosasl/http_spnego_test.go create mode 100644 agents/go-common/gosasl/sasl.go create mode 100644 agents/go-common/gosasl/sasl_gssapi_test.go create mode 100644 agents/go-common/gosasl/sasl_test.go diff --git a/.github/scripts/bump-agent-versions.mjs b/.github/scripts/bump-agent-versions.mjs index 46708b35cf..be7b434d1e 100644 --- a/.github/scripts/bump-agent-versions.mjs +++ b/.github/scripts/bump-agent-versions.mjs @@ -51,6 +51,7 @@ function fileContainsCommonDependency(path, moduleExists, readModuleFile) { const nativeDriverDirectories = { cassandra: "cassandra-go", duckdb: "duckdb", + hive: "hive-go", oracle: "oracle-go", kingbase: "kingbase-go", iotdb: "iotdb", @@ -59,7 +60,14 @@ const nativeDriverDirectories = { rabbitmq: "rabbitmq", tdengine: "tdengine", }; -const nativeDriverModules = new Set(["cassandra", "duckdb", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "rabbitmq", "tdengine"]); +const nativeDriverModules = new Set(["cassandra", "duckdb", "hive", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "rabbitmq", "tdengine"]); +const nativeDriverSharedPaths = { + hive: [ + "agents/go-common/go-gssapi", + "agents/go-common/gohive", + "agents/go-common/gosasl", + ], +}; function resolveAgentModule(moduleName, { legacyStandaloneModules, moduleExists, readModuleFile }) { let checkDir = null; @@ -82,6 +90,7 @@ function resolveAgentModule(moduleName, { legacyStandaloneModules, moduleExists, return { checkDir, modulePath, + sharedPaths: nativeDriverSharedPaths[moduleName] ?? [], javaBuild: hasBuildGradle, nativeBuild: nativeDriverModules.has(moduleName), commonDependent: hasBuildGradle && (explicitlyDependsOnCommon || !legacyStandaloneModules.has(moduleName)), @@ -137,7 +146,8 @@ export function evaluateAgentVersionBump({ } for (const { moduleName, module } of resolvedModules) { - const moduleChanged = pathChanged(changedFiles, module.modulePath); + const moduleChanged = [module.modulePath, ...module.sharedPaths] + .some((path) => pathChanged(changedFiles, path)); // Only modules that package agents/common need installer-visible updates // for shared Java runtime changes; native and standalone agents do not. const commonAffectsModule = commonChanged && module.commonDependent; diff --git a/.github/scripts/bump-agent-versions.test.mjs b/.github/scripts/bump-agent-versions.test.mjs index 533e1fcd5c..f5be89613a 100644 --- a/.github/scripts/bump-agent-versions.test.mjs +++ b/.github/scripts/bump-agent-versions.test.mjs @@ -83,6 +83,26 @@ test("bumps Cassandra from its native Go source directory", () => { assert.deepEqual(result.nativeModules, ["cassandra"]); }); +test("bumps Hive from native and shared Kerberos source directories", () => { + for (const changedFile of [ + "agents/drivers/hive-go/main.go", + "agents/go-common/gohive/driver.go", + "agents/go-common/gosasl/gssapi.go", + "agents/go-common/go-gssapi/krb5/krb5.go", + ]) { + const result = evaluateAgentVersionBump({ + versions: { hive: "0.1.43" }, + changedFiles: [changedFile], + moduleExists: (path) => path === "agents/drivers/hive-go", + readModuleFile: () => "", + }); + + assert.equal(result.versions.hive, "0.1.44"); + assert.deepEqual(result.javaModules, []); + assert.deepEqual(result.nativeModules, ["hive"]); + } +}); + test("bumps Neo4j from its native Go source directory", () => { const result = evaluateAgentVersionBump({ versions: { neo4j: "0.1.39" }, diff --git a/.github/workflows/agents-release.yml b/.github/workflows/agents-release.yml index 62bbe49c1c..78269b7a1c 100644 --- a/.github/workflows/agents-release.yml +++ b/.github/workflows/agents-release.yml @@ -329,6 +329,46 @@ jobs: name: cassandra-native path: "release-native/dbx-agent-cassandra-*" + build-hive-native: + needs: [bump-versions] + if: ${{ contains(fromJSON(needs.bump-versions.outputs.native_modules), 'hive') }} + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version: "1.23.x" + - name: Test Hive native agent + working-directory: agents/drivers/hive-go + run: go test ./... + - name: Cross-compile Hive native agent + shell: bash + run: | + mkdir -p release-native + cd agents/drivers/hive-go + 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-hive-${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: hive-native + path: "release-native/dbx-agent-hive-*" + build-kingbase-native: needs: [bump-versions] if: ${{ contains(fromJSON(needs.bump-versions.outputs.native_modules), 'kingbase') }} @@ -828,7 +868,7 @@ jobs: retention-days: 1 release: - needs: [bump-versions, commit-versions, build-agents, build-oracle-native, build-xugu-native, build-rabbitmq-native, build-cassandra-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-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: @@ -951,6 +991,7 @@ jobs: xugu) echo "虚谷 XuguDB" ;; rabbitmq) echo "RabbitMQ" ;; cassandra) echo "Apache Cassandra" ;; + hive) echo "Apache Hive" ;; neo4j) echo "Neo4j" ;; iotdb) echo "Apache IoTDB" ;; tdengine) echo "TDengine" ;; @@ -1017,7 +1058,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 cassandra tdengine; do + for name in oracle xugu kingbase vastbase neo4j iotdb duckdb rabbitmq cassandra hive tdengine; do version=$(get_module_version "$name") [ -f "release/dbx-agent-${name}-${version}.jar" ] && continue native_json=$(generate_native_platforms "$name" "$version") @@ -1088,6 +1129,7 @@ jobs: xugu) echo "虚谷 XuguDB" ;; rabbitmq) echo "RabbitMQ" ;; cassandra) echo "Apache Cassandra" ;; + hive) echo "Apache Hive" ;; neo4j) echo "Neo4j" ;; iotdb) echo "Apache IoTDB" ;; tdengine) echo "TDengine" ;; @@ -1117,6 +1159,8 @@ jobs: LOG_PATH="agents/drivers/vastbase-go/" elif [ "$name" = "cassandra" ]; then LOG_PATH="agents/drivers/cassandra-go/" + elif [ "$name" = "hive" ]; then + LOG_PATH="agents/drivers/hive-go/" elif [ "$name" = "neo4j" ]; then LOG_PATH="agents/drivers/neo4j-go/" elif [ "$name" = "iotdb" ]; then @@ -1127,6 +1171,13 @@ jobs: LOG_PATH="agents/$name/" fi LOG_PATHS=("$LOG_PATH") + if [ "$name" = "hive" ]; then + LOG_PATHS+=( + "agents/go-common/go-gssapi/" + "agents/go-common/gohive/" + "agents/go-common/gosasl/" + ) + fi while IFS= read -r line; do [ -n "$line" ] && NOTES="${NOTES}- ${line}"$'\n' done < <(git log --oneline "$PREV_TAG"..HEAD -- "${LOG_PATHS[@]}" | sed 's/^[0-9a-f]* //') diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3c85b3f5a6..9eff4e5fdb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -744,6 +744,10 @@ jobs: run: go test ./... working-directory: agents/drivers/cassandra-go + - name: Hive native agent tests + run: go test ./... + working-directory: agents/drivers/hive-go + - name: Vastbase native agent tests run: go test ./... working-directory: agents/drivers/vastbase-go @@ -772,6 +776,10 @@ jobs: run: CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-cassandra-linux-x64 . working-directory: agents/drivers/cassandra-go + - name: Hive native agent build + run: CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-hive-linux-x64 . + working-directory: agents/drivers/hive-go + - name: Vastbase native agent build run: CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -trimpath -ldflags="-s -w" -o /tmp/dbx-agent-vastbase-linux-x64 . working-directory: agents/drivers/vastbase-go diff --git a/agents/README.md b/agents/README.md index 03d858345c..fe4f45c66c 100644 --- a/agents/README.md +++ b/agents/README.md @@ -30,7 +30,7 @@ Each agent runs as a standalone process and communicates with DBX via stdin/stdo | h2 | H2 | H2 JDBC | | snowflake | Snowflake | Snowflake JDBC | | trino | Trino (Presto) | Trino JDBC | -| hive | Apache Hive | Hive JDBC | +| hive | Apache Hive | Native Go HS2 agent | | db2 | IBM DB2 | DB2 JDBC | | informix | IBM Informix | Informix JDBC | | neo4j | Neo4j | Official Neo4j Go Driver native agent | @@ -49,7 +49,7 @@ Each agent runs as a standalone process and communicates with DBX via stdin/stdo ## Multi-JRE Support -Most Java agents target JRE 21. Native agents, such as `cassandra`, `duckdb`, `iotdb`, `oracle`, `kingbase`, `tdengine`, `xugu`, and `rabbitmq`, 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`, and `rabbitmq`, do not require a JRE. DBX downloads and manages the JRE 21 installation automatically for Java agents. ## JDBC Connection Pooling @@ -75,7 +75,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/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`, and `drivers/rabbitmq` (amqp091-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`, and `drivers/rabbitmq` (amqp091-go) 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`. @@ -88,6 +88,7 @@ Requires JDK 21 (Gradle toolchain auto-downloads if needed). ./gradlew shadowJar (cd drivers/oracle-go && go build -o agent .) (cd drivers/cassandra-go && go build -o agent .) +(cd drivers/hive-go && go build -o agent .) (cd drivers/iotdb && go build -o agent .) (cd drivers/kingbase-go && go build -o agent .) (cd drivers/vastbase-go && go build -o agent .) @@ -96,7 +97,7 @@ Requires JDK 21 (Gradle toolchain auto-downloads if needed). (cd drivers/rabbitmq && go build -o agent .) ``` -Output JARs are in `drivers/{module}/build/libs/`. Native agents build from `drivers/cassandra-go`, `drivers/duckdb`, `drivers/iotdb`, `drivers/oracle-go`, `drivers/kingbase-go`, `drivers/vastbase-go`, `drivers/tdengine`, `drivers/xugu`, and `drivers/rabbitmq`. +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`, and `drivers/rabbitmq`. ### Local DBX Runtime Test @@ -110,7 +111,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`, `iotdb`, `oracle`, `kingbase`, `tdengine`, `xugu`, and `rabbitmq` 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`, and `rabbitmq` 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 c9f4a4e5a2..076c7b5c83 100644 --- a/agents/README.zh-CN.md +++ b/agents/README.zh-CN.md @@ -30,7 +30,7 @@ DBX 的 Agent 驱动 —— 通过 JDBC 和原生数据库驱动支持各种数 | h2 | H2 | H2 JDBC | | snowflake | Snowflake | Snowflake JDBC | | trino | Trino (Presto) | Trino JDBC | -| hive | Apache Hive | Hive JDBC | +| hive | Apache Hive | 原生 Go HS2 agent | | db2 | IBM DB2 | DB2 JDBC | | informix | IBM Informix | Informix JDBC | | neo4j | Neo4j | 官方 Neo4j Go Driver 原生 Agent | @@ -49,7 +49,7 @@ DBX 的 Agent 驱动 —— 通过 JDBC 和原生数据库驱动支持各种数 ## 多 JRE 支持 -多数 Java agent 以 JRE 21 为目标。原生 agent（如 `cassandra`、`duckdb`、`iotdb`、`oracle`、`kingbase`、`tdengine`、`xugu` 和 `rabbitmq`）不需要 JRE。对 Java agent，DBX 会自动下载并管理 JRE 21 安装。 +多数 Java agent 以 JRE 21 为目标。原生 agent（如 `cassandra`、`duckdb`、`hive`、`iotdb`、`oracle`、`kingbase`、`tdengine`、`xugu` 和 `rabbitmq`）不需要 JRE。对 Java agent，DBX 会自动下载并管理 JRE 21 安装。 ## JDBC 连接池 @@ -75,7 +75,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/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）。无需 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）。无需 JRE 下载与管理。 - **Java/JDBC** —— 当某数据库只有 JDBC 驱动，或原生驱动不成熟、缺乏维护时的默认兜底方案。多数 agent 仍属此类。 原生 agent 实现与 Java agent 相同的 JSON-RPC 契约和 `versions.json` 登记；它发布的是 `agent` 可执行文件而非 `agent.jar`。若同一数据库同时保留原生和 Java 源码实现，默认只发布原生产物；只有 Java 变体以独立兼容配置登记时才同时发布，例如 `oracle-legacy` / `oracle-10g`。 @@ -88,6 +88,7 @@ HikariCP 会直接打进启用连接池的 Agent JAR。已经使用 DBX 托管 J ./gradlew shadowJar (cd drivers/oracle-go && go build -o agent .) (cd drivers/cassandra-go && go build -o agent .) +(cd drivers/hive-go && go build -o agent .) (cd drivers/iotdb && go build -o agent .) (cd drivers/kingbase-go && go build -o agent .) (cd drivers/vastbase-go && go build -o agent .) @@ -96,7 +97,7 @@ HikariCP 会直接打进启用连接池的 Agent JAR。已经使用 DBX 托管 J (cd drivers/rabbitmq && go build -o agent .) ``` -产物 JAR 在 `drivers/{module}/build/libs/`。原生 agent 从 `drivers/cassandra-go`、`drivers/duckdb`、`drivers/iotdb`、`drivers/oracle-go`、`drivers/kingbase-go`、`drivers/vastbase-go`、`drivers/tdengine`、`drivers/xugu` 和 `drivers/rabbitmq` 构建。 +产物 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` 构建。 ### 本地 DBX 运行时测试 @@ -110,7 +111,7 @@ cp agents/drivers//build/libs/*-all.jar ~/.dbx/agents/drivers/ 重启 DBX 或断开重连数据库，使新 agent 进程加载替换后的 JAR。 -`cassandra`、`iotdb`、`oracle`、`kingbase`、`tdengine`、`xugu` 和 `rabbitmq` 等原生 agent 使用可执行文件而非 `agent.jar`。TDengine 从 `drivers/tdengine/Cargo.toml` 构建 `target/release/dbx-tdengine-driver`。 +`cassandra`、`hive`、`iotdb`、`oracle`、`kingbase`、`tdengine`、`xugu` 和 `rabbitmq` 等原生 agent 使用可执行文件而非 `agent.jar`。TDengine 从 `drivers/tdengine/Cargo.toml` 构建 `target/release/dbx-tdengine-driver`。 ## 版本管理 diff --git a/agents/build.gradle b/agents/build.gradle index ea4fe12c61..f218cc04df 100644 --- a/agents/build.gradle +++ b/agents/build.gradle @@ -6,7 +6,7 @@ def infrastructureProjects = ['common', 'test-support'] as Set def legacyStandaloneProjects = ['mongodb', 'kafka', 'rocketmq'] as Set def pooledJdbcProjects = [ 'access', 'bigquery', 'dameng', 'databend', 'databricks', 'db2', 'exasol', - 'firebird', 'gbase8a', 'gbase8s', 'goldendb', 'h2', 'h2-legacy', 'highgo', 'hive', + 'firebird', 'gbase8a', 'gbase8s', 'goldendb', 'h2', 'h2-legacy', 'highgo', 'informix', 'iris', 'kylin', 'oceanbase-oracle', 'oscar', 'saphana', 'snowflake', 'spark', 'sqlserver-legacy', 'sundb', 'teradata', 'trino', 'uxdb', 'vertica', 'yashandb' diff --git a/agents/drivers/hive-go/MIGRATION_PARITY.md b/agents/drivers/hive-go/MIGRATION_PARITY.md new file mode 100644 index 0000000000..fcd8cdd22f --- /dev/null +++ b/agents/drivers/hive-go/MIGRATION_PARITY.md @@ -0,0 +1,154 @@ +# Hive JDBC to Go migration parity + +Status date: 2026-08-11. + +The migration is not complete until a capability is implemented, covered by +automated tests, validated against a real compatible server, and included in +the DBX native-agent build and release path. A unit test alone does not count as +production parity. + +## Baseline + +- DBX Java baseline: Apache Hive JDBC standalone 4.0.1. +- Compatibility reference: DBeaver keeps separate Hive 2 legacy and Hive 4+ + JDBC profiles. The Go migration must not infer Hive 2 support from Hive 3/4. +- DBX protocol baseline: the Go Agent implements the same stdin/stdout JSON-RPC + methods used by the Java Agent. +- HS2 client protocol: Go requests `HIVE_CLI_SERVICE_PROTOCOL_V10`; Hive 3.1.3 + and Hive 4.2.0 accepted it in live validation. +- The current completion pass validates the Go Agent only. JDBC is retained as + a historical behavior reference, not as a candidate in the secure discovery + or Kerberos validation path. + +## Current matrix + +| Capability | Go code | Automated | Linux live | Windows live | Status | +| --- | --- | --- | --- | --- | --- | +| Hive 3.1.3 binary NOSASL | yes | yes | yes | n/a | parity smoke passed | +| Hive 4.2.0 binary NOSASL | yes | yes | yes | n/a | parity smoke passed | +| Spark 3.5.7 Thrift Server | yes | yes | yes | n/a | Go passed; Java requires a non-empty user in this fixture | +| Hive 2.x | probable protocol compatibility | partial | no | no | unsupported until a real Hive 2 server passes | +| Kyuubi | probable HS2 compatibility | partial | no | no | unsupported until a real Kyuubi server passes | +| Binary PLAIN (`NONE`) | yes | yes | yes | no | Java/Go parity passed on Hive 4.2.0 | +| Binary LDAP/CUSTOM PLAIN | yes | yes | no | no | needs a real authentication backend | +| Binary Kerberos `auth` | yes | yes | yes | no | Go keytab login, query, metadata, and clean shutdown passed on Hive 4.2.0 | +| Binary Kerberos `auth-int` | yes | yes | yes | no | Go integrity-protected query passed on Hive 4.2.0 | +| Binary Kerberos `auth-conf` | yes | yes | yes | no | Go confidentiality-protected query passed on Hive 4.2.0 | +| HTTP PLAIN/Basic | yes | yes | yes | no | Java/Go parity passed on Hive 4.2.0 | +| HTTP NOSASL | yes | yes | no | no | needs a real HS2 HTTP fixture | +| HTTP LDAP/CUSTOM | yes | yes | no | no | needs a real authentication backend | +| HTTP Kerberos/SPNEGO | yes | yes | no | no | needs KDC + HS2 validation | +| HTTP Kerberos TLS channel binding | yes | yes | no | no | needs TLS + KDC validation | +| HTTP JWT bearer | yes | yes | no | no | header and cookie retry behavior are covered; real HS2 JWT validation is pending | +| HTTP browser SSO | yes | yes | no | no | pre-issued token and interactive 302/browser/loopback callback flow are implemented; real IdP validation is pending | +| HTTP delegation-token header | yes | yes | no | no | `X-Hive-Delegation-Token` behavior is covered; real HS2 token validation is pending | +| HTTP cookie auth, XSRF/CSRF, and request tracking | yes | yes | partial | no | Java-compatible headers, static/server cookies, 401 credential retry, and `X-Request-ID` are covered; the retained HTTP fixture failed to restart because of its stale PID state | +| One-way TLS | yes | yes | no | no | PEM/JKS/PKCS12 parsing is tested; handshake is not | +| Mutual TLS | yes | yes | no | no | PEM/JKS/PKCS12 parsing is tested; handshake is not | +| Binary delegation token (`DIGEST-MD5`) | yes | yes | no | no | token decoding is tested; HS2 exchange is not | +| ZooKeeper service discovery | yes | yes | yes | no | two-node discovery and reconnect failover passed | +| ZooKeeper stale-node handling | yes | yes | yes | no | Java/Go both passed 12 sequential connects with one stale node | +| ZooKeeper digest ACL | yes | yes | no | no | needs a secured ZooKeeper fixture | +| ZooKeeper TLS | yes | yes | no | no | trust/key store parsing is tested; handshake is not | +| ZooKeeper Kerberos SASL | yes | MiniKDC protocol test | yes | no | required-SASL ZooKeeper discovery into Kerberos HS2 passed | +| ZooKeeper active/passive HA mode | yes | yes | no | no | needs an active/passive HS2 fixture | +| Windows Kerberos SSPI | yes | Windows x64 cross-build | n/a | no | PE32+ amd64 build passed; Windows domain live validation is still required | +| Keytab Kerberos | yes | yes | yes | no | MiniKDC + Hive 4.2.0 live validation passed | +| Ccache and password Kerberos | yes | yes | no | no | credential-source parsing is tested; real HS2 login is pending | +| JDBC URL session/hiveConf/hiveVar sections | yes | yes | yes | no | Hive 4.2.0 session values passed | +| Proxy user and compatibility session variables | yes | yes | partial | no | parsing/open-session mapping passed | +| Query values and column type semantics | yes | yes | yes | no | Hive 4.2.0 type matrix matches Java except improved binary hex output | +| Metadata databases/tables/columns/DDL | yes | yes | partial | no | database/table smoke and `visible_schemas` filtering passed; full metadata matrix pending | +| Paged reads | yes | yes | yes | no | Hive 3.1.3 and 4.2.0 parity passed | +| Failed SQL is not replayed | yes | yes | yes | no | failed statement followed by successful query passed | +| Cancellation and timeout | yes | yes | no | no | real long-running query validation pending | +| Large result and large complex values | yes | yes | partial | no | functional samples passed; boundary fixture pending | +| JDBC client compatibility properties | yes | yes | partial | no | fetch/message sizing, retries, init file, application name, HTTP headers/cookies, request tracking, and browser settings are mapped | +| Native DBX install/launch | yes | yes | local artifact smoke | no | DBX tests prove native launch without a JRE and replacement of a stale Hive `agent.jar`; packaged desktop upgrade remains pending | +| Native CI/release artifacts | yes | yes | local build | cross-build | Hive version bumping, registry packaging, release notes, CI tests, and six native targets are wired | + +## JDBC 4.0.1 client feature coverage + +The native Agent now maps these Hive JDBC 4.0.1 client behaviors: + +- JWT bearer authentication and delegation-token HTTP headers. +- Browser SSO with either a pre-issued bearer token or the JDBC-compatible + interactive 302 redirect, local callback listener, browser launch, token, and + client-identifier retry flow. +- Configurable cookie authentication and cookie name, including 401 retry and + static `http.cookie.*` authentication cookies. +- `http.header.*`, `http.cookie.*`, the JDBC XSRF/CSRF headers, and + `requestTrack` / `X-Request-ID`. +- `retries`, `retryInterval`, `initFile`, connection-level `fetchSize`, + `socketTimeout`, and `thrift.client.max.message.size`. +- `applicationName` / `ApplicationName`, `wmPool`, proxy user, session variables, + HiveConf, and HiveVar OpenSession mappings. +- Browser response port/timeout and the JDBC browser SSL requirement override. + +Java `kerberosAuthType=fromSubject` has no literal Go `Subject` object. Its +native equivalent is the connection-scoped credential abstraction: Windows +SSPI on Windows and the default credential cache on Unix, with explicit ccache, +keytab, or password sources still supported. + +One Java-specific secret source remains intentionally non-silent: +`storePasswordPath` points at a Hadoop credential-provider/JCEKS store. The Go +Agent rejects this case unless `trustStorePassword` / `keyStorePassword` is +provided explicitly; it does not pretend that the Java credential provider was +read successfully. + +The remaining migration work is therefore live compatibility validation and +native DBX delivery verification, not another Java implementation. + +## Live evidence + +The Linux x86-64 Go binary used for the secure Kerberos validation has SHA-256: + +```text +c41cb7c1192748d70dfaf575123059f78a42d1f1fd0b1d6952769ccd3dcab8d6 +``` + +The previous Linux x86-64 native artifact after the HTTP, Browser SSO, +init-file, and release-path completion pass had SHA-256: + +```text +2053c4d127a2bb3fd67eb31b995998cce749b7adec7b13548e768f53435a2850 +``` + +The current Linux x86-64 native artifact after the DBX visible-schema and +native-upgrade completion pass has SHA-256: + +```text +ea1924508688fc5f9ab3abab914fc0cc9a0a8c811bbfd95a14a4c57a82f4696d +``` + +Validation result SHA-256 values: + +```text +707846a387abce3b3a4f282e22afcb514f5e2760eee04f2a2f7f822740acc9fe functional Hive 3.1.3 / 4.2.0 / Spark smoke +8774dbbda55a5fc0aada1862659da29fcd0be2a1e7d88c28196981a7e6913479 Hive 4.2.0 Binary PLAIN +8774dbbda55a5fc0aada1862659da29fcd0be2a1e7d88c28196981a7e6913479 Hive 4.2.0 HTTP PLAIN/Basic +3d1d2d8278b3c792d0ac8c109a2de59cda86cb0a329fb9f74e87d62603b7af60 ZooKeeper two-node reconnect failover +8e97211cf9f4c1b95feb1c498230fe725b4e5e2b01c068ef0cb59ced11adb7d6 ZooKeeper stale-node handling +521137fdc96a04462956f01d6e83806122736f5da1da49d4a2dbb8d80a397d72 Hive 4.2.0 Binary Kerberos auth +bc16ca7661072a83e3bf1066e8a9db1de58c62a4dea6caca88baccdb5dd89219 Hive 4.2.0 Binary Kerberos auth-int +8309580d892b2e92e79122b590d4e3a1811e513c821ff02e5f386304ff1e6c3f Hive 4.2.0 Binary Kerberos auth-conf +f1bc8ac45e523cf21873f89343912f791d144e680c4832132fa4f419d7e32838 Kerberos ZooKeeper discovery into Kerberos Hive 4.2.0 +d81b3506acace2b16270d7ee806de6f5b165cfcae4aea17bddbcb57fb67a021e final native artifact against ZooKeeper-discovered Hive 4.2.0 +6ea6919c3239f4c5b0486429ab27cbd64f60a4d9e2c2fcd2df8d4ca161c2c186 current native artifact against ZooKeeper-discovered Hive 4.2.0 +fcd9069d1a6dfaeee3a478a3187a8a63f5e699175fc03f90edeb543274fc5e97 current visible-schema live validation +``` + +The two PLAIN result files intentionally have the same hash because they record +the same logical Java/Go result through different transports. + +## Next gates + +1. HTTP JWT, delegation-token, Browser SSO/IdP, Kerberos, and TLS channel-binding live validation. +2. One-way TLS and mutual-TLS handshakes. +3. ZooKeeper digest ACL, TLS, and active/passive HA. +4. Cancellation, timeout, and large-result live semantics. +5. Kyuubi and a real Hive 2.x deployment. +6. Windows x64 SSPI validation with the same KDC/HS2 fixture. +7. Packaged desktop install/launch and an actual user-data upgrade from the + previous Java artifact; automated core tests already cover native selection + and stale `agent.jar` replacement. diff --git a/agents/drivers/hive-go/bench/README.md b/agents/drivers/hive-go/bench/README.md new file mode 100644 index 0000000000..9d4d84a000 --- /dev/null +++ b/agents/drivers/hive-go/bench/README.md @@ -0,0 +1,101 @@ +# Hive Agent benchmark + +This benchmark compares the same DBX JSON-RPC operations through the native +Go Hive Agent and the archived JDBC Hive Agent. Both candidates run on the same +host and connect to the same HiveServer2 instance. + +Use `functional_probe.py` before running performance benchmarks. It validates +connection, sessions, query values, metadata, pagination, failed-SQL semantics, +clean shutdown, and Java/Go result parity without concurrency load. + +The runner measures: + +- process startup and fresh connection latency; +- artifact size, idle RSS, and peak RSS; +- `SELECT 1`-shape lookup and 100/1,000/10,000-row decoding; +- `list_databases`, `list_tables`, and complete paged reads; +- 1, 8, and 32 concurrent DBX Agent sessions; +- mean, p50, p95, p99, throughput, and clean shutdown behavior. + +Candidate order rotates between rounds to reduce warm-cache and server-order +bias. Startup and connection samples always use a fresh Agent process. + +## Prepare the fixture + +The defaults expect `dbx_agent_bench.agent_bench` with exactly 10,000 or more +rows and columns named `id` and `payload`: + +```sql +CREATE DATABASE IF NOT EXISTS dbx_agent_bench; +CREATE TABLE IF NOT EXISTS dbx_agent_bench.agent_bench ( + id BIGINT, + payload STRING +) STORED AS ORC; +``` + +Populate deterministic rows before running the benchmark. Keep the fixture, +HiveServer2 configuration, Agent host, and Java runtime unchanged between +candidates. + +## Run + +Functional parity probe: + +```bash +GO_AGENT=/tmp/dbx-hive-bench/hive-agent-linux-amd64 \ +JDBC_AGENT_JAR=/tmp/dbx-hive-bench/dbx-agent-hive.jar \ +JAVA_BIN=/tmp/dbx-hive-bench/jre21/bin/java \ +HIVE_HOST=127.0.0.1 \ +HIVE_PORT=10000 \ +HIVE_DATABASE=dbx_agent_bench \ +HIVE_URL_PARAMS=auth=noSasl \ +python3 agents/drivers/hive-go/bench/functional_probe.py \ + > /tmp/dbx-hive-bench/functional-result.json +``` + +The functional probe runs the Go candidate only by default. Set +`BENCH_CANDIDATES=go,jdbc` only when an explicit historical-JDBC comparison is +needed. + +Performance benchmark: + +```bash +GO_AGENT=/tmp/dbx-hive-bench/hive-agent-linux-amd64 \ +JDBC_AGENT_JAR=/tmp/dbx-hive-bench/dbx-agent-hive.jar \ +HIVE_HOST=127.0.0.1 \ +HIVE_PORT=10000 \ +HIVE_DATABASE=dbx_agent_bench \ +HIVE_URL_PARAMS=auth=noSasl \ +python3 agents/drivers/hive-go/bench/agent_compare.py \ + > /tmp/dbx-hive-bench/result.json +``` + +## Configuration + +- `BENCH_CANDIDATES`: performance benchmark default `go,jdbc`; functional probe + default `go`. +- `GO_AGENT_COMMAND`, `JDBC_AGENT_COMMAND`: optional full launch commands. +- `BENCH_STARTUPS`, `BENCH_CONNECTS`: fresh-process sample counts, default `8`. +- `BENCH_ROUNDS`: alternating steady-state rounds, default `3`. +- `BENCH_WARMUPS`: warmups before each workload, default `2`. +- `BENCH_CONCURRENCY`: comma-separated session counts, default `1,8,32`. +- `BENCH_CONCURRENCY_OPS_PER_WORKER`: operations per session, default `8`. +- `HIVE_HOST`, `HIVE_PORT`, `HIVE_DATABASE`, `HIVE_BENCH_TABLE`. +- `HIVE_USERNAME`, `HIVE_PASSWORD`, `HIVE_URL_PARAMS`, `HIVE_CONNECTION_STRING`. +- `HIVE_SSL`, `HIVE_CA_CERT_PATH`, `HIVE_CLIENT_CERT_PATH`, `HIVE_CLIENT_KEY_PATH`. +- `BENCH_*_SQL` and `BENCH_*_COUNT` override individual workloads. + +Run the benchmark on the Agent host. Do not compare a local Go process with a +remote JDBC process, use different HS2 endpoints, or mutate the fixture between +candidates. + +## Kerberos fixture + +`kdc_fixture` starts a test-only in-process KDC and writes a temporary +`krb5.conf` and keytab containing `alice` and `hive/localhost`. Never use these +credentials outside an isolated compatibility environment, and delete the +generated directory after validation. + +```bash +go run ./bench/kdc_fixture -dir /tmp/dbx-hive-kerberos +``` diff --git a/agents/drivers/hive-go/bench/agent_compare.py b/agents/drivers/hive-go/bench/agent_compare.py new file mode 100755 index 0000000000..f2c68d8994 --- /dev/null +++ b/agents/drivers/hive-go/bench/agent_compare.py @@ -0,0 +1,664 @@ +#!/usr/bin/env python3 +import json +import os +import queue +import shlex +import statistics +import subprocess +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path + + +@dataclass(frozen=True) +class Candidate: + name: str + command: list[str] + artifact: Path + rss_command: str = "" + + +class AgentProcess: + def __init__(self, candidate: Candidate): + self.candidate = candidate + self.process = subprocess.Popen( + candidate.command, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + ) + self.request_id = 0 + self.request_lock = threading.Lock() + self.write_lock = threading.Lock() + self.pending: dict[int, queue.Queue] = {} + self.ready = threading.Event() + self.saw_ready = False + self.exited = threading.Event() + self.stderr_lines: list[str] = [] + threading.Thread(target=self._read_stdout, daemon=True).start() + threading.Thread(target=self._drain_stderr, daemon=True).start() + if not self.ready.wait(env_float("BENCH_READY_TIMEOUT", 30.0)) or not self.saw_ready: + raise TimeoutError(self._failure("timed out waiting for agent readiness")) + + def _read_stdout(self) -> None: + assert self.process.stdout is not None + for line in self.process.stdout: + try: + response = json.loads(line) + except json.JSONDecodeError: + continue + if response.get("ready") is True: + self.saw_ready = True + self.ready.set() + continue + response_id = response.get("id") + if not isinstance(response_id, int): + continue + with self.request_lock: + response_queue = self.pending.get(response_id) + if response_queue is not None: + response_queue.put(response) + self.exited.set() + self.ready.set() + with self.request_lock: + pending = list(self.pending.values()) + for response_queue in pending: + response_queue.put(RuntimeError(self._failure("agent process exited"))) + + def _drain_stderr(self) -> None: + assert self.process.stderr is not None + for line in self.process.stderr: + self.stderr_lines.append(line.rstrip()) + + def call(self, method: str, params: dict | None = None) -> object: + if self.exited.is_set(): + raise RuntimeError(self._failure(f"agent exited before {method}")) + with self.request_lock: + self.request_id += 1 + request_id = self.request_id + response_queue: queue.Queue = queue.Queue(maxsize=1) + self.pending[request_id] = response_queue + request = { + "jsonrpc": "2.0", + "id": request_id, + "method": method, + "params": params or {}, + } + try: + assert self.process.stdin is not None + with self.write_lock: + self.process.stdin.write(json.dumps(request, separators=(",", ":")) + "\n") + self.process.stdin.flush() + response = response_queue.get(timeout=env_float("BENCH_RPC_TIMEOUT", 180.0)) + except queue.Empty as error: + raise TimeoutError(self._failure(f"timed out during {method}")) from error + finally: + with self.request_lock: + self.pending.pop(request_id, None) + if isinstance(response, Exception): + raise response + if response.get("error") is not None: + raise RuntimeError( + f"{self.candidate.name} {method}: " + f"{json.dumps(response['error'], ensure_ascii=False)}" + ) + return response.get("result") + + def rss_kib(self) -> int: + if self.candidate.rss_command: + output = subprocess.check_output( + self.candidate.rss_command, + shell=True, + text=True, + ).strip() + return int(output or "0") + status_path = Path(f"/proc/{self.process.pid}/status") + if status_path.is_file(): + for line in status_path.read_text().splitlines(): + if line.startswith("VmRSS:"): + return int(line.split()[1]) + output = subprocess.check_output( + ["ps", "-o", "rss=", "-p", str(self.process.pid)], + text=True, + ).strip() + return int(output or "0") + + def close(self) -> bool: + if self.process.poll() is not None: + return True + try: + self.call("shutdown") + except Exception: + pass + try: + self.process.wait(timeout=3) + return True + except subprocess.TimeoutExpired: + self.process.terminate() + try: + self.process.wait(timeout=2) + except subprocess.TimeoutExpired: + self.process.kill() + self.process.wait(timeout=5) + return False + + def _failure(self, message: str) -> str: + stderr = "\n".join(self.stderr_lines[-20:]) + return f"{self.candidate.name}: {message}\n{stderr}".rstrip() + + +class RSSMonitor: + def __init__(self, process: AgentProcess): + self.process = process + self.peak_kib = 0 + self.stop_event = threading.Event() + self.thread = threading.Thread(target=self._run, daemon=True) + + def __enter__(self): + self.peak_kib = self.process.rss_kib() + self.thread.start() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.stop_event.set() + self.thread.join(timeout=1) + try: + self.peak_kib = max(self.peak_kib, self.process.rss_kib()) + except (OSError, subprocess.SubprocessError, ValueError): + pass + + def _run(self) -> None: + interval = env_float("BENCH_RSS_INTERVAL", 0.02) + while not self.stop_event.wait(interval): + try: + self.peak_kib = max(self.peak_kib, self.process.rss_kib()) + except (OSError, subprocess.SubprocessError, ValueError): + return + + +def main() -> None: + candidates = configured_candidates() + connection = connection_params() + startup_iterations = env_int("BENCH_STARTUPS", 8) + connect_iterations = env_int("BENCH_CONNECTS", 8) + rounds = env_int("BENCH_ROUNDS", 3) + warmups = env_int("BENCH_WARMUPS", 2) + workloads = configured_workloads(connection["database"]) + concurrency_levels = env_int_list("BENCH_CONCURRENCY", [1, 8, 32]) + + startup_samples = {candidate.name: [] for candidate in candidates} + connect_samples = {candidate.name: [] for candidate in candidates} + workload_samples = { + candidate.name: {workload["name"]: [] for workload in workloads} + for candidate in candidates + } + concurrency_samples = { + candidate.name: {str(level): [] for level in concurrency_levels} + for candidate in candidates + } + process_metrics = {candidate.name: [] for candidate in candidates} + + for iteration in range(startup_iterations): + for candidate in rotated(candidates, iteration): + startup_samples[candidate.name].append(benchmark_startup(candidate)) + + for iteration in range(connect_iterations): + for candidate in rotated(candidates, iteration): + connect_samples[candidate.name].append(benchmark_connect(candidate, connection)) + + for round_index in range(rounds): + for candidate in rotated(candidates, round_index): + process = AgentProcess(candidate) + shutdown_clean = False + metrics = None + try: + process.call("connect", connection) + idle_rss_kib = process.rss_kib() + with RSSMonitor(process) as monitor: + for workload in workloads: + sample = benchmark_workload(process, workload, warmups) + workload_samples[candidate.name][workload["name"]].append(sample) + process.call("disconnect") + metrics = { + "idle_rss_kib": idle_rss_kib, + "peak_rss_kib": monitor.peak_kib, + } + finally: + shutdown_clean = process.close() + if metrics is not None: + metrics["shutdown_exited_within_3s"] = shutdown_clean + process_metrics[candidate.name].append(metrics) + + for level in concurrency_levels: + for candidate in rotated(candidates, round_index + level): + concurrency_samples[candidate.name][str(level)].append( + benchmark_concurrency(candidate, connection, level, warmups) + ) + + results = [] + for candidate in candidates: + name = candidate.name + results.append( + { + "candidate": name, + "command": candidate.command, + "artifact_bytes": candidate.artifact.stat().st_size, + "startup": summarize_latencies(startup_samples[name]), + "connect": summarize_latencies(connect_samples[name]), + "process": summarize_process_metrics(process_metrics[name]), + "workloads": [ + summarize_rounds(workload["name"], workload_samples[name][workload["name"]]) + for workload in workloads + ], + "concurrency": [ + summarize_rounds(f"concurrency_{level}", concurrency_samples[name][str(level)]) + for level in concurrency_levels + ], + } + ) + + output = { + "host": os.uname().nodename, + "server": env_default("HIVE_SERVER", f"{connection['host']}:{connection['port']}"), + "database": connection["database"], + "table": env_default("HIVE_BENCH_TABLE", "agent_bench"), + "startup_iterations": startup_iterations, + "connect_iterations": connect_iterations, + "rounds": rounds, + "warmups": warmups, + "concurrency_levels": concurrency_levels, + "results": results, + } + json.dump(output, sys.stdout, ensure_ascii=False, indent=2) + sys.stdout.write("\n") + + +def configured_candidates() -> list[Candidate]: + selected = { + item.strip() + for item in env_default("BENCH_CANDIDATES", "go,jdbc").split(",") + if item.strip() + } + candidates = [] + if "go" in selected: + artifact = required_path("GO_AGENT") + raw_command = os.getenv("GO_AGENT_COMMAND", "") + command = shlex.split(raw_command) if raw_command else [str(artifact)] + candidates.append( + Candidate("go-native", command, artifact, os.getenv("GO_RSS_COMMAND", "")) + ) + if "jdbc" in selected: + artifact = required_path("JDBC_AGENT_JAR") + raw_command = os.getenv("JDBC_AGENT_COMMAND", "") + command = ( + shlex.split(raw_command) + if raw_command + else [env_default("JAVA_BIN", "java"), "-jar", str(artifact)] + ) + candidates.append( + Candidate("jdbc-java", command, artifact, os.getenv("JDBC_RSS_COMMAND", "")) + ) + if not candidates: + raise ValueError("BENCH_CANDIDATES selected no candidates") + return candidates + + +def connection_params() -> dict: + return { + "host": env_default("HIVE_HOST", "127.0.0.1"), + "port": env_int("HIVE_PORT", 10000), + "database": env_default("HIVE_DATABASE", "dbx_agent_bench"), + "username": os.getenv("HIVE_USERNAME", ""), + "password": os.getenv("HIVE_PASSWORD", ""), + "url_params": env_default("HIVE_URL_PARAMS", "auth=noSasl"), + "connection_string": os.getenv("HIVE_CONNECTION_STRING", ""), + "ssl": env_bool("HIVE_SSL", False), + "ca_cert_path": os.getenv("HIVE_CA_CERT_PATH", ""), + "client_cert_path": os.getenv("HIVE_CLIENT_CERT_PATH", ""), + "client_key_path": os.getenv("HIVE_CLIENT_KEY_PATH", ""), + "connect_timeout_secs": env_int("HIVE_CONNECT_TIMEOUT", 30), + } + + +def configured_workloads(database: str) -> list[dict]: + table = env_default("HIVE_BENCH_TABLE", "agent_bench") + qualified = f"`{database}`.`{table}`" + workloads = [ + query_workload("select_one", "SELECT 1 AS value", 1, 40), + query_workload("rows_100", f"SELECT id, payload FROM {qualified} LIMIT 100", 100, 20), + query_workload("rows_1000", f"SELECT id, payload FROM {qualified} LIMIT 1000", 1000, 10), + query_workload("rows_10000", f"SELECT id, payload FROM {qualified} LIMIT 10000", 10000, 3), + { + "name": "list_databases", + "kind": "rpc", + "method": "list_databases", + "params": {}, + "count": env_int("BENCH_LIST_DATABASES_COUNT", 20), + }, + { + "name": "list_tables", + "kind": "rpc", + "method": "list_tables", + "params": {"schema": database}, + "count": env_int("BENCH_LIST_TABLES_COUNT", 20), + }, + { + "name": "page_10000_by_500", + "kind": "paged", + "sql": env_default("BENCH_PAGE_SQL", f"SELECT id, payload FROM {qualified} LIMIT 10000"), + "max_rows": 10000, + "page_size": env_int("BENCH_PAGE_SIZE", 500), + "count": env_int("BENCH_PAGE_COUNT", 3), + }, + ] + return workloads + + +def query_workload(name: str, fallback_sql: str, max_rows: int, fallback_count: int) -> dict: + suffix = name.upper() + return { + "name": name, + "kind": "rpc", + "method": "execute_query", + "params": { + "sql": env_default(f"BENCH_{suffix}_SQL", fallback_sql), + "maxRows": max_rows, + "fetchSize": min(max_rows, env_int("BENCH_FETCH_SIZE", 1000)), + }, + "count": env_int(f"BENCH_{suffix}_COUNT", fallback_count), + } + + +def benchmark_startup(candidate: Candidate) -> float: + started = time.perf_counter() + process = AgentProcess(candidate) + elapsed = elapsed_ms(started) + process.close() + return elapsed + + +def benchmark_connect(candidate: Candidate, connection: dict) -> float: + process = AgentProcess(candidate) + try: + started = time.perf_counter() + process.call("connect", connection) + return elapsed_ms(started) + finally: + process.close() + + +def benchmark_workload(process: AgentProcess, workload: dict, warmups: int) -> dict: + for _ in range(warmups): + execute_workload(process, workload) + samples = [] + started = time.perf_counter() + for _ in range(workload["count"]): + operation_started = time.perf_counter() + execute_workload(process, workload) + samples.append(elapsed_ms(operation_started)) + elapsed = time.perf_counter() - started + return sample_result(workload["count"], elapsed, samples) + + +def benchmark_concurrency( + candidate: Candidate, + connection: dict, + concurrency: int, + warmups: int, +) -> dict: + process = AgentProcess(candidate) + session_ids = [f"bench-{concurrency}-{index}" for index in range(concurrency)] + qualified = ( + f"`{connection['database']}`." + f"`{env_default('HIVE_BENCH_TABLE', 'agent_bench')}`" + ) + workload = { + "kind": "rpc", + "method": "execute_query", + "params": { + "sql": env_default( + "BENCH_CONCURRENCY_SQL", + "SELECT 1 AS value", + ), + "maxRows": 1, + }, + } + operations_per_worker = env_int("BENCH_CONCURRENCY_OPS_PER_WORKER", 8) + try: + for session_id in session_ids: + process.call("open_session", {**connection, "agentSessionId": session_id}) + for session_id in session_ids: + for _ in range(warmups): + execute_workload(process, workload, session_id) + with RSSMonitor(process) as monitor: + started = time.perf_counter() + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = [ + executor.submit( + concurrency_worker, + process, + workload, + session_id, + operations_per_worker, + ) + for session_id in session_ids + ] + samples = [sample for future in futures for sample in future.result()] + elapsed = time.perf_counter() - started + result = sample_result(len(samples), elapsed, samples) + result["concurrency"] = concurrency + result["peak_rss_kib"] = monitor.peak_kib + return result + finally: + for session_id in session_ids: + try: + process.call("close_session", {"agentSessionId": session_id}) + except Exception: + pass + process.close() + + +def concurrency_worker( + process: AgentProcess, + workload: dict, + session_id: str, + operations: int, +) -> list[float]: + samples = [] + for _ in range(operations): + started = time.perf_counter() + execute_workload(process, workload, session_id) + samples.append(elapsed_ms(started)) + return samples + + +def execute_workload( + process: AgentProcess, + workload: dict, + agent_session_id: str = "", +) -> object: + params = dict(workload.get("params", {})) + if agent_session_id: + params["agentSessionId"] = agent_session_id + if workload["kind"] == "rpc": + return process.call(workload["method"], params) + if workload["kind"] != "paged": + raise ValueError(f"unknown workload kind: {workload['kind']}") + first = process.call( + "execute_query_page", + { + "sql": workload["sql"], + "maxRows": workload["max_rows"], + "pageSize": workload["page_size"], + **({"agentSessionId": agent_session_id} if agent_session_id else {}), + }, + ) + rows = len(first.get("rows", [])) + session_id = first.get("session_id") + has_more = first.get("has_more", False) + try: + while has_more: + page = process.call( + "fetch_query_page", + { + "sessionId": session_id, + "pageSize": workload["page_size"], + **({"agentSessionId": agent_session_id} if agent_session_id else {}), + }, + ) + rows += len(page.get("rows", [])) + session_id = page.get("session_id") + has_more = page.get("has_more", False) + finally: + if session_id: + process.call( + "close_query_session", + { + "sessionId": session_id, + **({"agentSessionId": agent_session_id} if agent_session_id else {}), + }, + ) + if rows != workload["max_rows"]: + raise RuntimeError( + f"{process.candidate.name} paged query returned {rows} rows, " + f"expected {workload['max_rows']}" + ) + return rows + + +def sample_result(count: int, elapsed: float, samples: list[float]) -> dict: + summary = summarize_latencies(samples) + summary.update( + { + "count": count, + "elapsed_ms": elapsed * 1000, + "ops_per_sec": count / elapsed, + } + ) + return summary + + +def summarize_latencies(samples: list[float]) -> dict: + ordered = sorted(samples) + return { + "samples_ms": samples, + "mean_ms": statistics.mean(samples), + "p50_ms": percentile(ordered, 0.50), + "p95_ms": percentile(ordered, 0.95), + "p99_ms": percentile(ordered, 0.99), + "min_ms": ordered[0], + "max_ms": ordered[-1], + } + + +def summarize_rounds(name: str, rounds: list[dict]) -> dict: + latencies = [sample for round_result in rounds for sample in round_result["samples_ms"]] + elapsed = sum(round_result["elapsed_ms"] for round_result in rounds) / 1000 + result = summarize_latencies(latencies) + result.update( + { + "name": name, + "rounds": rounds, + "count": len(latencies), + "elapsed_ms": elapsed * 1000, + "ops_per_sec": len(latencies) / elapsed, + } + ) + peak_values = [round_result.get("peak_rss_kib", 0) for round_result in rounds] + if any(peak_values): + result["peak_rss_kib"] = max(peak_values) + return result + + +def summarize_process_metrics(samples: list[dict]) -> dict: + return { + "idle_rss_kib": summarize_numbers([sample["idle_rss_kib"] for sample in samples]), + "peak_rss_kib": summarize_numbers([sample["peak_rss_kib"] for sample in samples]), + "shutdown_exited_within_3s": all( + sample["shutdown_exited_within_3s"] for sample in samples + ), + "rounds": samples, + } + + +def summarize_numbers(values: list[int]) -> dict: + return { + "min": min(values), + "median": statistics.median(values), + "max": max(values), + } + + +def rotated(values: list[Candidate], offset: int) -> list[Candidate]: + if not values: + return [] + shift = offset % len(values) + return values[shift:] + values[:shift] + + +def percentile(values: list[float], fraction: float) -> float: + if not values: + return 0.0 + index = min(len(values) - 1, max(0, round((len(values) - 1) * fraction))) + return values[index] + + +def elapsed_ms(started: float) -> float: + return (time.perf_counter() - started) * 1000 + + +def required_path(name: str) -> Path: + value = os.getenv(name, "") + if not value: + raise ValueError(f"{name} is required") + path = Path(value).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(path) + return path + + +def env_default(name: str, fallback: str) -> str: + return os.getenv(name, "") or fallback + + +def env_int(name: str, fallback: int) -> int: + value = int(env_default(name, str(fallback))) + if value < 1: + raise ValueError(f"{name} must be positive") + return value + + +def env_int_list(name: str, fallback: list[int]) -> list[int]: + raw = os.getenv(name, "") + values = fallback if not raw else [int(value.strip()) for value in raw.split(",")] + if not values or any(value < 1 for value in values): + raise ValueError(f"{name} must contain positive integers") + return values + + +def env_float(name: str, fallback: float) -> float: + value = float(env_default(name, str(fallback))) + if value <= 0: + raise ValueError(f"{name} must be positive") + return value + + +def env_bool(name: str, fallback: bool) -> bool: + raw = os.getenv(name) + if raw is None or raw == "": + return fallback + normalized = raw.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + raise ValueError(f"{name} must be a boolean") + + +if __name__ == "__main__": + main() diff --git a/agents/drivers/hive-go/bench/functional_probe.py b/agents/drivers/hive-go/bench/functional_probe.py new file mode 100644 index 0000000000..66562a0946 --- /dev/null +++ b/agents/drivers/hive-go/bench/functional_probe.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +import hashlib +import json +import os +import sys +from pathlib import Path + +from agent_compare import AgentProcess, configured_candidates, connection_params, env_default, env_int + + +def main() -> None: + os.environ.setdefault("BENCH_CANDIDATES", "go") + connection = connection_params() + candidates = configured_candidates() + results = {candidate.name: probe_candidate(candidate, connection) for candidate in candidates} + output = { + "server": env_default("HIVE_SERVER", f"{connection['host']}:{connection['port']}"), + "connection": sanitized_connection(connection), + "artifacts": { + candidate.name: { + "path": str(candidate.artifact), + "sha256": sha256(candidate.artifact), + "size_bytes": candidate.artifact.stat().st_size, + } + for candidate in candidates + }, + "results": results, + "parity": compare_results(results), + } + json.dump(output, sys.stdout, ensure_ascii=False, indent=2) + sys.stdout.write("\n") + if any(not result.get("ok") for result in results.values()) or not output["parity"]["ok"]: + raise SystemExit(1) + + +def probe_candidate(candidate, connection: dict) -> dict: + process = None + session_id = f"functional-probe-{candidate.name}" + result = {"ok": False} + try: + process = AgentProcess(candidate) + result["test_connection"] = process.call("test_connection", connection) + process.call("open_session", {"agentSessionId": session_id, **connection}) + result["validate_session"] = process.call("validate_session", {"agentSessionId": session_id}) + result["select_one"] = normalized_query( + process.call( + "execute_query", + { + "agentSessionId": session_id, + "sql": env_default("PROBE_SELECT_SQL", "SELECT 1 AS value"), + "maxRows": env_int("PROBE_SELECT_MAX_ROWS", 10), + "fetchSize": env_int("PROBE_FETCH_SIZE", 10), + }, + ) + ) + result["databases"] = sorted( + item.get("name", "") + for item in process.call("list_databases", {"agentSessionId": session_id}) + ) + schema = env_default("PROBE_SCHEMA", connection["database"]) + result["tables"] = sorted( + item.get("name", "") + for item in process.call( + "list_tables", + {"agentSessionId": session_id, "schema": schema}, + ) + ) + result["paging"] = probe_paging(process, session_id) + result["invalid_sql"] = probe_failure_semantics(process, session_id) + result["after_failure"] = normalized_query( + process.call( + "execute_query", + { + "agentSessionId": session_id, + "sql": env_default("PROBE_AFTER_FAILURE_SQL", "SELECT 2 AS value"), + "maxRows": 10, + "fetchSize": env_int("PROBE_FETCH_SIZE", 10), + }, + ) + ) + process.call("close_session", {"agentSessionId": session_id}) + result["ok"] = True + except Exception as error: + result["error"] = str(error) + finally: + if process is not None: + result["clean_shutdown"] = process.close() + return result + + +def probe_paging(process: AgentProcess, agent_session_id: str) -> dict: + page_size = env_int("PROBE_PAGE_SIZE", 2) + first = process.call( + "execute_query_page", + { + "agentSessionId": agent_session_id, + "sql": env_default( + "PROBE_PAGE_SQL", + "SELECT id, payload FROM dbx_agent_bench.agent_bench LIMIT 3", + ), + "maxRows": env_int("PROBE_PAGE_MAX_ROWS", 3), + "pageSize": page_size, + }, + ) + pages = [first] + query_session_id = first.get("session_id") + while pages[-1].get("has_more"): + pages.append( + process.call( + "fetch_query_page", + { + "agentSessionId": agent_session_id, + "sessionId": query_session_id, + "pageSize": page_size, + }, + ) + ) + return { + "columns": first.get("columns", []), + "column_types": first.get("column_types", []), + "rows": [row for page in pages for row in page.get("rows", [])], + "page_count": len(pages), + "has_more_final": pages[-1].get("has_more", False), + "truncated": any(page.get("truncated", False) for page in pages), + } + + +def probe_failure_semantics(process: AgentProcess, agent_session_id: str) -> dict: + try: + process.call( + "execute_query", + { + "agentSessionId": agent_session_id, + "sql": env_default( + "PROBE_INVALID_SQL", + "SELECT * FROM dbx_missing_table_for_failure_semantics", + ), + "maxRows": 10, + "fetchSize": env_int("PROBE_FETCH_SIZE", 10), + }, + ) + except Exception as error: + return {"failed": True, "error": str(error)} + return {"failed": False, "error": ""} + + +def normalized_query(result: dict) -> dict: + return { + "columns": result.get("columns", []), + "column_types": result.get("column_types", []), + "rows": result.get("rows", []), + "truncated": result.get("truncated", False), + } + + +def compare_results(results: dict) -> dict: + successful = [result for result in results.values() if result.get("ok")] + if len(successful) < 2: + return {"ok": len(results) == 1 and len(successful) == 1, "differences": []} + baseline = successful[0] + differences = [] + for field in ["select_one", "databases", "tables", "paging", "after_failure"]: + expected = baseline.get(field) + for candidate_name, candidate_result in results.items(): + if candidate_result.get("ok") and candidate_result.get(field) != expected: + differences.append( + { + "candidate": candidate_name, + "field": field, + "expected": expected, + "actual": candidate_result.get(field), + } + ) + for candidate_name, candidate_result in results.items(): + if candidate_result.get("ok") and not candidate_result.get("invalid_sql", {}).get("failed"): + differences.append( + { + "candidate": candidate_name, + "field": "invalid_sql.failed", + "expected": True, + "actual": False, + } + ) + return {"ok": not differences, "differences": differences} + + +def sanitized_connection(connection: dict) -> dict: + result = dict(connection) + if result.get("password"): + result["password"] = "***" + return result + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +if __name__ == "__main__": + main() diff --git a/agents/drivers/hive-go/bench/kdc_fixture/main.go b/agents/drivers/hive-go/bench/kdc_fixture/main.go new file mode 100644 index 0000000000..667d0eecca --- /dev/null +++ b/agents/drivers/hive-go/bench/kdc_fixture/main.go @@ -0,0 +1,99 @@ +package main + +import ( + "encoding/json" + "flag" + "fmt" + "log" + "os" + "os/signal" + "path/filepath" + "syscall" + + "github.com/jcmturner/krb5test" +) + +type fixtureInfo struct { + Realm string `json:"realm"` + Address string `json:"address"` + ConfigPath string `json:"config_path"` + KeytabPath string `json:"keytab_path"` + ClientPrincipal string `json:"client_principal"` + ServicePrincipal string `json:"service_principal"` + ZooKeeperPrincipal string `json:"zookeeper_principal"` +} + +func main() { + directory := flag.String("dir", "", "directory for generated Kerberos fixture files") + flag.Parse() + if *directory == "" { + log.Fatal("-dir is required") + } + if err := os.MkdirAll(*directory, 0o700); err != nil { + log.Fatal(err) + } + + logger := log.New(os.Stderr, "kdc: ", log.LstdFlags) + kdc, err := krb5test.NewKDC(map[string][]string{ + "alice": nil, + "hive/localhost": nil, + "zookeeper/localhost": nil, + }, logger) + if err != nil { + log.Fatal(err) + } + kdc.KRB5Conf.LibDefaults.UDPPreferenceLimit = 1 + kdc.Start() + defer kdc.Close() + + configPath := filepath.Join(*directory, "krb5.conf") + keytabPath := filepath.Join(*directory, "fixture.keytab") + config := fmt.Sprintf(`[libdefaults] + default_realm = %s + dns_lookup_realm = false + dns_lookup_kdc = false + rdns = false + udp_preference_limit = 1 + default_tgs_enctypes = aes256-cts-hmac-sha1-96 + default_tkt_enctypes = aes256-cts-hmac-sha1-96 + permitted_enctypes = aes256-cts-hmac-sha1-96 + +[realms] + %s = { + kdc = %s + } + +[domain_realm] + .localhost = %s + localhost = %s +`, kdc.Realm, kdc.Realm, kdc.TCPListener.Addr().String(), kdc.Realm, kdc.Realm) + if err := os.WriteFile(configPath, []byte(config), 0o644); err != nil { + log.Fatal(err) + } + keytab, err := kdc.Keytab.Marshal() + if err != nil { + log.Fatal(err) + } + if err := os.WriteFile(keytabPath, keytab, 0o600); err != nil { + log.Fatal(err) + } + info := fixtureInfo{ + Realm: kdc.Realm, + Address: kdc.TCPListener.Addr().String(), + ConfigPath: configPath, + KeytabPath: keytabPath, + ClientPrincipal: "alice@" + kdc.Realm, + ServicePrincipal: "hive/localhost@" + kdc.Realm, + ZooKeeperPrincipal: "zookeeper/localhost@" + kdc.Realm, + } + if err := json.NewEncoder(os.Stdout).Encode(info); err != nil { + log.Fatal(err) + } + if err := os.Stdout.Sync(); err != nil { + log.Fatal(err) + } + + signals := make(chan os.Signal, 1) + signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM) + <-signals +} diff --git a/agents/drivers/hive-go/config.go b/agents/drivers/hive-go/config.go new file mode 100644 index 0000000000..63e995b594 --- /dev/null +++ b/agents/drivers/hive-go/config.go @@ -0,0 +1,1180 @@ +package main + +import ( + "bufio" + "crypto/tls" + "crypto/x509" + "encoding/base64" + "errors" + "fmt" + "io" + "net" + "net/url" + "os" + "path/filepath" + "regexp" + "runtime" + "strconv" + "strings" + "time" +) + +const ( + defaultHivePort = 10000 + defaultHiveDatabase = "default" + defaultHiveHTTPPath = "cliservice" + defaultHiveService = "hive" + defaultZooKeeperNamespace = "hiveserver2" + defaultConnectTimeout = 15 * time.Second + defaultRetryInterval = time.Second + defaultBrowserSSOTimeout = 120 * time.Second + defaultCookieName = "hive.server2.auth" +) + +type connectParams struct { + Host string `json:"host"` + Port int `json:"port"` + Database string `json:"database"` + Username string `json:"username"` + Password string `json:"password"` + URLParams string `json:"url_params"` + ConnectionString string `json:"connection_string"` + SSL bool `json:"ssl"` + CACertPath string `json:"ca_cert_path"` + ClientCertPath string `json:"client_cert_path"` + ClientKeyPath string `json:"client_key_path"` + ConnectTimeout int `json:"connect_timeout_secs"` + AgentJavaOptions []string `json:"agent_java_options"` + SessionRole string `json:"sessionRole"` +} + +type endpoint struct { + Host string + Port int + TransportMode string + HTTPPath string + Auth string + Principal string + SSL bool +} + +func (value endpoint) address() string { + return net.JoinHostPort(value.Host, strconv.Itoa(value.Port)) +} + +type kerberosConfig struct { + Enabled bool + ServerPrincipal string + ServerPrincipalExplicit bool + ClientPrincipal string + Service string + ServerName string + Realm string + ConfigPath string + JAASConfigPath string + KeytabPath string + CCachePath string + Password string + AuthorizationID string + QOP string + UseKeytab bool + UseTicketCache bool + UseSSPI bool + CanonicalHostname bool + ChannelBinding bool + DisablePAFXFAST bool +} + +type zooKeeperKerberosConfig struct { + Enabled bool + Service string + ServerPrincipal string + Realm string + CanonicalHostname bool +} + +type connectionConfig struct { + Endpoints []endpoint + Database string + Username string + Password string + Auth string + TransportMode string + HTTPPath string + TLSConfig *tls.Config + AuthExplicit bool + TransportModeExplicit bool + HTTPPathExplicit bool + TLSExplicit bool + HiveConfiguration map[string]string + ServiceDiscoveryMode string + ZooKeeperNamespace string + ZooKeeperAuthScheme string + ZooKeeperAuth string + ZooKeeperTLSConfig *tls.Config + ZooKeeperKerberos zooKeeperKerberosConfig + ConnectTimeout time.Duration + SocketTimeout time.Duration + FetchSize int + MaxMessageSize int32 + Retries int + RetryInterval time.Duration + HTTPHeaders map[string]string + HTTPCookies map[string]string + RequestTracking bool + CookieAuth bool + CookieName string + JWT string + DelegationToken string + BrowserToken string + BrowserClientID string + BrowserResponsePort int + BrowserResponseTimeout time.Duration + BrowserDisableSSLCheck bool + InitStatements []string + Kerberos kerberosConfig +} + +func parseConnectionConfig(params connectParams) (connectionConfig, error) { + config := connectionConfig{ + Database: strings.TrimSpace(params.Database), + Username: params.Username, + Password: params.Password, + Auth: "NONE", + TransportMode: "binary", + HTTPPath: defaultHiveHTTPPath, + HiveConfiguration: map[string]string{}, + FetchSize: defaultFetchSize, + Retries: 1, + RetryInterval: defaultRetryInterval, + BrowserResponseTimeout: defaultBrowserSSOTimeout, + HTTPHeaders: map[string]string{}, + HTTPCookies: map[string]string{}, + CookieAuth: true, + CookieName: defaultCookieName, + ZooKeeperNamespace: defaultZooKeeperNamespace, + ZooKeeperKerberos: zooKeeperKerberosConfig{ + Service: "zookeeper", + CanonicalHostname: true, + }, + ConnectTimeout: defaultConnectTimeout, + Kerberos: kerberosConfig{ + Service: defaultHiveService, + QOP: "auth", + CanonicalHostname: true, + }, + } + if config.Database == "" { + config.Database = defaultHiveDatabase + } + if params.ConnectTimeout > 0 { + config.ConnectTimeout = time.Duration(params.ConnectTimeout) * time.Second + } + + parsed, err := parseHiveConnectionString(params.ConnectionString) + if err != nil { + return connectionConfig{}, err + } + if len(parsed.endpoints) > 0 { + config.Endpoints = parsed.endpoints + } else { + port := params.Port + if port <= 0 { + port = defaultHivePort + } + if host := strings.TrimSpace(params.Host); host != "" { + for _, value := range splitEndpoints(host) { + parsedEndpoint, endpointErr := parseEndpoint(value, port) + if endpointErr != nil { + return connectionConfig{}, endpointErr + } + config.Endpoints = append(config.Endpoints, parsedEndpoint) + } + } + } + if parsed.database != "" { + config.Database = parsed.database + } + if parsed.username != "" { + config.Username = parsed.username + } + if parsed.password != "" { + config.Password = parsed.password + } + + urlSections := parseHiveParameterSections(params.URLParams) + values := mergeHiveParameters(parsed.parameters, urlSections.session) + hiveConfs := mergeHiveAssignments(parsed.hiveConfs, urlSections.hiveConfs) + hiveVars := mergeHiveAssignments(parsed.hiveVars, urlSections.hiveVars) + if value, exists := firstParameter(values, "user", "username"); exists { + config.Username = value + } + if value, exists := firstParameter(values, "password"); exists { + config.Password = value + } + if err := applyHiveParameters(&config, values, hiveConfs); err != nil { + return connectionConfig{}, err + } + applyOpenSessionVariables(&config, values, hiveConfs, hiveVars) + if err := applyDelegationToken(&config, values); err != nil { + return connectionConfig{}, err + } + applyKerberosJavaOptions(&config.Kerberos, params.AgentJavaOptions) + applyZooKeeperKerberosJavaOptions(&config.ZooKeeperKerberos, params.AgentJavaOptions) + applyKerberosEnvironment(&config.Kerberos) + if err := finalizeKerberosConfig(&config); err != nil { + return connectionConfig{}, err + } + if len(config.Endpoints) == 0 { + return connectionConfig{}, errors.New("Hive host is required") + } + + tlsConfig, err := buildTLSConfig(params, values, config.Endpoints[0].Host) + if err != nil { + return connectionConfig{}, err + } + config.TLSConfig = tlsConfig + zooKeeperTLSConfig, err := buildZooKeeperTLSConfig(values) + if err != nil { + return connectionConfig{}, err + } + config.ZooKeeperTLSConfig = zooKeeperTLSConfig + return config, nil +} + +type parsedHiveConnection struct { + endpoints []endpoint + database string + username string + password string + parameters map[string]string + hiveConfs map[string]string + hiveVars map[string]string +} + +type hiveParameterSections struct { + session map[string]string + hiveConfs map[string]string + hiveVars map[string]string +} + +func newParsedHiveConnection() parsedHiveConnection { + return parsedHiveConnection{ + parameters: map[string]string{}, + hiveConfs: map[string]string{}, + hiveVars: map[string]string{}, + } +} + +func parseHiveConnectionString(raw string) (parsedHiveConnection, error) { + value := strings.TrimSpace(raw) + if value == "" { + return newParsedHiveConnection(), nil + } + if strings.HasPrefix(strings.ToLower(value), "jdbc:hive2://") { + value = value[len("jdbc:hive2://"):] + } else if strings.HasPrefix(strings.ToLower(value), "hive://") { + parsedURL, err := url.Parse(value) + if err != nil { + return parsedHiveConnection{}, fmt.Errorf("invalid Hive connection string: %w", err) + } + result := newParsedHiveConnection() + if parsedURL.User != nil { + result.username = parsedURL.User.Username() + result.password, _ = parsedURL.User.Password() + } + port := defaultHivePort + if parsedURL.Port() != "" { + parsedPort, err := strconv.Atoi(parsedURL.Port()) + if err != nil { + return parsedHiveConnection{}, fmt.Errorf("invalid Hive port: %w", err) + } + port = parsedPort + } + result.endpoints = []endpoint{{Host: parsedURL.Hostname(), Port: port}} + result.database = strings.Trim(parsedURL.Path, "/") + for key, entries := range parsedURL.Query() { + if len(entries) > 0 { + setCaseInsensitive(result.parameters, key, entries[len(entries)-1]) + } + } + return result, nil + } else { + return parsedHiveConnection{}, errors.New("Hive connection string must start with jdbc:hive2:// or hive://") + } + + result := newParsedHiveConnection() + if fragment := strings.IndexByte(value, '#'); fragment >= 0 { + result.hiveVars = parseHiveAssignments(value[fragment+1:], false) + value = value[:fragment] + } + if query := strings.IndexByte(value, '?'); query >= 0 { + result.hiveConfs = parseHiveAssignments(value[query+1:], false) + value = value[:query] + } + pathStart := strings.IndexByte(value, '/') + authority := value + pathAndParams := "" + if pathStart >= 0 { + authority = value[:pathStart] + pathAndParams = value[pathStart+1:] + } + if at := strings.LastIndex(authority, "@"); at >= 0 { + credentials := authority[:at] + authority = authority[at+1:] + if colon := strings.IndexByte(credentials, ':'); colon >= 0 { + result.username, _ = url.QueryUnescape(credentials[:colon]) + result.password, _ = url.QueryUnescape(credentials[colon+1:]) + } else { + result.username, _ = url.QueryUnescape(credentials) + } + } + for _, value := range splitEndpoints(authority) { + parsedEndpoint, err := parseEndpoint(value, defaultHivePort) + if err != nil { + return parsedHiveConnection{}, err + } + result.endpoints = append(result.endpoints, parsedEndpoint) + } + if separator := strings.IndexByte(pathAndParams, ';'); separator >= 0 { + result.database = strings.TrimSpace(pathAndParams[:separator]) + result.parameters = mergeHiveParameters(result.parameters, parseHiveParameters(pathAndParams[separator+1:])) + } else { + result.database = strings.TrimSpace(pathAndParams) + } + return result, nil +} + +func splitEndpoints(value string) []string { + parts := strings.Split(value, ",") + result := make([]string, 0, len(parts)) + for _, part := range parts { + if trimmed := strings.TrimSpace(part); trimmed != "" { + result = append(result, trimmed) + } + } + return result +} + +func parseEndpoint(value string, defaultPort int) (endpoint, error) { + value = strings.TrimSpace(value) + if value == "" { + return endpoint{}, errors.New("Hive endpoint is empty") + } + host := value + port := defaultPort + if parsedHost, parsedPort, err := net.SplitHostPort(value); err == nil { + host = parsedHost + parsed, parseErr := strconv.Atoi(parsedPort) + if parseErr != nil { + return endpoint{}, fmt.Errorf("invalid Hive endpoint %q: %w", value, parseErr) + } + port = parsed + } else if strings.Count(value, ":") == 1 { + parts := strings.SplitN(value, ":", 2) + parsed, parseErr := strconv.Atoi(parts[1]) + if parseErr != nil { + return endpoint{}, fmt.Errorf("invalid Hive endpoint %q: %w", value, parseErr) + } + host = parts[0] + port = parsed + } else if strings.HasPrefix(value, "[") && strings.HasSuffix(value, "]") { + host = strings.Trim(value, "[]") + } + if strings.TrimSpace(host) == "" || port <= 0 || port > 65535 { + return endpoint{}, fmt.Errorf("invalid Hive endpoint %q", value) + } + return endpoint{Host: host, Port: port}, nil +} + +func parseHiveParameters(raw string) map[string]string { + return parseHiveAssignments(raw, false) +} + +func parseHiveAssignments(raw string, lowercaseKeys bool) map[string]string { + result := map[string]string{} + trimmed := strings.Trim(strings.TrimSpace(raw), "?#&;") + for _, part := range strings.FieldsFunc(trimmed, func(char rune) bool { return char == ';' || char == '&' }) { + part = strings.TrimSpace(part) + if part == "" { + continue + } + key, value, found := strings.Cut(part, "=") + key = strings.TrimSpace(key) + if decoded, err := url.QueryUnescape(key); err == nil { + key = decoded + } + if lowercaseKeys { + key = strings.ToLower(key) + } + if key == "" { + continue + } + if found { + if decoded, err := url.QueryUnescape(strings.TrimSpace(value)); err == nil { + value = decoded + } + } else { + value = "" + } + if lowercaseKeys { + result[key] = value + } else { + setCaseInsensitive(result, key, value) + } + } + return result +} + +func parseHiveParameterSections(raw string) hiveParameterSections { + value := strings.TrimSpace(raw) + sections := hiveParameterSections{ + session: map[string]string{}, + hiveConfs: map[string]string{}, + hiveVars: map[string]string{}, + } + if fragment := strings.IndexByte(value, '#'); fragment >= 0 { + sections.hiveVars = parseHiveAssignments(value[fragment+1:], false) + value = value[:fragment] + } + if query := strings.IndexByte(value, '?'); query >= 0 { + sections.hiveConfs = parseHiveAssignments(value[query+1:], false) + value = value[:query] + } + sections.session = parseHiveParameters(value) + return sections +} + +func mergeHiveParameters(first, second map[string]string) map[string]string { + result := make(map[string]string, len(first)+len(second)) + for key, value := range first { + setCaseInsensitive(result, key, value) + } + for key, value := range second { + setCaseInsensitive(result, key, value) + } + return result +} + +func setCaseInsensitive(values map[string]string, key, value string) { + for existing := range values { + if strings.EqualFold(existing, key) { + delete(values, existing) + } + } + values[key] = value +} + +func mergeHiveAssignments(first, second map[string]string) map[string]string { + result := make(map[string]string, len(first)+len(second)) + for key, value := range first { + result[key] = value + } + for key, value := range second { + result[key] = value + } + return result +} + +func applyHiveParameters(config *connectionConfig, values, hiveConfs map[string]string) error { + if value := parameter(values, "auth"); value != "" { + config.Auth = strings.ToUpper(value) + config.AuthExplicit = true + } + if value := firstNonEmpty(parameter(values, "transportmode"), hiveAssignmentValue(hiveConfs, "hive.server2.transport.mode")); value != "" { + config.TransportMode = strings.ToLower(value) + config.TransportModeExplicit = true + } + if value := firstNonEmpty(parameter(values, "httppath"), hiveAssignmentValue(hiveConfs, "hive.server2.thrift.http.path")); value != "" { + config.HTTPPath = strings.TrimPrefix(value, "/") + config.HTTPPathExplicit = true + } + if value := parameter(values, "servicediscoverymode"); value != "" { + config.ServiceDiscoveryMode = strings.ToLower(value) + } + if value := parameter(values, "zookeepernamespace"); value != "" { + config.ZooKeeperNamespace = strings.Trim(value, "/") + } + if strings.EqualFold(config.ServiceDiscoveryMode, "zookeeperha") && parameter(values, "zookeepernamespace") == "" { + config.ZooKeeperNamespace = "hs2ActivePassiveHA" + } + if hasParameter(values, "ssl") { + config.TLSExplicit = true + } + config.ZooKeeperAuthScheme = parameter(values, "zookeeperauthscheme") + config.ZooKeeperAuth = parameter(values, "zookeeperauth") + config.HTTPHeaders = prefixedParameters(values, "http.header.") + config.HTTPCookies = prefixedParameters(values, "http.cookie.") + config.RequestTracking = parameterBool(values, "requesttrack") + if value, exists := firstParameter(values, "cookieauth"); exists { + config.CookieAuth = !strings.EqualFold(value, "false") + } + config.CookieName = firstNonEmpty(parameter(values, "cookiename"), defaultCookieName) + config.JWT = firstNonEmpty(parameter(values, "jwt"), os.Getenv("JWT")) + config.BrowserToken = firstNonEmpty(parameter(values, "browsertoken"), parameter(values, "token")) + config.BrowserClientID = parameter(values, "browserclientidentifier") + if value := parameter(values, "browserresponseport"); value != "" { + parsed, err := strconv.Atoi(value) + if err != nil || parsed < 0 || parsed > 65535 { + return fmt.Errorf("invalid Hive browserResponsePort %q: expected 0-65535", value) + } + config.BrowserResponsePort = parsed + } + if value := parameter(values, "browserresponsetimeout"); value != "" { + parsed, err := strconv.ParseInt(value, 10, 64) + if err != nil || parsed <= 0 { + return fmt.Errorf("invalid Hive browserResponseTimeout %q: expected positive seconds", value) + } + config.BrowserResponseTimeout = time.Duration(parsed) * time.Second + } + config.BrowserDisableSSLCheck = parameterBool(values, "browserdisablesslcheck") + if strings.EqualFold(config.Auth, "JWT") && config.JWT == "" { + return errors.New("Hive JWT authentication requires jwt or the JWT environment variable") + } + if value := parameter(values, "fetchsize"); value != "" { + parsed, err := strconv.Atoi(value) + if err != nil || parsed <= 0 { + return fmt.Errorf("invalid Hive fetchSize %q: expected a positive integer", value) + } + config.FetchSize = parsed + } + if value := parameter(values, "sockettimeout"); value != "" { + parsed, err := strconv.ParseInt(value, 10, 64) + if err != nil { + return fmt.Errorf("invalid Hive socketTimeout %q: expected seconds", value) + } + if parsed > 0 { + config.SocketTimeout = time.Duration(parsed) * time.Second + } + } + if value := parameter(values, "thrift.client.max.message.size"); value != "" { + parsed, err := strconv.ParseInt(value, 10, 32) + if err != nil { + return fmt.Errorf("invalid Hive thrift.client.max.message.size %q: expected bytes", value) + } + if parsed > 0 { + config.MaxMessageSize = int32(parsed) + } + } + if value := parameter(values, "retries"); value != "" { + parsed, err := strconv.Atoi(value) + if err == nil && parsed > 0 { + config.Retries = parsed + } + } + if value := parameter(values, "retryinterval"); value != "" { + parsed, err := strconv.ParseInt(value, 10, 64) + if err == nil && parsed >= 0 { + config.RetryInterval = time.Duration(parsed) * time.Millisecond + } + } + if value := parameter(values, "initfile"); value != "" { + statements, err := readHiveInitFile(value) + if err != nil { + return err + } + config.InitStatements = statements + } + + kerberos := &config.Kerberos + kerberos.ServerPrincipal = parameter(values, "principal") + kerberos.ServerPrincipalExplicit = kerberos.ServerPrincipal != "" + kerberos.ClientPrincipal = firstNonEmpty( + parameter(values, "kerberosprincipal"), + parameter(values, "clientprincipal"), + parameter(values, "userprincipal"), + ) + kerberos.Service = firstNonEmpty(parameter(values, "service"), serviceFromPrincipal(kerberos.ServerPrincipal), defaultHiveService) + kerberos.ServerName = parameter(values, "servername") + kerberos.Realm = firstNonEmpty(parameter(values, "realm"), realmFromPrincipal(kerberos.ClientPrincipal)) + kerberos.ConfigPath = firstNonEmpty(parameter(values, "krb5conf"), parameter(values, "kerberosconfig")) + kerberos.JAASConfigPath = parameter(values, "jaasconfig") + kerberos.KeytabPath = parameter(values, "keytab") + kerberos.CCachePath = firstNonEmpty(parameter(values, "ccache"), parameter(values, "credentialcache")) + kerberos.AuthorizationID = firstNonEmpty(parameter(values, "authorizationid"), parameter(values, "proxyuser")) + kerberos.QOP = firstNonEmpty( + parameter(values, "hive.server2.thrift.sasl.qop"), + hiveAssignmentValue(hiveConfs, "hive.server2.thrift.sasl.qop"), + parameter(values, "sasl.qop"), + parameter(values, "saslqop"), + "auth", + ) + kerberos.UseKeytab = parameterBool(values, "usekeytab") || kerberos.KeytabPath != "" + kerberos.UseTicketCache = parameterBool(values, "useticketcache") || kerberos.CCachePath != "" + kerberos.UseSSPI = parameterBool(values, "usesspi") + if hasParameter(values, "kerberosenablecanonicalhostnamecheck") { + kerberos.CanonicalHostname = parameterBool(values, "kerberosenablecanonicalhostnamecheck") + } + kerberos.ChannelBinding = parameterBool(values, "kerberoschannelbinding") || + parameterBool(values, "tlschannelbinding") || + parameterBool(values, "channelbinding") + kerberos.DisablePAFXFAST = parameterBool(values, "disablepafxfast") + if kerberos.ServerPrincipal != "" || strings.EqualFold(config.Auth, "KERBEROS") { + kerberos.Enabled = true + config.Auth = "KERBEROS" + } + + zooKeeperKerberos := &config.ZooKeeperKerberos + zooKeeperKerberos.Enabled = kerberos.ServerPrincipalExplicit + if value, exists := firstParameter(values, "hive.zookeeper.use.kerberos", "hiveconf:hive.zookeeper.use.kerberos"); exists { + zooKeeperKerberos.Enabled = booleanValue(value) + } else if value := hiveAssignmentValue(hiveConfs, "hive.zookeeper.use.kerberos"); value != "" { + zooKeeperKerberos.Enabled = booleanValue(value) + } + if value, exists := firstParameter(values, "zookeeper.sasl.client"); exists && !booleanValue(value) { + zooKeeperKerberos.Enabled = false + } + zooKeeperKerberos.Service = firstNonEmpty(parameter(values, "zookeeper.sasl.client.username"), "zookeeper") + zooKeeperKerberos.ServerPrincipal = parameter(values, "zookeeper.server.principal") + zooKeeperKerberos.Realm = parameter(values, "zookeeper.server.realm") + if hasParameter(values, "zookeeper.sasl.client.canonicalize.hostname") { + zooKeeperKerberos.CanonicalHostname = parameterBool(values, "zookeeper.sasl.client.canonicalize.hostname") + } + + return nil +} + +func applyOpenSessionVariables(config *connectionConfig, values, hiveConfs, hiveVars map[string]string) { + for key, value := range values { + lowerKey := strings.ToLower(key) + switch { + case strings.HasPrefix(lowerKey, "hiveconf:"): + config.HiveConfiguration["set:hiveconf:"+key[len("hiveconf:"):]] = value + case strings.HasPrefix(lowerKey, "hivevar:"): + config.HiveConfiguration["set:hivevar:"+key[len("hivevar:"):]] = value + } + } + for key, value := range hiveConfs { + if strings.EqualFold(key, "hive.server2.transport.mode") || strings.EqualFold(key, "hive.server2.thrift.http.path") { + continue + } + config.HiveConfiguration["set:hiveconf:"+key] = value + } + for key, value := range hiveVars { + config.HiveConfiguration["set:hivevar:"+key] = value + } + if proxyUser := firstNonEmpty(parameter(values, "proxyuser"), parameter(values, "hive.server2.proxy.user")); proxyUser != "" { + config.HiveConfiguration["hive.server2.proxy.user"] = proxyUser + } + if value := parameter(values, "hivecreateasexternallegacy"); value != "" { + config.HiveConfiguration["set:hiveconf:hive.create.as.external.legacy"] = strings.ToLower(value) + } + if value := parameter(values, "wmpool"); value != "" { + config.HiveConfiguration["set:hivevar:wmpool"] = value + } + if value := firstNonEmpty(parameter(values, "applicationname"), parameter(values, "ApplicationName")); value != "" { + config.HiveConfiguration["set:hivevar:wmapp"] = value + } +} + +func hiveAssignmentValue(values map[string]string, key string) string { + for candidate, value := range values { + if strings.EqualFold(strings.TrimSpace(candidate), key) { + return value + } + } + return "" +} + +func applyDelegationToken(config *connectionConfig, values map[string]string) error { + if !strings.EqualFold(config.Auth, "DELEGATIONTOKEN") && !strings.EqualFold(config.Auth, "DELEGATION_TOKEN") { + return nil + } + token := firstNonEmpty(parameter(values, "delegationtoken"), parameter(values, "token"), config.Password) + if token == "" { + return errors.New("Hive delegation token authentication requires delegationToken, token, or password") + } + config.DelegationToken = token + identifier, password, err := decodeHadoopDelegationToken(token) + if err != nil { + return fmt.Errorf("decode Hive delegation token: %w", err) + } + config.Username = base64.StdEncoding.EncodeToString(identifier) + config.Password = base64.StdEncoding.EncodeToString(password) + return nil +} + +func decodeHadoopDelegationToken(value string) ([]byte, []byte, error) { + encoded := strings.Join(strings.Fields(strings.TrimSpace(value)), "") + if encoded == "" { + return nil, nil, errors.New("token is empty") + } + var decoded []byte + var decodeErr error + for _, encoding := range []*base64.Encoding{ + base64.RawURLEncoding, + base64.URLEncoding, + base64.RawStdEncoding, + base64.StdEncoding, + } { + decoded, decodeErr = encoding.DecodeString(encoded) + if decodeErr == nil { + break + } + } + if decodeErr != nil { + return nil, nil, decodeErr + } + reader := strings.NewReader(string(decoded)) + identifier, err := readHadoopByteArray(reader) + if err != nil { + return nil, nil, fmt.Errorf("identifier: %w", err) + } + password, err := readHadoopByteArray(reader) + if err != nil { + return nil, nil, fmt.Errorf("password: %w", err) + } + if len(identifier) == 0 || len(password) == 0 { + return nil, nil, errors.New("token identifier and password must be non-empty") + } + if _, err := readHadoopByteArray(reader); err != nil { + return nil, nil, fmt.Errorf("kind: %w", err) + } + if _, err := readHadoopByteArray(reader); err != nil { + return nil, nil, fmt.Errorf("service: %w", err) + } + if reader.Len() != 0 { + return nil, nil, errors.New("token contains trailing data") + } + return identifier, password, nil +} + +func readHadoopByteArray(reader io.ByteReader) ([]byte, error) { + length, err := readHadoopVInt(reader) + if err != nil { + return nil, err + } + if length < 0 { + return nil, fmt.Errorf("negative length %d", length) + } + if length > 64*1024*1024 { + return nil, fmt.Errorf("length %d exceeds limit", length) + } + value := make([]byte, int(length)) + byteReader, ok := reader.(io.Reader) + if !ok { + return nil, errors.New("reader cannot read token payload") + } + if _, err := io.ReadFull(byteReader, value); err != nil { + return nil, err + } + return value, nil +} + +func readHadoopVInt(reader io.ByteReader) (int64, error) { + firstByte, err := reader.ReadByte() + if err != nil { + return 0, err + } + first := int8(firstByte) + if first >= -112 { + return int64(first), nil + } + length := -111 - int(first) + negative := false + if first < -120 { + length = -119 - int(first) + negative = true + } + var value int64 + for index := 0; index < length-1; index++ { + current, readErr := reader.ReadByte() + if readErr != nil { + return 0, readErr + } + value = value<<8 | int64(current) + } + if negative { + value = ^value + } + return value, nil +} + +func applyKerberosJavaOptions(config *kerberosConfig, options []string) { + for _, option := range options { + trimmed := strings.TrimSpace(option) + switch { + case strings.HasPrefix(trimmed, "-Djava.security.krb5.conf="): + config.ConfigPath = javaSystemPropertyValue(strings.TrimPrefix(trimmed, "-Djava.security.krb5.conf=")) + case strings.HasPrefix(trimmed, "-Djava.security.auth.login.config="): + config.JAASConfigPath = javaSystemPropertyValue(strings.TrimPrefix(trimmed, "-Djava.security.auth.login.config=")) + } + } +} + +func applyZooKeeperKerberosJavaOptions(config *zooKeeperKerberosConfig, options []string) { + for _, option := range options { + trimmed := strings.TrimSpace(option) + keyValue := strings.TrimPrefix(trimmed, "-D") + key, value, found := strings.Cut(keyValue, "=") + if !strings.HasPrefix(trimmed, "-D") || !found { + continue + } + value = javaSystemPropertyValue(value) + switch strings.ToLower(strings.TrimSpace(key)) { + case "hive.zookeeper.use.kerberos": + config.Enabled = booleanValue(value) + case "zookeeper.sasl.client": + if !booleanValue(value) { + config.Enabled = false + } + case "zookeeper.sasl.client.username": + config.Service = firstNonEmpty(value, "zookeeper") + case "zookeeper.sasl.client.canonicalize.hostname": + config.CanonicalHostname = booleanValue(value) + case "zookeeper.server.principal": + config.ServerPrincipal = strings.TrimSpace(value) + case "zookeeper.server.realm": + config.Realm = strings.TrimSpace(value) + } + } +} + +func javaSystemPropertyValue(value string) string { + value = strings.TrimSpace(value) + if len(value) >= 2 && value[0] == '"' && value[len(value)-1] == '"' { + return value[1 : len(value)-1] + } + return value +} + +func applyKerberosEnvironment(config *kerberosConfig) { + config.ConfigPath = firstNonEmpty(config.ConfigPath, os.Getenv("KRB5_CONFIG")) + config.CCachePath = firstNonEmpty(config.CCachePath, os.Getenv("KRB5CCNAME")) + config.KeytabPath = firstNonEmpty(config.KeytabPath, os.Getenv("KRB5_CLIENT_KTNAME"), os.Getenv("KRB5_KTNAME")) +} + +func finalizeKerberosConfig(config *connectionConfig) error { + kerberos := &config.Kerberos + if !kerberos.Enabled { + return nil + } + kerberos.Password = config.Password + kerberos.ConfigPath = normalizeKerberosReference(kerberos.ConfigPath) + kerberos.CCachePath = normalizeKerberosReference(kerberos.CCachePath) + kerberos.KeytabPath = normalizeKerberosReference(kerberos.KeytabPath) + kerberos.JAASConfigPath = normalizeKerberosReference(kerberos.JAASConfigPath) + if kerberos.JAASConfigPath != "" { + if err := applyKerberosJAASFile(kerberos); err != nil { + return err + } + kerberos.KeytabPath = normalizeKerberosReference(kerberos.KeytabPath) + kerberos.CCachePath = normalizeKerberosReference(kerberos.CCachePath) + } + if kerberos.ConfigPath == "" { + if candidate := defaultKerberosConfigPath(); fileExists(candidate) { + kerberos.ConfigPath = candidate + } + } + if !kerberos.UseTicketCache && kerberos.CCachePath == "" { + if candidate := defaultKerberosCCachePath(); fileExists(candidate) { + kerberos.CCachePath = candidate + kerberos.UseTicketCache = true + } + } + if runtime.GOOS == "windows" && kerberos.ConfigPath == "" && kerberos.KeytabPath == "" && kerberos.CCachePath == "" { + kerberos.UseSSPI = true + } + if kerberos.UseSSPI { + return nil + } + if kerberos.ConfigPath == "" { + return errors.New("Kerberos requires krb5.conf or Windows SSPI") + } + if kerberos.ClientPrincipal == "" && !kerberos.UseTicketCache && !kerberos.UseKeytab { + kerberos.ClientPrincipal = strings.TrimSpace(config.Username) + } + if kerberos.KeytabPath != "" { + kerberos.UseKeytab = true + } + if kerberos.CCachePath != "" { + kerberos.UseTicketCache = true + } + kerberos.Realm = firstNonEmpty(kerberos.Realm, realmFromPrincipal(kerberos.ClientPrincipal)) + if !kerberos.UseTicketCache && !kerberos.UseKeytab && (kerberos.ClientPrincipal == "" || kerberos.Password == "") { + return errors.New("Kerberos requires SSPI, credential cache, keytab, or principal and password") + } + return nil +} + +var jaasOptionPattern = regexp.MustCompile(`(?i)\b(principal|keytab|ticketcache|usekeytab|useticketcache)\s*=\s*("(?:\\.|[^"])*"|'(?:\\.|[^'])*'|[^\s;]+)`) + +func applyKerberosJAASFile(config *kerberosConfig) error { + contents, err := os.ReadFile(config.JAASConfigPath) + if err != nil { + return fmt.Errorf("read Kerberos JAAS config: %w", err) + } + text := string(contents) + module := strings.Index(strings.ToLower(text), "krb5loginmodule") + if module < 0 { + return errors.New("Kerberos JAAS config contains no Krb5LoginModule") + } + block := text[module:] + if end := strings.IndexByte(block, ';'); end >= 0 { + block = block[:end] + } + for _, match := range jaasOptionPattern.FindAllStringSubmatch(block, -1) { + key := strings.ToLower(match[1]) + value := decodeJAASValue(match[2]) + switch key { + case "principal": + if config.ClientPrincipal == "" { + config.ClientPrincipal = value + } + case "keytab": + if config.KeytabPath == "" { + config.KeytabPath = value + } + case "ticketcache": + if config.CCachePath == "" { + config.CCachePath = value + } + case "usekeytab": + config.UseKeytab = config.UseKeytab || parseJAASBool(value) + case "useticketcache": + config.UseTicketCache = config.UseTicketCache || parseJAASBool(value) + } + } + return nil +} + +func decodeJAASValue(value string) string { + value = strings.TrimSpace(value) + if len(value) >= 2 && ((value[0] == '"' && value[len(value)-1] == '"') || (value[0] == '\'' && value[len(value)-1] == '\'')) { + value = value[1 : len(value)-1] + } + value = strings.ReplaceAll(value, `\\`, `\`) + value = strings.ReplaceAll(value, `\"`, `"`) + value = strings.ReplaceAll(value, `\'`, `'`) + return value +} + +func parseJAASBool(value string) bool { + switch strings.ToLower(strings.TrimSpace(value)) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + +func fileExists(path string) bool { + if strings.TrimSpace(path) == "" { + return false + } + info, err := os.Stat(path) + return err == nil && !info.IsDir() +} + +func normalizeKerberosReference(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "" + } + if strings.HasPrefix(strings.ToUpper(value), "FILE:") { + value = value[5:] + } + if strings.HasPrefix(value, "~/") { + if home, err := os.UserHomeDir(); err == nil { + value = filepath.Join(home, value[2:]) + } + } + return filepath.Clean(value) +} + +func buildTLSConfig(params connectParams, values map[string]string, serverName string) (*tls.Config, error) { + enabled := params.SSL || parameterBool(values, "ssl") || strings.EqualFold(parameter(values, "ssl"), "true") + if !enabled { + return nil, nil + } + config := &tls.Config{MinVersion: tls.VersionTLS12, ServerName: serverName} + if parameterBool(values, "sslinsecureskipverify") || parameterBool(values, "allowselfsigned") { + config.InsecureSkipVerify = true + } + var customRoots *x509.CertPool + credentialProviderPath := parameter(values, "storepasswordpath") + if path := strings.TrimSpace(params.CACertPath); path != "" { + contents, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("read Hive CA certificate: %w", err) + } + customRoots = x509.NewCertPool() + if !customRoots.AppendCertsFromPEM(contents) { + return nil, errors.New("Hive CA certificate contains no certificates") + } + } + trustStoreLocation := parameter(values, "ssltruststore") + if trustStoreLocation != "" { + if parameter(values, "truststorepassword") == "" && credentialProviderPath != "" { + return nil, errors.New("Hive storePasswordPath uses the Java Hadoop credential-provider format; configure trustStorePassword explicitly for the native agent") + } + certificates, err := loadTrustStore( + trustStoreLocation, + parameter(values, "truststorepassword"), + parameter(values, "truststoretype"), + ) + if err != nil { + return nil, fmt.Errorf("load Hive truststore: %w", err) + } + if customRoots == nil { + customRoots = x509.NewCertPool() + } + for _, certificate := range certificates { + customRoots.AddCert(certificate) + } + } + config.RootCAs = customRoots + if params.ClientCertPath != "" || params.ClientKeyPath != "" { + if params.ClientCertPath == "" || params.ClientKeyPath == "" { + return nil, errors.New("Hive client certificate and key must be configured together") + } + certificate, err := tls.LoadX509KeyPair(params.ClientCertPath, params.ClientKeyPath) + if err != nil { + return nil, fmt.Errorf("load Hive client certificate: %w", err) + } + config.Certificates = []tls.Certificate{certificate} + } + keyStoreLocation := parameter(values, "sslkeystore") + if keyStoreLocation != "" { + if parameter(values, "keystorepassword") == "" && credentialProviderPath != "" { + return nil, errors.New("Hive storePasswordPath uses the Java Hadoop credential-provider format; configure keyStorePassword explicitly for the native agent") + } + certificate, err := loadClientKeyStore( + keyStoreLocation, + parameter(values, "keystorepassword"), + parameter(values, "keystoretype"), + ) + if err != nil { + return nil, fmt.Errorf("load Hive keystore: %w", err) + } + config.Certificates = append(config.Certificates, certificate) + } + if parameterBool(values, "twoway") { + if keyStoreLocation == "" && len(config.Certificates) == 0 { + return nil, errors.New("Hive two-way TLS requires sslKeyStore or a client certificate") + } + if trustStoreLocation == "" && config.RootCAs == nil { + return nil, errors.New("Hive two-way TLS requires sslTrustStore or a CA certificate") + } + } + return config, nil +} + +func parameter(values map[string]string, key string) string { + for candidate, value := range values { + if strings.EqualFold(strings.TrimSpace(candidate), key) { + return strings.TrimSpace(value) + } + } + return "" +} + +func parameterBool(values map[string]string, key string) bool { + return booleanValue(parameter(values, key)) +} + +func booleanValue(value string) bool { + value = strings.ToLower(strings.TrimSpace(value)) + return value == "1" || value == "true" || value == "yes" || value == "on" +} + +func firstParameter(values map[string]string, keys ...string) (string, bool) { + for _, key := range keys { + for candidate, value := range values { + if strings.EqualFold(strings.TrimSpace(candidate), key) { + return strings.TrimSpace(value), true + } + } + } + return "", false +} + +func hasParameter(values map[string]string, key string) bool { + _, exists := firstParameter(values, key) + return exists +} + +func prefixedParameters(values map[string]string, prefix string) map[string]string { + result := map[string]string{} + for key, value := range values { + if len(key) <= len(prefix) || !strings.EqualFold(key[:len(prefix)], prefix) { + continue + } + name := strings.TrimSpace(key[len(prefix):]) + if name != "" { + result[name] = value + } + } + return result +} + +func readHiveInitFile(path string) ([]string, error) { + file, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("read Hive initFile: %w", err) + } + defer file.Close() + + var script strings.Builder + scanner := bufio.NewScanner(file) + scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" || strings.HasPrefix(line, "#") || strings.HasPrefix(line, "--") { + continue + } + script.WriteString(line) + script.WriteByte(' ') + } + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("read Hive initFile: %w", err) + } + + statements := make([]string, 0) + for _, statement := range strings.Split(script.String(), ";") { + if trimmed := strings.TrimSpace(statement); trimmed != "" { + statements = append(statements, trimmed) + } + } + return statements, nil +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if trimmed := strings.TrimSpace(value); trimmed != "" { + return trimmed + } + } + return "" +} + +func serviceFromPrincipal(principal string) string { + value := strings.TrimSpace(principal) + if separator := strings.IndexByte(value, '/'); separator > 0 { + return value[:separator] + } + return "" +} + +func hostFromPrincipal(principal string) string { + value := strings.TrimSpace(principal) + separator := strings.IndexByte(value, '/') + if separator < 0 { + return "" + } + value = value[separator+1:] + if realm := strings.IndexByte(value, '@'); realm >= 0 { + value = value[:realm] + } + if value == "_HOST" { + return "" + } + return value +} + +func realmFromPrincipal(principal string) string { + if separator := strings.LastIndexByte(principal, '@'); separator >= 0 { + return strings.TrimSpace(principal[separator+1:]) + } + return "" +} diff --git a/agents/drivers/hive-go/config_test.go b/agents/drivers/hive-go/config_test.go new file mode 100644 index 0000000000..09276b1f48 --- /dev/null +++ b/agents/drivers/hive-go/config_test.go @@ -0,0 +1,474 @@ +package main + +import ( + "bytes" + "crypto/x509" + "encoding/base64" + "net/url" + "os" + "path/filepath" + "reflect" + "runtime" + "strings" + "testing" + "time" + + pkcs12 "software.sslmate.com/src/go-pkcs12" +) + +func writeHadoopVInt(buffer *bytes.Buffer, value int64) { + if value >= -112 && value <= 127 { + buffer.WriteByte(byte(int8(value))) + return + } + lengthMarker := int8(-112) + encoded := value + if value < 0 { + encoded = ^value + lengthMarker = -120 + } + temporary := encoded + for temporary != 0 { + temporary >>= 8 + lengthMarker-- + } + buffer.WriteByte(byte(lengthMarker)) + length := -int(lengthMarker) + if lengthMarker < -120 { + length -= 120 + } else { + length -= 112 + } + for index := length; index != 0; index-- { + shift := uint((index - 1) * 8) + buffer.WriteByte(byte(encoded >> shift)) + } +} + +func encodeHadoopToken(identifier, password, kind, service []byte) string { + var buffer bytes.Buffer + for _, value := range [][]byte{identifier, password, kind, service} { + writeHadoopVInt(&buffer, int64(len(value))) + buffer.Write(value) + } + return base64.RawURLEncoding.EncodeToString(buffer.Bytes()) +} + +func TestParseDirectJDBCConnection(t *testing.T) { + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hive.example.com:10001/analytics;transportMode=http;httpPath=gateway;ssl=true", + Username: "alice", + Password: "secret", + }) + if err != nil { + t.Fatal(err) + } + if len(config.Endpoints) != 1 || config.Endpoints[0] != (endpoint{Host: "hive.example.com", Port: 10001}) { + t.Fatalf("unexpected endpoints: %#v", config.Endpoints) + } + if config.Database != "analytics" || config.TransportMode != "http" || config.HTTPPath != "gateway" { + t.Fatalf("unexpected config: %#v", config) + } + if config.TLSConfig == nil { + t.Fatal("expected TLS config") + } +} + +func TestParseHTTPKerberosChannelBinding(t *testing.T) { + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hive.example.com:10001/default;transportMode=http;ssl=true;auth=KERBEROS;principal=HTTP/_HOST@EXAMPLE.COM;kerberosChannelBinding=true", + Username: "alice@EXAMPLE.COM", + Password: "secret", + AgentJavaOptions: []string{"-Djava.security.krb5.conf=/etc/krb5.conf"}, + }) + if err != nil { + t.Fatal(err) + } + if !config.Kerberos.ChannelBinding { + t.Fatal("expected HTTP Kerberos TLS channel binding") + } +} + +func TestParseZooKeeperKerberosJDBCConnection(t *testing.T) { + configPath := "/etc/krb5.conf" + if runtime.GOOS == "windows" { + configPath = `C:\ProgramData\MIT\Kerberos5\krb5.ini` + } + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://zk1.example.com:2181,zk2.example.com:2181/default;serviceDiscoveryMode=zooKeeper;zooKeeperNamespace=kyuubi;principal=hive/_HOST@EXAMPLE.COM;hive.server2.thrift.sasl.qop=auth-conf", + Username: "alice@EXAMPLE.COM", + Password: "secret", + AgentJavaOptions: []string{"-Djava.security.krb5.conf=" + configPath}, + }) + if err != nil { + t.Fatal(err) + } + if config.ServiceDiscoveryMode != "zookeeper" || config.ZooKeeperNamespace != "kyuubi" { + t.Fatalf("unexpected ZooKeeper config: %#v", config) + } + if len(config.Endpoints) != 2 || config.Endpoints[1].Host != "zk2.example.com" || config.Endpoints[1].Port != 2181 { + t.Fatalf("unexpected endpoints: %#v", config.Endpoints) + } + if !config.Kerberos.Enabled || config.Kerberos.Service != "hive" || config.Kerberos.QOP != "auth-conf" { + t.Fatalf("unexpected Kerberos config: %#v", config.Kerberos) + } + if config.Kerberos.ConfigPath != configPath { + t.Fatalf("unexpected krb5 config path: %q", config.Kerberos.ConfigPath) + } + if !config.ZooKeeperKerberos.Enabled || config.ZooKeeperKerberos.Service != "zookeeper" || !config.ZooKeeperKerberos.CanonicalHostname { + t.Fatalf("unexpected ZooKeeper Kerberos config: %#v", config.ZooKeeperKerberos) + } +} + +func TestParseZooKeeperKerberosCompatibilityProperties(t *testing.T) { + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://zk.example.com:2181/default;serviceDiscoveryMode=zooKeeper;principal=hive/_HOST@HIVE.EXAMPLE.COM;zookeeper.sasl.client.username=zkservice;zookeeper.server.realm=ZK.EXAMPLE.COM;zookeeper.sasl.client.canonicalize.hostname=false", + Username: "alice@EXAMPLE.COM", + Password: "secret", + AgentJavaOptions: []string{ + "-Djava.security.krb5.conf=/etc/krb5.conf", + "-Dzookeeper.server.principal=zookeeper/zk.example.com@EXPLICIT.EXAMPLE.COM", + }, + }) + if err != nil { + t.Fatal(err) + } + want := zooKeeperKerberosConfig{ + Enabled: true, + Service: "zkservice", + ServerPrincipal: "zookeeper/zk.example.com@EXPLICIT.EXAMPLE.COM", + Realm: "ZK.EXAMPLE.COM", + CanonicalHostname: false, + } + if config.ZooKeeperKerberos != want { + t.Fatalf("ZooKeeper Kerberos config = %#v, want %#v", config.ZooKeeperKerberos, want) + } +} + +func TestZooKeeperKerberosCanBeDisabledExplicitly(t *testing.T) { + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://zk.example.com:2181/default;serviceDiscoveryMode=zooKeeper;principal=hive/_HOST@EXAMPLE.COM;hive.zookeeper.use.kerberos=true", + Username: "alice@EXAMPLE.COM", + Password: "secret", + AgentJavaOptions: []string{ + "-Djava.security.krb5.conf=/etc/krb5.conf", + "-Dzookeeper.sasl.client=false", + }, + }) + if err != nil { + t.Fatal(err) + } + if config.ZooKeeperKerberos.Enabled { + t.Fatalf("ZooKeeper Kerberos should be disabled: %#v", config.ZooKeeperKerberos) + } +} + +func TestURLParamsOverrideConnectionString(t *testing.T) { + config, err := parseConnectionConfig(connectParams{ + Host: "hive.example.com", + Port: 10000, + ConnectionString: "jdbc:hive2://old.example.com:10000/default;transportMode=binary", + URLParams: "transportMode=http;httpPath=proxy", + }) + if err != nil { + t.Fatal(err) + } + if config.TransportMode != "http" || config.HTTPPath != "proxy" { + t.Fatalf("URL params did not override connection string: %#v", config) + } +} + +func TestParseStandardJDBCURLSectionsAndCredentials(t *testing.T) { + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hs2.example.com:10001/analytics;user=alice;password=p%40ss?hive.server2.transport.mode=http;hive.server2.thrift.http.path=proxy;hive.exec.dynamic.partition=true#SourceTable=events", + URLParams: "user=bob?hive.exec.dynamic.partition=false#SourceTable=override", + }) + if err != nil { + t.Fatal(err) + } + if config.Username != "bob" || config.Password != "p@ss" { + t.Fatalf("unexpected credentials: %q / %q", config.Username, config.Password) + } + if config.TransportMode != "http" || config.HTTPPath != "proxy" { + t.Fatalf("deprecated Hive conf transport settings were not applied: %#v", config) + } + want := map[string]string{ + "set:hiveconf:hive.exec.dynamic.partition": "false", + "set:hivevar:SourceTable": "override", + } + if !reflect.DeepEqual(config.HiveConfiguration, want) { + t.Fatalf("unexpected OpenSession configuration: %#v", config.HiveConfiguration) + } +} + +func TestOpenSessionCompatibilityVariablesFromSessionParams(t *testing.T) { + config, err := parseConnectionConfig(connectParams{ + Host: "hs2.example.com", + URLParams: "proxyUser=alice;hiveCreateAsExternalLegacy=TRUE;wmPool=etl;hiveconf:hive.exec.compress.output=true;hivevar:source=events", + }) + if err != nil { + t.Fatal(err) + } + want := map[string]string{ + "hive.server2.proxy.user": "alice", + "set:hiveconf:hive.create.as.external.legacy": "true", + "set:hiveconf:hive.exec.compress.output": "true", + "set:hivevar:source": "events", + "set:hivevar:wmpool": "etl", + } + if !reflect.DeepEqual(config.HiveConfiguration, want) { + t.Fatalf("unexpected OpenSession compatibility variables: %#v", config.HiveConfiguration) + } +} + +func TestParseHiveJDBCClientCompatibilityOptions(t *testing.T) { + directory := t.TempDir() + initPath := filepath.Join(directory, "hive-init.sql") + if err := os.WriteFile(initPath, []byte("# ignored\nSET hive.exec.dynamic.partition=true;\n-- ignored\nUSE analytics;\n"), 0o600); err != nil { + t.Fatal(err) + } + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hs2.example.com:10001/default;transportMode=http;auth=jwt;jwt=signed-token;fetchSize=77;socketTimeout=9;thrift.client.max.message.size=1048576;retries=3;retryInterval=250;requestTrack=true;cookieAuth=false;cookieName=CustomAuth;http.header.X-Trace-ID=trace-value;http.cookie.SessionID=cookie-value;applicationName=dbx-hive;initFile=" + url.QueryEscape(initPath), + }) + if err != nil { + t.Fatal(err) + } + if config.Auth != "JWT" || config.JWT != "signed-token" { + t.Fatalf("unexpected JWT config: %#v", config) + } + if config.FetchSize != 77 || config.SocketTimeout != 9*time.Second || config.MaxMessageSize != 1048576 { + t.Fatalf("unexpected client sizing config: %#v", config) + } + if config.Retries != 3 || config.RetryInterval != 250*time.Millisecond { + t.Fatalf("unexpected retry config: %#v", config) + } + if config.CookieAuth || config.CookieName != "CustomAuth" { + t.Fatalf("unexpected cookie auth config: %#v", config) + } + if config.HTTPHeaders["X-Trace-ID"] != "trace-value" || config.HTTPCookies["SessionID"] != "cookie-value" { + t.Fatalf("HTTP header or cookie case was not preserved: %#v / %#v", config.HTTPHeaders, config.HTTPCookies) + } + if !config.RequestTracking { + t.Fatal("requestTrack was not enabled") + } + if config.HiveConfiguration["set:hivevar:wmapp"] != "dbx-hive" { + t.Fatalf("applicationName was not mapped: %#v", config.HiveConfiguration) + } + if !reflect.DeepEqual(config.InitStatements, []string{"SET hive.exec.dynamic.partition=true", "USE analytics"}) { + t.Fatalf("unexpected init statements: %#v", config.InitStatements) + } +} + +func TestJWTCanComeFromEnvironment(t *testing.T) { + t.Setenv("JWT", "environment-token") + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hs2.example.com:10001/default;transportMode=http;auth=jwt", + }) + if err != nil { + t.Fatal(err) + } + if config.JWT != "environment-token" { + t.Fatalf("unexpected JWT: %q", config.JWT) + } +} + +func TestBrowserAuthSupportsInteractiveSSOParameters(t *testing.T) { + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hs2.example.com:10001/default;transportMode=http;auth=browser;browserResponsePort=18080;browserResponseTimeout=45;browserDisableSslCheck=true", + }) + if err != nil { + t.Fatal(err) + } + if config.BrowserToken != "" || config.BrowserResponsePort != 18080 || config.BrowserResponseTimeout != 45*time.Second || !config.BrowserDisableSSLCheck { + t.Fatalf("unexpected browser SSO config: %#v", config) + } +} + +func TestInvalidHiveJDBCClientSizesAreRejected(t *testing.T) { + for _, params := range []string{ + "fetchSize=zero", + "socketTimeout=soon", + "thrift.client.max.message.size=huge", + } { + _, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hs2.example.com/default;" + params, + }) + if err == nil { + t.Fatalf("expected %q to be rejected", params) + } + } +} + +func TestJavaCredentialProviderPasswordPathIsNotSilentlyIgnored(t *testing.T) { + _, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hs2.example.com/default;ssl=true;sslTrustStore=/tmp/truststore.jks;storePasswordPath=jceks://file/tmp/hive.jceks", + }) + if err == nil || !strings.Contains(err.Error(), "configure trustStorePassword explicitly") { + t.Fatalf("unexpected credential-provider error: %v", err) + } +} + +func TestKerberosAcceptsSaslQopCompatibilityAlias(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "krb5.conf") + if err := os.WriteFile(configPath, []byte("[libdefaults]\n default_realm = EXAMPLE.COM\n"), 0o600); err != nil { + t.Fatal(err) + } + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hs2.example.com/default;auth=kerberos;principal=hive/_HOST@EXAMPLE.COM;kerberosPrincipal=alice@EXAMPLE.COM;sasl.qop=auth-int", + Password: "secret", + AgentJavaOptions: []string{"-Djava.security.krb5.conf=" + configPath}, + }) + if err != nil { + t.Fatal(err) + } + if config.Kerberos.QOP != "auth-int" { + t.Fatalf("unexpected SASL QOP: %q", config.Kerberos.QOP) + } +} + +func TestDelegationTokenProducesDigestCredentials(t *testing.T) { + identifier := []byte("token-identifier") + password := []byte("token-password") + token := encodeHadoopToken(identifier, password, []byte("HIVE_DELEGATION_TOKEN"), []byte("hs2.example.com:10000")) + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hs2.example.com/default;auth=delegationToken;delegationToken=" + token, + }) + if err != nil { + t.Fatal(err) + } + if config.Auth != "DELEGATIONTOKEN" { + t.Fatalf("unexpected auth: %q", config.Auth) + } + if config.DelegationToken != token { + t.Fatalf("raw delegation token was not retained: %q", config.DelegationToken) + } + if config.Username != base64.StdEncoding.EncodeToString(identifier) || config.Password != base64.StdEncoding.EncodeToString(password) { + t.Fatalf("unexpected delegation token credentials: %q / %q", config.Username, config.Password) + } +} + +func TestDelegationTokenRejectsMalformedValue(t *testing.T) { + _, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hs2.example.com/default;auth=delegationToken;delegationToken=not-a-token", + }) + if err == nil || !strings.Contains(err.Error(), "decode Hive delegation token") { + t.Fatalf("expected delegation token error, got %v", err) + } +} + +func TestParseIPv6Endpoint(t *testing.T) { + value, err := parseEndpoint("[2001:db8::1]:10000", defaultHivePort) + if err != nil { + t.Fatal(err) + } + if value.Host != "2001:db8::1" || value.Port != 10000 { + t.Fatalf("unexpected endpoint: %#v", value) + } +} + +func TestKerberosSeparatesServerAndClientPrincipals(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "krb5.conf") + if err := os.WriteFile(configPath, []byte("[libdefaults]\n default_realm = CLIENT.EXAMPLE.COM\n"), 0o600); err != nil { + t.Fatal(err) + } + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://alias.example.com:10000/default;auth=kerberos;principal=hive/_HOST@SERVER.EXAMPLE.COM;kerberosPrincipal=alice@CLIENT.EXAMPLE.COM;kerberosEnableCanonicalHostnameCheck=false", + Password: "secret", + AgentJavaOptions: []string{`-Djava.security.krb5.conf="` + configPath + `"`}, + }) + if err != nil { + t.Fatal(err) + } + if config.Kerberos.ServerPrincipal != "hive/_HOST@SERVER.EXAMPLE.COM" { + t.Fatalf("unexpected server principal: %q", config.Kerberos.ServerPrincipal) + } + if config.Kerberos.ClientPrincipal != "alice@CLIENT.EXAMPLE.COM" { + t.Fatalf("unexpected client principal: %q", config.Kerberos.ClientPrincipal) + } + if config.Kerberos.Realm != "CLIENT.EXAMPLE.COM" || config.Kerberos.CanonicalHostname { + t.Fatalf("unexpected Kerberos options: %#v", config.Kerberos) + } +} + +func TestKerberosReadsLegacyJAASKeytabOptions(t *testing.T) { + directory := t.TempDir() + configPath := filepath.Join(directory, "krb5.conf") + jaasPath := filepath.Join(directory, "hive jaas.conf") + keytabPath := filepath.Join(directory, "alice.keytab") + if err := os.WriteFile(configPath, []byte("[libdefaults]\n default_realm = EXAMPLE.COM\n"), 0o600); err != nil { + t.Fatal(err) + } + jaas := `HiveClient { + com.sun.security.auth.module.Krb5LoginModule required + useKeyTab=true + keyTab="` + keytabPath + `" + principal="alice@EXAMPLE.COM" + doNotPrompt=true; +};` + if err := os.WriteFile(jaasPath, []byte(jaas), 0o600); err != nil { + t.Fatal(err) + } + config, err := parseConnectionConfig(connectParams{ + ConnectionString: "jdbc:hive2://hs2.example.com/default;auth=kerberos;principal=hive/_HOST@EXAMPLE.COM", + AgentJavaOptions: []string{ + "-Djava.security.krb5.conf=" + configPath, + `-Djava.security.auth.login.config="` + jaasPath + `"`, + }, + }) + if err != nil { + t.Fatal(err) + } + if !config.Kerberos.UseKeytab || config.Kerberos.KeytabPath != keytabPath { + t.Fatalf("JAAS keytab was not preserved: %#v", config.Kerberos) + } + if config.Kerberos.ClientPrincipal != "alice@EXAMPLE.COM" { + t.Fatalf("JAAS principal was not preserved: %#v", config.Kerberos) + } +} + +func TestBuildHiveTLSConfigFromPKCS12Stores(t *testing.T) { + privateKey, certificate, _ := testZooKeeperCertificate(t) + password := "changeit" + keyStore, err := pkcs12.Modern.Encode(privateKey, certificate, nil, password) + if err != nil { + t.Fatal(err) + } + trustStore, err := pkcs12.Modern.EncodeTrustStore([]*x509.Certificate{certificate}, password) + if err != nil { + t.Fatal(err) + } + directory := t.TempDir() + keyPath := filepath.Join(directory, "client.p12") + trustPath := filepath.Join(directory, "trust.p12") + if err := os.WriteFile(keyPath, keyStore, 0o600); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(trustPath, trustStore, 0o600); err != nil { + t.Fatal(err) + } + config, err := buildTLSConfig(connectParams{}, map[string]string{ + "ssl": "true", + "twoway": "true", + "ssltruststore": trustPath, + "truststorepassword": password, + "truststoretype": "PKCS12", + "sslkeystore": keyPath, + "keystorepassword": password, + "keystoretype": "PKCS12", + }, "hs2.example.com") + if err != nil { + t.Fatal(err) + } + if config == nil || config.RootCAs == nil || len(config.Certificates) != 1 { + t.Fatalf("unexpected Hive TLS config: %#v", config) + } +} + +func TestBuildHiveTwoWayTLSRequiresTrustAndKeyMaterial(t *testing.T) { + if _, err := buildTLSConfig(connectParams{}, map[string]string{ + "ssl": "true", + "twoway": "true", + }, "hs2.example.com"); err == nil { + t.Fatal("expected missing two-way TLS material error") + } +} diff --git a/agents/drivers/hive-go/connector.go b/agents/drivers/hive-go/connector.go new file mode 100644 index 0000000000..865a452fb7 --- /dev/null +++ b/agents/drivers/hive-go/connector.go @@ -0,0 +1,194 @@ +package main + +import ( + "context" + "crypto/tls" + "database/sql" + "database/sql/driver" + "errors" + "fmt" + "strings" + "time" + + "github.com/beltran/gosasl" + gohive "github.com/t8y2/dbx/agents/go-common/gohive" +) + +type connectorFactory func(endpoint) driver.Connector + +type discoveryConnector struct { + discovery endpointDiscovery + factory connectorFactory + driver driver.Driver + retries int + retryInterval time.Duration +} + +func newDiscoveryConnector(config connectionConfig) *discoveryConnector { + return &discoveryConnector{ + discovery: newEndpointDiscovery(config), + factory: func(target endpoint) driver.Connector { + tlsConfig := config.TLSConfig + if target.SSL && tlsConfig == nil && !config.TLSExplicit { + tlsConfig = &tls.Config{MinVersion: tls.VersionTLS12, ServerName: target.Host} + } + if tlsConfig != nil { + tlsConfig = tlsConfig.Clone() + if tlsConfig.ServerName == "" || tlsConfig.ServerName == config.Endpoints[0].Host { + tlsConfig.ServerName = target.Host + } + } + hiveConfiguration := make(map[string]string, len(config.HiveConfiguration)+1) + for key, value := range config.HiveConfiguration { + hiveConfiguration[key] = value + } + if config.Kerberos.Enabled { + hiveConfiguration["hive.server2.thrift.sasl.qop"] = config.Kerberos.QOP + } + transportMode := config.TransportMode + if target.TransportMode != "" && !config.TransportModeExplicit { + transportMode = target.TransportMode + } + httpPath := config.HTTPPath + if target.HTTPPath != "" && !config.HTTPPathExplicit { + httpPath = target.HTTPPath + } + auth := config.Auth + if target.Auth != "" && !config.AuthExplicit { + auth = target.Auth + } + service := kerberosServiceForEndpoint(config, target) + return gohive.NewConnector(gohive.Config{ + Host: target.Host, + Port: target.Port, + Auth: normalizeHiveAuth(auth), + Username: config.Username, + Password: config.Password, + Database: config.Database, + TransportMode: transportMode, + HTTPPath: httpPath, + Service: service, + HTTPKerberosChannelBinding: config.Kerberos.ChannelBinding, + GSSAPIOptions: gssapiOptionsFromKerberos(config.Kerberos), + TLSConfig: tlsConfig, + HiveConfiguration: hiveConfiguration, + ConnectTimeout: config.ConnectTimeout, + SocketTimeout: config.SocketTimeout, + HTTPTimeout: config.SocketTimeout, + FetchSize: int64(config.FetchSize), + MaxMessageSize: config.MaxMessageSize, + HTTPHeaders: config.HTTPHeaders, + HTTPCookies: config.HTTPCookies, + RequestTracking: config.RequestTracking, + DisableCookieAuth: !config.CookieAuth, + CookieName: config.CookieName, + JWT: config.JWT, + DelegationToken: config.DelegationToken, + BrowserToken: config.BrowserToken, + BrowserClientID: config.BrowserClientID, + BrowserResponsePort: config.BrowserResponsePort, + BrowserResponseTimeout: config.BrowserResponseTimeout, + BrowserDisableSSLCheck: config.BrowserDisableSSLCheck, + }) + }, + driver: &gohive.Driver{}, + retries: max(config.Retries, 1), + retryInterval: config.RetryInterval, + } +} + +func gssapiOptionsFromKerberos(config kerberosConfig) gosasl.GSSAPIOptions { + return gosasl.GSSAPIOptions{ + ConfigPath: config.ConfigPath, + CCachePath: config.CCachePath, + KeytabPath: config.KeytabPath, + Principal: config.ClientPrincipal, + Password: config.Password, + QOP: config.QOP, + AuthorizationID: config.AuthorizationID, + ServiceHost: config.ServerName, + UseCCache: config.UseTicketCache, + UseKeytab: config.UseKeytab, + UseSSPI: config.UseSSPI, + CanonicalizeHost: config.CanonicalHostname, + DisablePAFXFAST: config.DisablePAFXFAST, + } +} + +func kerberosServiceForEndpoint(config connectionConfig, target endpoint) string { + service := firstNonEmpty(config.Kerberos.ServerPrincipal, config.Kerberos.Service) + if target.Principal != "" && !config.Kerberos.ServerPrincipalExplicit { + service = target.Principal + } + return service +} + +func (connector *discoveryConnector) Connect(ctx context.Context) (driver.Conn, error) { + var failures []string + for attempt := 0; attempt < max(connector.retries, 1); attempt++ { + rejected := map[string]bool{} + for { + endpoints, err := connector.discovery.Endpoints(ctx, rejected) + if err != nil { + failures = append(failures, fmt.Sprintf("discovery attempt %d: %v", attempt+1, err)) + break + } + if len(endpoints) == 0 { + break + } + for _, target := range endpoints { + connection, connectErr := connector.factory(target).Connect(ctx) + if connectErr == nil { + return connection, nil + } + rejected[target.address()] = true + failures = append(failures, fmt.Sprintf("attempt %d %s: %v", attempt+1, target.address(), connectErr)) + } + break + } + if attempt+1 < connector.retries && connector.retryInterval > 0 { + timer := time.NewTimer(connector.retryInterval) + select { + case <-ctx.Done(): + timer.Stop() + return nil, ctx.Err() + case <-timer.C: + } + } + } + if len(failures) == 0 { + return nil, errors.New("Hive discovery returned no endpoints") + } + return nil, fmt.Errorf("all HiveServer2 endpoints failed: %s", strings.Join(failures, "; ")) +} + +func (connector *discoveryConnector) Driver() driver.Driver { + return connector.driver +} + +func openHiveDatabase(config connectionConfig) *sql.DB { + database := sql.OpenDB(newDiscoveryConnector(config)) + database.SetMaxOpenConns(1) + database.SetMaxIdleConns(1) + return database +} + +func normalizeHiveAuth(value string) string { + normalized := strings.ToUpper(strings.TrimSpace(value)) + switch normalized { + case "": + return "NONE" + case "NOSASL", "NO_SASL": + return "NOSASL" + case "KERBEROS", "GSSAPI": + return "KERBEROS" + case "LDAP": + return "LDAP" + case "CUSTOM": + return "CUSTOM" + case "DIGEST-MD5", "DELEGATIONTOKEN", "DELEGATION_TOKEN": + return "DIGEST-MD5" + default: + return normalized + } +} diff --git a/agents/drivers/hive-go/connector_test.go b/agents/drivers/hive-go/connector_test.go new file mode 100644 index 0000000000..89a2d35310 --- /dev/null +++ b/agents/drivers/hive-go/connector_test.go @@ -0,0 +1,155 @@ +package main + +import ( + "context" + "database/sql/driver" + "errors" + "io" + "reflect" + "testing" +) + +type staticDiscovery struct { + values []endpoint +} + +func (discovery staticDiscovery) Endpoints(_ context.Context, rejected map[string]bool) ([]endpoint, error) { + result := make([]endpoint, 0, len(discovery.values)) + for _, value := range discovery.values { + if !rejected[value.address()] { + result = append(result, value) + } + } + return result, nil +} + +type fakeConnector struct { + connection driver.Conn + err error +} + +func (connector fakeConnector) Connect(context.Context) (driver.Conn, error) { + return connector.connection, connector.err +} + +func (fakeConnector) Driver() driver.Driver { return fakeDriver{} } + +type fakeDriver struct{} + +func (fakeDriver) Open(string) (driver.Conn, error) { return &fakeConnection{}, nil } + +type fakeConnection struct{} + +func (*fakeConnection) Prepare(string) (driver.Stmt, error) { return nil, errors.New("unsupported") } +func (*fakeConnection) Close() error { return nil } +func (*fakeConnection) Begin() (driver.Tx, error) { return nil, errors.New("unsupported") } + +type emptyRows struct{} + +func (emptyRows) Columns() []string { return nil } +func (emptyRows) Close() error { return nil } +func (emptyRows) Next([]driver.Value) error { return io.EOF } + +func TestDiscoveryConnectorFailsOver(t *testing.T) { + first := endpoint{Host: "first", Port: 10000} + second := endpoint{Host: "second", Port: 10000} + connected := &fakeConnection{} + var attempts []endpoint + connector := &discoveryConnector{ + discovery: staticDiscovery{values: []endpoint{first, second}}, + driver: fakeDriver{}, + factory: func(value endpoint) driver.Connector { + attempts = append(attempts, value) + if value == first { + return fakeConnector{err: errors.New("unavailable")} + } + return fakeConnector{connection: connected} + }, + } + value, err := connector.Connect(context.Background()) + if err != nil { + t.Fatal(err) + } + if value != connected || !reflect.DeepEqual(attempts, []endpoint{first, second}) { + t.Fatalf("unexpected failover result: value=%#v attempts=%#v", value, attempts) + } +} + +func TestDiscoveryConnectorRetriesAllEndpoints(t *testing.T) { + target := endpoint{Host: "hs2", Port: 10000} + connected := &fakeConnection{} + attempts := 0 + connector := &discoveryConnector{ + discovery: staticDiscovery{values: []endpoint{target}}, + driver: fakeDriver{}, + retries: 2, + factory: func(endpoint) driver.Connector { + attempts++ + if attempts == 1 { + return fakeConnector{err: errors.New("temporarily unavailable")} + } + return fakeConnector{connection: connected} + }, + } + value, err := connector.Connect(context.Background()) + if err != nil { + t.Fatal(err) + } + if value != connected || attempts != 2 { + t.Fatalf("unexpected retry result: value=%#v attempts=%d", value, attempts) + } +} + +func TestNormalizeHiveAuth(t *testing.T) { + for input, expected := range map[string]string{ + "": "NONE", + "noSasl": "NOSASL", + "kerberos": "KERBEROS", + "delegationToken": "DIGEST-MD5", + "vendor-auth-mode": "VENDOR-AUTH-MODE", + } { + if actual := normalizeHiveAuth(input); actual != expected { + t.Fatalf("normalizeHiveAuth(%q) = %q, expected %q", input, actual, expected) + } + } +} + +func TestKerberosServiceForEndpoint(t *testing.T) { + target := endpoint{Host: "hs2.example.com", Port: 10000, Principal: "hive/hs2.example.com@EXAMPLE.COM"} + discovered := connectionConfig{Kerberos: kerberosConfig{Service: "hive"}} + if value := kerberosServiceForEndpoint(discovered, target); value != target.Principal { + t.Fatalf("discovered principal was ignored: %s", value) + } + explicit := connectionConfig{Kerberos: kerberosConfig{ + Service: "hive", + ServerPrincipal: "hive/_HOST@USER.EXAMPLE.COM", + ServerPrincipalExplicit: true, + }} + if value := kerberosServiceForEndpoint(explicit, target); value != explicit.Kerberos.ServerPrincipal { + t.Fatalf("explicit principal was overwritten: %s", value) + } +} + +func TestKerberosUsesConnectionScopedGSSAPIOptions(t *testing.T) { + config := kerberosConfig{ + Enabled: true, + ServerPrincipal: "hive/_HOST@EXAMPLE.COM", + ClientPrincipal: "alice@EXAMPLE.COM", + ServerName: "canonical.example.com", + CanonicalHostname: true, + ConfigPath: "/etc/krb5.conf", + CCachePath: "/tmp/alice.ccache", + KeytabPath: "/tmp/alice.keytab", + Password: "secret", + AuthorizationID: "proxy-user", + QOP: "auth-conf", + UseTicketCache: true, + UseKeytab: true, + UseSSPI: true, + DisablePAFXFAST: true, + } + options := gssapiOptionsFromKerberos(config) + if options.ConfigPath != "/etc/krb5.conf" || options.Principal != "alice@EXAMPLE.COM" || options.Password != "secret" || options.AuthorizationID != "proxy-user" || options.ServiceHost != "canonical.example.com" || !options.CanonicalizeHost || options.CCachePath != "/tmp/alice.ccache" || options.KeytabPath != "/tmp/alice.keytab" || options.QOP != "auth-conf" || !options.UseCCache || !options.UseKeytab || !options.UseSSPI || !options.DisablePAFXFAST { + t.Fatalf("unexpected GSSAPI options: %#v", options) + } +} diff --git a/agents/drivers/hive-go/discovery.go b/agents/drivers/hive-go/discovery.go new file mode 100644 index 0000000000..94e01f156c --- /dev/null +++ b/agents/drivers/hive-go/discovery.go @@ -0,0 +1,392 @@ +package main + +import ( + "context" + "crypto/tls" + "encoding/json" + "errors" + "fmt" + "math/rand/v2" + "net" + "net/url" + "strconv" + "strings" + "time" + + "github.com/go-zookeeper/zk" +) + +type endpointDiscovery interface { + Endpoints(context.Context, map[string]bool) ([]endpoint, error) +} + +type directDiscovery struct { + endpoints []endpoint +} + +func (discovery directDiscovery) Endpoints(_ context.Context, rejected map[string]bool) ([]endpoint, error) { + return shuffledEndpoints(discovery.endpoints, rejected), nil +} + +type zooKeeperDiscovery struct { + servers []endpoint + namespace string + discoveryMode string + authScheme string + auth string + timeout time.Duration + dialer func([]string, time.Duration) (zooKeeperClient, <-chan zk.Event, error) +} + +type zooKeeperClient interface { + AddAuth(string, []byte) error + Children(string) ([]string, *zk.Stat, error) + Get(string) ([]byte, *zk.Stat, error) + Close() +} + +func newEndpointDiscovery(config connectionConfig) endpointDiscovery { + if strings.EqualFold(config.ServiceDiscoveryMode, "zookeeper") || strings.EqualFold(config.ServiceDiscoveryMode, "zookeeperha") { + return &zooKeeperDiscovery{ + servers: append([]endpoint(nil), config.Endpoints...), + namespace: config.ZooKeeperNamespace, + discoveryMode: config.ServiceDiscoveryMode, + authScheme: config.ZooKeeperAuthScheme, + auth: config.ZooKeeperAuth, + timeout: config.ConnectTimeout, + dialer: newZooKeeperDialer(config), + } + } + return directDiscovery{endpoints: append([]endpoint(nil), config.Endpoints...)} +} + +func newZooKeeperDialer(config connectionConfig) func([]string, time.Duration) (zooKeeperClient, <-chan zk.Event, error) { + if config.ZooKeeperKerberos.Enabled { + return func(servers []string, timeout time.Duration) (zooKeeperClient, <-chan zk.Event, error) { + return connectKerberosZooKeeper(servers, timeout, config.ZooKeeperTLSConfig, config) + } + } + tlsConfig := config.ZooKeeperTLSConfig + if tlsConfig == nil { + return func(servers []string, timeout time.Duration) (zooKeeperClient, <-chan zk.Event, error) { + return zk.Connect(servers, timeout, zk.WithLogInfo(false)) + } + } + return func(servers []string, timeout time.Duration) (zooKeeperClient, <-chan zk.Event, error) { + return zk.Connect( + servers, + timeout, + zk.WithLogInfo(false), + zk.WithDialer(func(network, address string, dialTimeout time.Duration) (net.Conn, error) { + config := tlsConfig.Clone() + if config.ServerName == "" { + host, _, err := net.SplitHostPort(address) + if err == nil { + config.ServerName = host + } + } + dialer := &net.Dialer{Timeout: dialTimeout} + return tls.DialWithDialer(dialer, network, address, config) + }), + ) + } +} + +func (discovery *zooKeeperDiscovery) Endpoints(ctx context.Context, rejected map[string]bool) ([]endpoint, error) { + addresses := make([]string, 0, len(discovery.servers)) + for _, server := range discovery.servers { + addresses = append(addresses, server.address()) + } + timeout := discovery.timeout + if timeout <= 0 { + timeout = defaultConnectTimeout + } + connection, events, err := discovery.dialer(addresses, timeout) + if err != nil { + return nil, fmt.Errorf("connect to ZooKeeper: %w", err) + } + defer connection.Close() + if err := waitForZooKeeperSession(ctx, events, timeout); err != nil { + return nil, err + } + if discovery.authScheme != "" || discovery.auth != "" { + if discovery.authScheme == "" || discovery.auth == "" { + return nil, errors.New("ZooKeeper auth scheme and credentials must be configured together") + } + if err := connection.AddAuth(discovery.authScheme, []byte(discovery.auth)); err != nil { + return nil, fmt.Errorf("authenticate to ZooKeeper: %w", err) + } + } + resolved := make([]endpoint, 0) + var listedPath string + var nodeFailures []string + for _, path := range discovery.paths() { + children, _, childrenErr := connection.Children(path) + if errors.Is(childrenErr, zk.ErrNoNode) { + continue + } + if childrenErr != nil { + return nil, fmt.Errorf("list ZooKeeper namespace %s: %w", path, childrenErr) + } + listedPath = path + for _, child := range children { + data, _, dataErr := connection.Get(path + "/" + child) + if dataErr != nil { + if errors.Is(dataErr, zk.ErrNoNode) { + continue + } + nodeFailures = append(nodeFailures, fmt.Sprintf("%s/%s: %v", path, child, dataErr)) + continue + } + value, parseErr := parseHiveServerRegistration(child, data) + if parseErr == nil { + resolved = append(resolved, value) + } else { + nodeFailures = append(nodeFailures, fmt.Sprintf("%s/%s: %v", path, child, parseErr)) + } + } + if len(resolved) > 0 { + break + } + } + resolved = shuffledEndpoints(uniqueEndpoints(resolved), rejected) + if len(resolved) == 0 { + if listedPath == "" { + return nil, fmt.Errorf("HiveServer2 ZooKeeper namespace not found; tried %s", strings.Join(discovery.paths(), ", ")) + } + if len(nodeFailures) > 0 { + return nil, fmt.Errorf("no usable HiveServer2 nodes in ZooKeeper namespace %s: %s", listedPath, strings.Join(nodeFailures, "; ")) + } + return nil, fmt.Errorf("no available HiveServer2 nodes in ZooKeeper namespace %s", listedPath) + } + return resolved, nil +} + +func (discovery *zooKeeperDiscovery) paths() []string { + namespace := strings.Trim(discovery.namespace, "/") + if strings.EqualFold(discovery.discoveryMode, "zookeeperha") { + return []string{ + zooKeeperPath(namespace, "instances"), + zooKeeperPath(namespace+"-unsecure", "instances"), + zooKeeperPath(namespace+"-sasl", "instances"), + } + } + return []string{zooKeeperPath(namespace)} +} + +func zooKeeperPath(parts ...string) string { + cleaned := make([]string, 0, len(parts)) + for _, part := range parts { + if value := strings.Trim(part, "/"); value != "" { + cleaned = append(cleaned, value) + } + } + if len(cleaned) == 0 { + return "/" + } + return "/" + strings.Join(cleaned, "/") +} + +func waitForZooKeeperSession(ctx context.Context, events <-chan zk.Event, timeout time.Duration) error { + timer := time.NewTimer(timeout) + defer timer.Stop() + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return errors.New("ZooKeeper connection timed out before a session was established") + case event, ok := <-events: + if !ok { + return errors.New("ZooKeeper event stream closed before a session was established") + } + if event.Err != nil { + return fmt.Errorf("ZooKeeper connection event: %w", event.Err) + } + switch event.State { + case zk.StateHasSession: + return nil + case zk.StateAuthFailed: + return errors.New("ZooKeeper authentication failed") + case zk.StateExpired: + return errors.New("ZooKeeper session expired during connection") + } + } + } +} + +func parseHiveServerRegistration(child string, data []byte) (endpoint, error) { + candidates := []string{strings.TrimSpace(string(data)), strings.TrimSpace(child)} + for _, candidate := range candidates { + if candidate == "" { + continue + } + if value, err := endpointFromRegistrationJSON(candidate); err == nil { + return value, nil + } + parameters := parseHiveParameters(candidate) + for _, key := range []string{"serveruri", "hiveserver2uri", "server_uri"} { + if raw := parameter(parameters, key); raw != "" { + return parseRegisteredEndpoint(raw) + } + } + if value, err := endpointFromPublishedHiveConfig(parameters); err == nil { + return value, nil + } + if strings.Contains(candidate, "=") { + continue + } + if value, err := parseRegisteredEndpoint(candidate); err == nil { + return value, nil + } + } + return endpoint{}, fmt.Errorf("unsupported HiveServer2 ZooKeeper registration %q", child) +} + +func endpointFromRegistrationJSON(value string) (endpoint, error) { + var object map[string]any + if json.Unmarshal([]byte(value), &object) != nil { + return endpoint{}, errors.New("not JSON") + } + for _, key := range []string{"serverUri", "server_uri", "hiveServer2Uri", "uri"} { + if raw, ok := object[key].(string); ok && strings.TrimSpace(raw) != "" { + return parseRegisteredEndpoint(raw) + } + } + if serviceRecordEndpoint, ok := endpointFromServiceRecord(object); ok { + return serviceRecordEndpoint, nil + } + host, _ := object["host"].(string) + if host == "" { + host, _ = object["hostname"].(string) + } + port := 0 + switch value := object["port"].(type) { + case float64: + port = int(value) + case string: + port, _ = strconv.Atoi(value) + } + if host != "" && port > 0 { + return endpoint{Host: host, Port: port}, nil + } + return endpoint{}, errors.New("JSON registration has no endpoint") +} + +func endpointFromServiceRecord(object map[string]any) (endpoint, bool) { + internal, _ := object["internal"].([]any) + for _, rawEndpoint := range internal { + published, _ := rawEndpoint.(map[string]any) + if !strings.EqualFold(registrationStringValue(published["api"]), "activeEndpoint") { + continue + } + addresses, _ := published["addresses"].([]any) + for _, rawAddress := range addresses { + address, _ := rawAddress.(map[string]any) + host := registrationStringValue(address["host"]) + port, _ := strconv.Atoi(registrationStringValue(address["port"])) + if host != "" && port > 0 { + result := endpoint{Host: host, Port: port} + applyPublishedHiveConfig(&result, object) + return result, true + } + } + } + return endpoint{}, false +} + +func registrationStringValue(value any) string { + switch typed := value.(type) { + case string: + return strings.TrimSpace(typed) + case float64: + return strconv.FormatInt(int64(typed), 10) + case json.Number: + return typed.String() + default: + return "" + } +} + +func applyPublishedHiveConfig(target *endpoint, parameters map[string]any) { + value := func(key string) string { + if nested, ok := parameters["attributes"].(map[string]any); ok { + if result := registrationStringValue(nested[key]); result != "" { + return result + } + } + return registrationStringValue(parameters[key]) + } + target.TransportMode = strings.ToLower(value("hive.server2.transport.mode")) + target.HTTPPath = strings.TrimPrefix(value("hive.server2.thrift.http.path"), "/") + target.Auth = strings.ToUpper(value("hive.server2.authentication")) + target.Principal = value("hive.server2.authentication.kerberos.principal") + target.SSL = strings.EqualFold(value("hive.server2.use.ssl"), "true") +} + +func endpointFromPublishedHiveConfig(parameters map[string]string) (endpoint, error) { + host := firstNonEmpty( + parameter(parameters, "hive.server2.thrift.bind.host"), + parameter(parameters, "host"), + ) + transportMode := strings.ToLower(parameter(parameters, "hive.server2.transport.mode")) + portValue := parameter(parameters, "hive.server2.thrift.port") + if transportMode == "http" { + portValue = firstNonEmpty(parameter(parameters, "hive.server2.thrift.http.port"), portValue) + } + port, err := strconv.Atoi(portValue) + if host == "" || err != nil || port <= 0 { + return endpoint{}, errors.New("published HiveServer2 configuration has no valid host and port") + } + return endpoint{ + Host: host, + Port: port, + TransportMode: transportMode, + HTTPPath: strings.TrimPrefix(parameter(parameters, "hive.server2.thrift.http.path"), "/"), + Auth: strings.ToUpper(parameter(parameters, "hive.server2.authentication")), + Principal: parameter(parameters, "hive.server2.authentication.kerberos.principal"), + SSL: parameterBool(parameters, "hive.server2.use.ssl"), + }, nil +} + +func parseRegisteredEndpoint(value string) (endpoint, error) { + value = strings.TrimSpace(value) + if parsed, err := url.Parse(value); err == nil && parsed.Hostname() != "" { + port := defaultHivePort + if parsed.Port() != "" { + parsedPort, parseErr := strconv.Atoi(parsed.Port()) + if parseErr != nil { + return endpoint{}, parseErr + } + port = parsedPort + } + return endpoint{Host: parsed.Hostname(), Port: port}, nil + } + return parseEndpoint(value, defaultHivePort) +} + +func shuffledEndpoints(values []endpoint, rejected map[string]bool) []endpoint { + result := make([]endpoint, 0, len(values)) + for _, value := range values { + if !rejected[value.address()] { + result = append(result, value) + } + } + rand.Shuffle(len(result), func(first, second int) { + result[first], result[second] = result[second], result[first] + }) + return result +} + +func uniqueEndpoints(values []endpoint) []endpoint { + seen := map[string]bool{} + result := make([]endpoint, 0, len(values)) + for _, value := range values { + key := value.address() + if !seen[key] { + seen[key] = true + result = append(result, value) + } + } + return result +} diff --git a/agents/drivers/hive-go/discovery_test.go b/agents/drivers/hive-go/discovery_test.go new file mode 100644 index 0000000000..dfa14f9e43 --- /dev/null +++ b/agents/drivers/hive-go/discovery_test.go @@ -0,0 +1,221 @@ +package main + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/go-zookeeper/zk" +) + +func TestParseHiveServerRegistrationFromChildName(t *testing.T) { + value, err := parseHiveServerRegistration("serverUri=hs2.example.com:10000;version=4.0.1", nil) + if err != nil { + t.Fatal(err) + } + if value != (endpoint{Host: "hs2.example.com", Port: 10000}) { + t.Fatalf("unexpected endpoint: %#v", value) + } +} + +func TestParseHiveServerRegistrationFromData(t *testing.T) { + value, err := parseHiveServerRegistration("instance-0001", []byte(`{"serverUri":"thrift://kyuubi.example.com:10009"}`)) + if err != nil { + t.Fatal(err) + } + if value != (endpoint{Host: "kyuubi.example.com", Port: 10009}) { + t.Fatalf("unexpected endpoint: %#v", value) + } +} + +func TestRejectedEndpointsAreExcluded(t *testing.T) { + values := shuffledEndpoints([]endpoint{{Host: "one", Port: 1}, {Host: "two", Port: 2}}, map[string]bool{"one:1": true}) + if len(values) != 1 || values[0].Host != "two" { + t.Fatalf("unexpected endpoints: %#v", values) + } +} + +func TestParsePublishedHiveServerConfiguration(t *testing.T) { + value, err := parseHiveServerRegistration("instance-1", []byte( + "hive.server2.thrift.bind.host=hs2.example.com;"+ + "hive.server2.transport.mode=http;"+ + "hive.server2.thrift.http.port=10001;"+ + "hive.server2.thrift.http.path=cliservice;"+ + "hive.server2.authentication=KERBEROS;"+ + "hive.server2.authentication.kerberos.principal=hive/_HOST@EXAMPLE.COM;"+ + "hive.server2.use.ssl=true", + )) + if err != nil { + t.Fatal(err) + } + if value.Host != "hs2.example.com" || value.Port != 10001 || value.TransportMode != "http" || value.HTTPPath != "cliservice" || value.Auth != "KERBEROS" || !value.SSL { + t.Fatalf("unexpected endpoint: %#v", value) + } +} + +func TestParseActivePassiveServiceRecord(t *testing.T) { + value, err := parseHiveServerRegistration("instance-0001", []byte(`{ + "hive.server2.transport.mode":"binary", + "hive.server2.authentication":"KERBEROS", + "hive.server2.authentication.kerberos.principal":"hive/_HOST@EXAMPLE.COM", + "internal":[{ + "api":"activeEndpoint", + "addresses":[{"host":"active.example.com","port":"10000"}] + }] + }`)) + if err != nil { + t.Fatal(err) + } + if value.Host != "active.example.com" || value.Port != 10000 || value.Auth != "KERBEROS" || value.TransportMode != "binary" { + t.Fatalf("unexpected active endpoint: %#v", value) + } +} + +func TestZooKeeperHAPaths(t *testing.T) { + discovery := &zooKeeperDiscovery{namespace: "hs2ActivePassiveHA", discoveryMode: "zooKeeperHA"} + want := []string{ + "/hs2ActivePassiveHA/instances", + "/hs2ActivePassiveHA-unsecure/instances", + "/hs2ActivePassiveHA-sasl/instances", + } + got := discovery.paths() + if len(got) != len(want) { + t.Fatalf("unexpected paths: %#v", got) + } + for index := range want { + if got[index] != want[index] { + t.Fatalf("path %d = %q, want %q", index, got[index], want[index]) + } + } +} + +func TestZooKeeperPathNormalizesEmptyNamespace(t *testing.T) { + if value := zooKeeperPath(""); value != "/" { + t.Fatalf("unexpected root path: %q", value) + } + if value := zooKeeperPath("", "instances"); value != "/instances" { + t.Fatalf("unexpected instances path: %q", value) + } +} + +func TestWaitForZooKeeperSession(t *testing.T) { + events := make(chan zk.Event, 2) + events <- zk.Event{State: zk.StateConnected} + events <- zk.Event{State: zk.StateHasSession} + if err := waitForZooKeeperSession(context.Background(), events, time.Second); err != nil { + t.Fatal(err) + } +} + +func TestWaitForZooKeeperSessionRejectsAuthFailure(t *testing.T) { + events := make(chan zk.Event, 1) + events <- zk.Event{State: zk.StateAuthFailed} + if err := waitForZooKeeperSession(context.Background(), events, time.Second); err == nil { + t.Fatal("expected authentication failure") + } +} + +type fakeZooKeeperClient struct { + children map[string][]string + data map[string][]byte + auth []byte + closed bool +} + +func (client *fakeZooKeeperClient) AddAuth(_ string, auth []byte) error { + client.auth = append([]byte(nil), auth...) + return nil +} + +func (client *fakeZooKeeperClient) Children(path string) ([]string, *zk.Stat, error) { + children, ok := client.children[path] + if !ok { + return nil, nil, zk.ErrNoNode + } + return append([]string(nil), children...), nil, nil +} + +func (client *fakeZooKeeperClient) Get(path string) ([]byte, *zk.Stat, error) { + data, ok := client.data[path] + if !ok { + return nil, nil, zk.ErrNoNode + } + return append([]byte(nil), data...), nil, nil +} + +func (client *fakeZooKeeperClient) Close() { client.closed = true } + +func TestZooKeeperDiscoveryUsesSessionAuthAndSkipsStaleNodes(t *testing.T) { + client := &fakeZooKeeperClient{ + children: map[string][]string{"/hiveserver2": {"stale", "live"}}, + data: map[string][]byte{ + "/hiveserver2/live": []byte("serverUri=live.example.com:10000"), + }, + } + discovery := &zooKeeperDiscovery{ + servers: []endpoint{{Host: "zk.example.com", Port: 2181}}, + namespace: "hiveserver2", + authScheme: "digest", + auth: "user:password", + timeout: time.Second, + discoveryMode: "zookeeper", + dialer: func([]string, time.Duration) (zooKeeperClient, <-chan zk.Event, error) { + events := make(chan zk.Event, 1) + events <- zk.Event{State: zk.StateHasSession} + return client, events, nil + }, + } + values, err := discovery.Endpoints(context.Background(), map[string]bool{}) + if err != nil { + t.Fatal(err) + } + if len(values) != 1 || values[0].Host != "live.example.com" { + t.Fatalf("unexpected endpoints: %#v", values) + } + if string(client.auth) != "user:password" || !client.closed { + t.Fatalf("auth=%q closed=%v", client.auth, client.closed) + } +} + +func TestZooKeeperHAFallsThroughStaleCandidatePath(t *testing.T) { + client := &fakeZooKeeperClient{ + children: map[string][]string{ + "/hs2ActivePassiveHA/instances": {"stale"}, + "/hs2ActivePassiveHA-unsecure/instances": {"live"}, + }, + data: map[string][]byte{ + "/hs2ActivePassiveHA-unsecure/instances/live": []byte("serverUri=active.example.com:10000"), + }, + } + discovery := &zooKeeperDiscovery{ + servers: []endpoint{{Host: "zk.example.com", Port: 2181}}, + namespace: "hs2ActivePassiveHA", + timeout: time.Second, + discoveryMode: "zookeeperHA", + dialer: func([]string, time.Duration) (zooKeeperClient, <-chan zk.Event, error) { + events := make(chan zk.Event, 1) + events <- zk.Event{State: zk.StateHasSession} + return client, events, nil + }, + } + values, err := discovery.Endpoints(context.Background(), map[string]bool{}) + if err != nil { + t.Fatal(err) + } + if len(values) != 1 || values[0].Host != "active.example.com" { + t.Fatalf("unexpected HA endpoints: %#v", values) + } +} + +func TestZooKeeperDiscoveryPropagatesDialError(t *testing.T) { + discovery := &zooKeeperDiscovery{ + servers: []endpoint{{Host: "zk.example.com", Port: 2181}}, + dialer: func([]string, time.Duration) (zooKeeperClient, <-chan zk.Event, error) { + return nil, nil, errors.New("dial failed") + }, + } + if _, err := discovery.Endpoints(context.Background(), nil); err == nil { + t.Fatal("expected dial error") + } +} diff --git a/agents/drivers/hive-go/go.mod b/agents/drivers/hive-go/go.mod new file mode 100644 index 0000000000..59f6ea08e5 --- /dev/null +++ b/agents/drivers/hive-go/go.mod @@ -0,0 +1,35 @@ +module github.com/t8y2/dbx/agents/drivers/hive-go + +go 1.23.0 + +require ( + github.com/beltran/gosasl v1.0.0 + github.com/go-zookeeper/zk v1.0.4 + github.com/golang-auth/go-gssapi/v2 v2.0.0 + github.com/jcmturner/krb5test v0.0.0-20201230140143-102e4b78cdb8 + github.com/pavlo-v-chernykh/keystore-go/v4 v4.5.0 + github.com/t8y2/dbx/agents/go-common/gohive v0.0.0 + software.sslmate.com/src/go-pkcs12 v0.7.3 +) + +require ( + github.com/alexbrainman/sspi v0.0.0-20250919150558-7d374ff0d59e // indirect + github.com/apache/thrift v0.22.0 // indirect + github.com/beltran/gohive/v2 v2.1.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 + github.com/pkg/errors v0.9.1 // indirect + golang.org/x/crypto v0.39.0 // indirect + golang.org/x/net v0.41.0 // indirect +) + +replace github.com/beltran/gosasl => ../../go-common/gosasl + +replace github.com/golang-auth/go-gssapi/v2 => ../../go-common/go-gssapi + +replace github.com/t8y2/dbx/agents/go-common/gohive => ../../go-common/gohive diff --git a/agents/drivers/hive-go/go.sum b/agents/drivers/hive-go/go.sum new file mode 100644 index 0000000000..04407f477b --- /dev/null +++ b/agents/drivers/hive-go/go.sum @@ -0,0 +1,68 @@ +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/apache/thrift v0.22.0 h1:r7mTJdj51TMDe6RtcmNdQxgn9XcyfGDOzegMDRg47uc= +github.com/apache/thrift v0.22.0/go.mod h1:1e7J/O1Ae6ZQMTYdy9xa3w9k+XHWPfRvdPyJeynQ+/g= +github.com/beltran/gohive/v2 v2.1.0 h1:+XPfODaLXsZkD8mEnKCYqinC4OmVMx4/oR4ys8IGKeI= +github.com/beltran/gohive/v2 v2.1.0/go.mod h1:MrxyI0sdriG52I+Bo2lJpGR75bxjcI/Hx3+MBnYqR10= +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.0/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM= +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.0/go.mod h1:T1hnNppQsBtxW0tCHMHTkAt8n/sABdzZgZdoFrZaZNM= +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/krb5test v0.0.0-20201230140143-102e4b78cdb8 h1:bWgHpg2hLK3m8zHGSLqa+UV8XkivFFjG5IniW5lWzIA= +github.com/jcmturner/krb5test v0.0.0-20201230140143-102e4b78cdb8/go.mod h1:eUQyrWMo2eMPWPezAYjLLwaiQHzJUD3su+1vAL10dOk= +github.com/jcmturner/rpc/v2 v2.0.2/go.mod h1:VUJYCIDm3PVOEHw8sgt091/20OJjskO/YJki3ELg/Hc= +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/pavlo-v-chernykh/keystore-go/v4 v4.5.0 h1:2nosf3P75OZv2/ZO/9Px5ZgZ5gbKrzA3joN1QMfOGMQ= +github.com/pavlo-v-chernykh/keystore-go/v4 v4.5.0/go.mod h1:lAVhWwbNaveeJmxrxuSTxMgKpF6DjnuVpn6T8WiBwYQ= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +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-20200117160349-530e935923ad/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/crypto v0.0.0-20201112155050-0c6587e931a9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/crypto v0.39.0 h1:SHs+kF4LP+f+p14esP5jAoDpHU8Gu/v9lFRK6IT5imM= +golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U= +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/go.mod h1:uSPa2vr4CLtc/ILN5odXGNXS6mhrKVzTaCXzk9m6W3k= +golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw= +golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA= +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= +software.sslmate.com/src/go-pkcs12 v0.7.3 h1:JBQD3FDqYjTeyDAeZQklj2ar88ykBLtALloPJHyAauU= +software.sslmate.com/src/go-pkcs12 v0.7.3/go.mod h1:Qiz0EyvDRJjjxGyUQa2cCNZn/wMyzrRJ/qcDXOQazLI= diff --git a/agents/drivers/hive-go/init_test.go b/agents/drivers/hive-go/init_test.go new file mode 100644 index 0000000000..9ad3edd58d --- /dev/null +++ b/agents/drivers/hive-go/init_test.go @@ -0,0 +1,65 @@ +package main + +import ( + "context" + "database/sql/driver" + "errors" + "reflect" + "strings" + "testing" + + "github.com/t8y2/dbx/agents/go-common/gohive" +) + +func TestRunHiveInitStatementsExecutesAndDrainsResults(t *testing.T) { + var resultRows *scriptedRows + behavior := &scriptedBehavior{} + behavior.query = func(ctx context.Context, query string) (driver.Rows, error) { + switch query { + case "SET hive.exec.dynamic.partition=true": + resultRows = newScriptedRows(ctx, []string{"set"}, []string{"STRING"}, [][]driver.Value{ + {"hive.exec.dynamic.partition=true"}, + }) + return resultRows, nil + case "USE analytics": + return nil, &gohive.NonQueryResult{} + default: + return nil, errors.New("unexpected init statement") + } + } + server := newScriptedServer(t, behavior) + + err := runHiveInitStatements( + context.Background(), + server.connection, + []string{"SET hive.exec.dynamic.partition=true", "USE analytics"}, + 77, + ) + if err != nil { + t.Fatal(err) + } + queries, executions, _, _ := behavior.snapshot() + if !reflect.DeepEqual(queries, []string{"SET hive.exec.dynamic.partition=true", "USE analytics"}) || len(executions) != 0 { + t.Fatalf("unexpected init statements: queries=%v executions=%v", queries, executions) + } + if resultRows == nil || !resultRows.isClosed() { + t.Fatal("initFile result set was not drained and closed") + } +} + +func TestRunHiveInitStatementsStopsOnFailure(t *testing.T) { + behavior := &scriptedBehavior{} + behavior.query = func(context.Context, string) (driver.Rows, error) { + return nil, errors.New("permission denied") + } + server := newScriptedServer(t, behavior) + + err := runHiveInitStatements(context.Background(), server.connection, []string{"USE restricted", "USE skipped"}, 100) + if err == nil || !strings.Contains(err.Error(), "execute Hive initFile statement: permission denied") { + t.Fatalf("unexpected initFile error: %v", err) + } + queries, _, _, _ := behavior.snapshot() + if !reflect.DeepEqual(queries, []string{"USE restricted"}) { + t.Fatalf("initFile continued after failure: %v", queries) + } +} diff --git a/agents/drivers/hive-go/kerberos_defaults_unix.go b/agents/drivers/hive-go/kerberos_defaults_unix.go new file mode 100644 index 0000000000..6a6e890932 --- /dev/null +++ b/agents/drivers/hive-go/kerberos_defaults_unix.go @@ -0,0 +1,16 @@ +//go:build !windows + +package main + +import ( + "fmt" + "os" +) + +func defaultKerberosConfigPath() string { + return "/etc/krb5.conf" +} + +func defaultKerberosCCachePath() string { + return fmt.Sprintf("/tmp/krb5cc_%d", os.Getuid()) +} diff --git a/agents/drivers/hive-go/kerberos_defaults_windows.go b/agents/drivers/hive-go/kerberos_defaults_windows.go new file mode 100644 index 0000000000..1ed65fac33 --- /dev/null +++ b/agents/drivers/hive-go/kerberos_defaults_windows.go @@ -0,0 +1,19 @@ +//go:build windows + +package main + +import ( + "os" + "path/filepath" +) + +func defaultKerberosConfigPath() string { + if windowsDirectory := os.Getenv("WINDIR"); windowsDirectory != "" { + return filepath.Join(windowsDirectory, "krb5.ini") + } + return `C:\Windows\krb5.ini` +} + +func defaultKerberosCCachePath() string { + return "" +} diff --git a/agents/drivers/hive-go/main.go b/agents/drivers/hive-go/main.go new file mode 100644 index 0000000000..7ad269e491 --- /dev/null +++ b/agents/drivers/hive-go/main.go @@ -0,0 +1,611 @@ +package main + +import ( + "bufio" + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "os" + "runtime" + "strconv" + "strings" + "sync" + "time" + _ "time/tzdata" +) + +const ( + protocolVersion = 2 + defaultMaxRows = 10000 + defaultPageSize = 1000 + defaultFetchSize = 50 + legacyAgentSessionID = "__legacy__" + maxAgentSessions = 256 + querySessionIdleTime = 10 * time.Minute +) + +type request struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + Params map[string]json.RawMessage `json:"params"` +} + +type response struct { + JSONRPC string `json:"jsonrpc,omitempty"` + ID json.RawMessage `json:"id,omitempty"` + Result any `json:"result,omitempty"` + Error *rpcError `json:"error,omitempty"` +} + +type queryOptions struct { + SQL string `json:"sql"` + Database string `json:"database"` + Schema string `json:"schema"` + MaxRows int `json:"maxRows"` + FetchSize int `json:"fetchSize"` + TimeoutSecs int `json:"timeoutSecs"` +} + +type queryResult struct { + Columns []string `json:"columns"` + ColumnTypes []string `json:"column_types"` + Rows [][]any `json:"rows"` + AffectedRows int64 `json:"affected_rows"` + ExecutionTimeMS int64 `json:"execution_time_ms"` + Truncated bool `json:"truncated"` +} + +type queryPageResult struct { + Columns []string `json:"columns"` + ColumnTypes []string `json:"column_types"` + Rows [][]any `json:"rows"` + AffectedRows int64 `json:"affected_rows"` + ExecutionTimeMS int64 `json:"execution_time_ms"` + Truncated bool `json:"truncated"` + SessionID *string `json:"session_id"` + HasMore bool `json:"has_more"` +} + +type querySession struct { + rows *sql.Rows + columns []string + columnTypes []string + pending []any + remaining int + cancel context.CancelFunc + lastAccessed time.Time +} + +type server struct { + params connectParams + config connectionConfig + + connectionMu sync.Mutex + database *sql.DB + connection *sql.Conn + + querySessions map[string]*querySession + nextSessionID uint64 + + activeMu sync.Mutex + activeCancel context.CancelFunc +} + +type agentSession struct { + server *server + mu sync.Mutex +} + +type runtimeServer struct { + mu sync.RWMutex + sessions map[string]*agentSession +} + +func main() { + configureRuntimeParallelism() + runtimeServer := newRuntimeServer() + encoder := json.NewEncoder(os.Stdout) + var encoderMu sync.Mutex + var requests sync.WaitGroup + fmt.Fprintln(os.Stdout, `{"ready":true}`) + + scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 0, 64*1024), 512*1024*1024) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + var envelope request + if json.Unmarshal([]byte(line), &envelope) == nil && envelope.Method == "shutdown" { + requests.Wait() + result, _ := runtimeServer.handleLine(line) + encoderMu.Lock() + _ = encoder.Encode(result) + encoderMu.Unlock() + return + } + requests.Add(1) + go func(line string) { + defer requests.Done() + result, _ := runtimeServer.handleLine(line) + encoderMu.Lock() + defer encoderMu.Unlock() + if err := encoder.Encode(result); err != nil { + fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err) + } + }(line) + } + requests.Wait() +} + +func configureRuntimeParallelism() { + if raw := strings.TrimSpace(os.Getenv("DBX_AGENT_HIVE_GOMAXPROCS")); raw != "" { + if configured, err := strconv.Atoi(raw); err == nil && configured > 0 { + runtime.GOMAXPROCS(configured) + return + } + } + if strings.TrimSpace(os.Getenv("GOMAXPROCS")) == "" { + runtime.GOMAXPROCS(min(runtime.NumCPU(), 4)) + } +} + +func newRuntimeServer() *runtimeServer { + return &runtimeServer{sessions: map[string]*agentSession{}} +} + +func (runtimeServer *runtimeServer) handleLine(line string) (response, bool) { + var request request + if err := json.Unmarshal([]byte(line), &request); err != nil { + return errorResponse(nil, "", "", err), false + } + if len(request.ID) == 0 { + request.ID = json.RawMessage("1") + } + result, shutdown, err := runtimeServer.dispatch(request.Method, request.Params) + if err != nil { + return errorResponse(request.ID, request.Method, stringParam(request.Params, "agentSessionId"), err), false + } + return response{JSONRPC: "2.0", ID: request.ID, Result: result}, shutdown +} + +func (runtimeServer *runtimeServer) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) { + switch method { + case "handshake": + return handshakeResult(true), false, nil + case "open_session": + id := stringParam(params, "agentSessionId") + if id == "" { + return nil, false, errors.New("agentSessionId is required") + } + var connection connectParams + if err := decodeParams(params, &connection); err != nil { + return nil, false, err + } + return map[string]bool{"ok": true}, false, runtimeServer.openSession(id, connection) + case "close_session": + return map[string]bool{"ok": true}, false, runtimeServer.closeSession(stringParam(params, "agentSessionId")) + case "validate_session": + session, err := runtimeServer.session(stringParam(params, "agentSessionId")) + if err != nil { + return nil, false, err + } + session.mu.Lock() + defer session.mu.Unlock() + return map[string]bool{"ok": true}, false, session.server.validateConnection() + case "cancel_session": + session, err := runtimeServer.session(stringParam(params, "agentSessionId")) + if err != nil { + return nil, false, err + } + session.server.cancelActiveQuery() + return map[string]bool{"ok": true}, false, nil + case "test_connection": + var connection connectParams + if err := decodeParams(params, &connection); err != nil { + return nil, false, err + } + result, err := testConnection(connection) + return result, false, err + case "connect": + var connection connectParams + if err := decodeParams(params, &connection); err != nil { + return nil, false, err + } + _ = runtimeServer.closeSession(legacyAgentSessionID) + return map[string]bool{"ok": true}, false, runtimeServer.openSession(legacyAgentSessionID, connection) + case "disconnect": + return map[string]bool{"ok": true}, false, runtimeServer.closeSession(legacyAgentSessionID) + case "shutdown": + return map[string]bool{"ok": true}, true, runtimeServer.closeAllSessions() + } + + sessionID := stringParam(params, "agentSessionId") + if sessionID == "" { + sessionID = legacyAgentSessionID + } + session, err := runtimeServer.session(sessionID) + if err != nil { + return nil, false, err + } + session.mu.Lock() + defer session.mu.Unlock() + return session.server.dispatch(method, params) +} + +func (runtimeServer *runtimeServer) openSession(id string, params connectParams) error { + runtimeServer.mu.Lock() + if len(runtimeServer.sessions) >= maxAgentSessions { + runtimeServer.mu.Unlock() + return fmt.Errorf("maximum Hive Agent sessions reached (%d)", maxAgentSessions) + } + if _, exists := runtimeServer.sessions[id]; exists { + runtimeServer.mu.Unlock() + return fmt.Errorf("Hive Agent session already exists: %s", id) + } + runtimeServer.mu.Unlock() + + server, err := newServer(params) + if err != nil { + return err + } + runtimeServer.mu.Lock() + defer runtimeServer.mu.Unlock() + if _, exists := runtimeServer.sessions[id]; exists { + _ = server.disconnect() + return fmt.Errorf("Hive Agent session already exists: %s", id) + } + if len(runtimeServer.sessions) >= maxAgentSessions { + _ = server.disconnect() + return fmt.Errorf("maximum Hive Agent sessions reached (%d)", maxAgentSessions) + } + runtimeServer.sessions[id] = &agentSession{server: server} + return nil +} + +func (runtimeServer *runtimeServer) session(id string) (*agentSession, error) { + runtimeServer.mu.RLock() + session := runtimeServer.sessions[id] + runtimeServer.mu.RUnlock() + if session == nil { + return nil, fmt.Errorf("Hive Agent session not found: %s", id) + } + return session, nil +} + +func (runtimeServer *runtimeServer) closeSession(id string) error { + if id == "" { + id = legacyAgentSessionID + } + runtimeServer.mu.Lock() + session := runtimeServer.sessions[id] + delete(runtimeServer.sessions, id) + runtimeServer.mu.Unlock() + if session == nil { + return nil + } + session.mu.Lock() + defer session.mu.Unlock() + return session.server.disconnect() +} + +func (runtimeServer *runtimeServer) closeAllSessions() error { + runtimeServer.mu.Lock() + sessions := runtimeServer.sessions + runtimeServer.sessions = map[string]*agentSession{} + runtimeServer.mu.Unlock() + var failures []string + for id, session := range sessions { + session.mu.Lock() + err := session.server.disconnect() + session.mu.Unlock() + if err != nil { + failures = append(failures, fmt.Sprintf("%s: %v", id, err)) + } + } + if len(failures) > 0 { + return fmt.Errorf("close Hive Agent sessions: %s", strings.Join(failures, "; ")) + } + return nil +} + +func newServer(params connectParams) (*server, error) { + config, err := parseConnectionConfig(params) + if err != nil { + return nil, err + } + server := &server{ + params: params, + config: config, + querySessions: map[string]*querySession{}, + } + if err := server.openConnection(); err != nil { + return nil, err + } + return server, nil +} + +func (server *server) openConnection() error { + server.connectionMu.Lock() + defer server.connectionMu.Unlock() + if server.connection != nil { + return nil + } + database := openHiveDatabase(server.config) + ctx, cancel := context.WithTimeout(context.Background(), server.connectionOpenTimeout()) + var connection *sql.Conn + connection, err := database.Conn(ctx) + if err == nil { + err = connection.PingContext(ctx) + } + cancel() + if err == nil { + err = runHiveInitStatements(context.Background(), connection, server.config.InitStatements, server.config.FetchSize) + } + if err != nil { + if connection != nil { + _ = connection.Close() + } + _ = database.Close() + return err + } + server.database = database + server.connection = connection + return nil +} + +func (server *server) connectionOpenTimeout() time.Duration { + timeout := server.config.ConnectTimeout + if strings.EqualFold(server.config.Auth, "BROWSER") && strings.TrimSpace(server.config.BrowserToken) == "" { + timeout += server.config.BrowserResponseTimeout + } + return timeout +} + +func runHiveInitStatements(ctx context.Context, connection *sql.Conn, statements []string, fetchSize int) error { + for _, statement := range statements { + rows, _, hasResultSet, err := executeHiveStatement(ctx, connection, statement, fetchSize) + if err != nil { + return fmt.Errorf("execute Hive initFile statement: %w", err) + } + if !hasResultSet { + continue + } + for rows.Next() { + } + if rowsErr := rows.Err(); rowsErr != nil { + _ = rows.Close() + return fmt.Errorf("read Hive initFile statement result: %w", rowsErr) + } + if closeErr := rows.Close(); closeErr != nil { + return fmt.Errorf("close Hive initFile statement result: %w", closeErr) + } + } + return nil +} + +func (server *server) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) { + switch method { + case "validate_connection": + return map[string]bool{"ok": true}, false, server.validateConnection() + case "connection_info": + result, err := server.connectionInfo() + return result, false, err + case "list_databases": + result, err := server.listDatabases() + return result, false, err + case "list_schemas": + result, err := server.listSchemas(stringSliceParam(params, "visible_schemas")) + return result, false, err + case "list_tables": + result, err := server.listTables(stringParam(params, "schema"), metadataListConstraintsFromParams(params)) + return result, false, err + case "get_table_comment": + result, err := server.getTableComment(stringParam(params, "schema"), stringParam(params, "table")) + return result, false, err + case "list_objects": + result, err := server.listObjects(stringParam(params, "schema"), metadataListConstraintsFromParams(params)) + return result, false, err + case "list_data_types": + result, err := server.listDataTypes() + return result, false, err + case "completion_assistant_search_v1": + var input completionAssistantRequest + if err := decodeParams(params, &input); err != nil { + return nil, false, err + } + result, err := server.completionAssistantSearch(input) + return result, false, err + case "get_columns": + result, err := server.getColumns(stringParam(params, "schema"), stringParam(params, "table")) + return result, false, err + case "list_indexes": + return []indexInfo{}, false, nil + case "list_foreign_keys": + return []foreignKeyInfo{}, false, nil + case "list_triggers": + return []triggerInfo{}, false, nil + case "list_constraints", "list_partitions", "list_subpartitions": + return []any{}, false, nil + case "get_object_source": + result, err := server.getTableDDL(stringParam(params, "schema"), firstNonEmpty(stringParam(params, "name"), stringParam(params, "table"))) + return result, false, err + case "get_table_ddl": + result, err := server.getTableDDL(stringParam(params, "schema"), stringParam(params, "table")) + return result, false, err + case "get_explain_info": + result, err := server.getExplainInfo(stringParam(params, "sql")) + return map[string]any{"plan": result, "has_actual_stats": false}, false, err + case "execute_query": + result, err := server.executeQuery(queryOptionsFromParams(params)) + return result, false, err + case "execute_query_page", "start_table_read": + result, err := server.executeQueryPage(queryOptionsFromParams(params), intParam(params, "pageSize")) + return result, false, err + case "fetch_query_page", "fetch_table_read_page": + result, err := server.fetchQueryPage(stringParam(params, "sessionId"), intParam(params, "pageSize")) + return result, false, err + case "close_query_session", "close_table_read_session": + return server.closeQuerySession(stringParam(params, "sessionId")), false, nil + case "execute_transaction": + result, err := server.executeStatements(params, true) + return result, false, err + case "execute_batch": + result, err := server.executeStatements(params, false) + return result, false, err + case "disconnect": + return map[string]bool{"ok": true}, false, server.disconnect() + case "shutdown": + return map[string]bool{"ok": true}, true, server.disconnect() + default: + return nil, false, fmt.Errorf("unknown method: %s", method) + } +} + +func handshakeResult(multiSession bool) map[string]any { + capabilities := []string{ + "connect", "test_connection", "metadata", "query", "paged_query", "transaction", "ddl", "structured_error_v1", + } + if multiSession { + capabilities = append(capabilities, "multi_session") + } + return map[string]any{ + "protocolVersion": protocolVersion, + "agentProtocolVersion": protocolVersion, + "capabilities": capabilities, + } +} + +func testConnection(params connectParams) (map[string]any, error) { + server, err := newServer(params) + if err != nil { + return nil, err + } + defer server.disconnect() + if err := server.validateConnection(); err != nil { + return nil, err + } + info, err := server.connectionInfo() + if err != nil { + return nil, err + } + return map[string]any{"ok": true, "info": info}, nil +} + +func (server *server) disconnect() error { + server.cancelActiveQuery() + if err := server.closeAllQuerySessions(); err != nil { + return err + } + server.connectionMu.Lock() + connection := server.connection + database := server.database + server.connection = nil + server.database = nil + server.connectionMu.Unlock() + var failures []string + if connection != nil { + if err := connection.Close(); err != nil { + failures = append(failures, err.Error()) + } + } + if database != nil { + if err := database.Close(); err != nil { + failures = append(failures, err.Error()) + } + } + if len(failures) > 0 { + return errors.New(strings.Join(failures, "; ")) + } + return nil +} + +func (server *server) requireConnection() (*sql.Conn, error) { + server.connectionMu.Lock() + connection := server.connection + server.connectionMu.Unlock() + if connection == nil { + return nil, errors.New("Hive connection is not open") + } + return connection, nil +} + +func (server *server) setActiveOperation(cancel context.CancelFunc) { + server.activeMu.Lock() + server.activeCancel = cancel + server.activeMu.Unlock() +} + +func (server *server) clearActiveOperation(cancel context.CancelFunc) { + cancel() + server.activeMu.Lock() + server.activeCancel = nil + server.activeMu.Unlock() +} + +func (server *server) cancelActiveQuery() { + server.activeMu.Lock() + cancel := server.activeCancel + server.activeMu.Unlock() + if cancel != nil { + cancel() + } +} + +func queryOptionsFromParams(params map[string]json.RawMessage) queryOptions { + return queryOptions{ + SQL: stringParam(params, "sql"), + Database: stringParam(params, "database"), + Schema: stringParam(params, "schema"), + MaxRows: intParam(params, "maxRows"), + FetchSize: intParam(params, "fetchSize"), + TimeoutSecs: intParam(params, "timeoutSecs"), + } +} + +func decodeParams(params map[string]json.RawMessage, target any) error { + data, err := json.Marshal(params) + if err != nil { + return err + } + return json.Unmarshal(data, target) +} + +func stringParam(params map[string]json.RawMessage, key string) string { + if raw, ok := params[key]; ok { + var value string + if json.Unmarshal(raw, &value) == nil { + return value + } + } + return "" +} + +func intParam(params map[string]json.RawMessage, key string) int { + if raw, ok := params[key]; ok { + var value int + if json.Unmarshal(raw, &value) == nil { + return value + } + } + return 0 +} + +func stringSliceParam(params map[string]json.RawMessage, key string) []string { + raw, ok := params[key] + if !ok { + return nil + } + var value []string + if json.Unmarshal(raw, &value) == nil { + return value + } + return nil +} + +func errorResponse(id json.RawMessage, method, sessionID string, err error) response { + return response{JSONRPC: "2.0", ID: id, Error: classifyRPCError(method, sessionID, err)} +} diff --git a/agents/drivers/hive-go/metadata.go b/agents/drivers/hive-go/metadata.go new file mode 100644 index 0000000000..05c1610fac --- /dev/null +++ b/agents/drivers/hive-go/metadata.go @@ -0,0 +1,753 @@ +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sort" + "strconv" + "strings" + + "github.com/t8y2/dbx/agents/go-common/gohive" +) + +const metadataQueryLimit = 100000 + +var hiveTypes = []string{ + "tinyint", "smallint", "int", "bigint", "boolean", "float", "double", "decimal", "string", "varchar", + "char", "binary", "date", "timestamp", "timestamp with local time zone", "interval_year_month", + "interval_day_time", "array", "map", "struct", "uniontype", "void", +} + +type databaseInfo struct { + Name string `json:"name"` +} + +type tableInfo struct { + Name string `json:"name"` + TableType string `json:"table_type"` + Comment *string `json:"comment"` + ParentSchema *string `json:"parent_schema,omitempty"` + ParentName *string `json:"parent_name,omitempty"` +} + +type objectInfo struct { + Name string `json:"name"` + ObjectType string `json:"object_type"` + Schema string `json:"schema"` + Comment *string `json:"comment"` + Valid *bool `json:"valid,omitempty"` +} + +type columnInfo struct { + Name string `json:"name"` + DataType string `json:"data_type"` + IsNullable bool `json:"is_nullable"` + ColumnDefault *string `json:"column_default"` + IsPrimaryKey bool `json:"is_primary_key"` + Extra *string `json:"extra"` + Comment *string `json:"comment"` + NumericPrecision *int `json:"numeric_precision"` + NumericScale *int `json:"numeric_scale"` + CharacterMaximumLength *int `json:"character_maximum_length"` +} + +type indexInfo struct { + Name string `json:"name"` + Columns []string `json:"columns"` + IsUnique bool `json:"is_unique"` + IsPrimary bool `json:"is_primary"` + Filter *string `json:"filter"` + IndexType *string `json:"index_type"` + IncludedColumns []string `json:"included_columns"` + Comment *string `json:"comment"` +} + +func (value indexInfo) MarshalJSON() ([]byte, error) { + type alias indexInfo + copy := alias(value) + if copy.Columns == nil { + copy.Columns = []string{} + } + if copy.IncludedColumns == nil { + copy.IncludedColumns = []string{} + } + return json.Marshal(copy) +} + +type foreignKeyInfo struct { + Name string `json:"name"` + Column string `json:"column"` + RefTable string `json:"ref_table"` + RefColumn string `json:"ref_column"` +} + +type triggerInfo struct { + Name string `json:"name"` + Event string `json:"event"` + Timing string `json:"timing"` +} + +type metadataListConstraints struct { + Filter string + Limit int + Offset int + ObjectTypes []string +} + +type completionAssistantRequest struct { + ConnectionID string `json:"connection_id"` + Database string `json:"database"` + Schema string `json:"schema"` + ObjectKinds []string `json:"object_kinds"` + Mask string `json:"mask"` + CaseSensitive bool `json:"case_sensitive"` + GlobalSearch bool `json:"global_search"` + MaxResults int `json:"max_results"` + ParentSchema string `json:"parent_schema"` + ParentName string `json:"parent_name"` + MatchMode string `json:"match_mode"` +} + +type completionAssistantCandidate struct { + Name string `json:"name"` + Kind string `json:"kind"` + Database *string `json:"database"` + Schema *string `json:"schema"` + ParentSchema *string `json:"parent_schema"` + ParentName *string `json:"parent_name"` + Comment *string `json:"comment"` + DataType *string `json:"data_type"` +} + +type completionAssistantResponse struct { + Candidates []completionAssistantCandidate `json:"candidates"` + Incomplete bool `json:"incomplete"` + FallbackUsed bool `json:"fallback_used"` +} + +func hiveDataTypes() []string { + return append([]string(nil), hiveTypes...) +} + +type hiveMetadataRows struct { + indexes map[string]int + rows [][]any +} + +func newHiveMetadataRows(result gohive.MetadataResult) hiveMetadataRows { + indexes := make(map[string]int, len(result.Columns)) + for index, column := range result.Columns { + indexes[normalizeMetadataColumn(column)] = index + } + rows := make([][]any, 0, len(result.Rows)) + for _, row := range result.Rows { + values := make([]any, len(row)) + for index, value := range row { + values[index] = value + } + rows = append(rows, values) + } + return hiveMetadataRows{indexes: indexes, rows: rows} +} + +func (rows hiveMetadataRows) value(row []any, names ...string) any { + for _, name := range names { + if index, ok := rows.indexes[normalizeMetadataColumn(name)]; ok && index >= 0 && index < len(row) { + return row[index] + } + } + return nil +} + +func normalizeMetadataColumn(value string) string { + return strings.NewReplacer("_", "", "-", "", " ", "").Replace(strings.ToUpper(strings.TrimSpace(value))) +} + +func (server *server) hiveMetadata(operation func(context.Context, gohive.MetadataProvider) (gohive.MetadataResult, error)) (gohive.MetadataResult, error) { + connection, err := server.requireConnection() + if err != nil { + return gohive.MetadataResult{}, err + } + ctx, cancel := context.WithCancel(context.Background()) + server.setActiveOperation(cancel) + defer server.clearActiveOperation(cancel) + var result gohive.MetadataResult + err = connection.Raw(func(rawConnection any) error { + provider, ok := rawConnection.(gohive.MetadataProvider) + if !ok { + return errors.New("Hive driver does not expose HiveServer2 metadata RPCs") + } + var operationErr error + result, operationErr = operation(ctx, provider) + return operationErr + }) + return result, err +} + +func (server *server) connectionInfo() (map[string]any, error) { + version := "" + if result, err := server.executeQuery(queryOptions{SQL: "SELECT VERSION()", MaxRows: 1, TimeoutSecs: 5}); err == nil && len(result.Rows) > 0 && len(result.Rows[0]) > 0 { + version = stringValue(result.Rows[0][0]) + } + username := server.config.Username + if result, err := server.executeQuery(queryOptions{SQL: "SELECT CURRENT_USER()", MaxRows: 1, TimeoutSecs: 5}); err == nil && len(result.Rows) > 0 && len(result.Rows[0]) > 0 { + if current := stringValue(result.Rows[0][0]); current != "" { + username = current + } + } + return map[string]any{ + "database": server.config.Database, + "schema": server.config.Database, + "username": username, + "version": version, + "sqlDialect": "HIVE", + "identifierQuote": "`", + "compatibilityMode": "hive", + "databaseInfo": map[string]string{ + "productName": "Apache Hive", + "productVersion": version, + "unquotedIdentifierCase": "mixed", + "quotedIdentifierCase": "mixed", + "driverName": "DBX Hive Go Agent", + "driverVersion": "gohive-v2.1.0", + }, + }, nil +} + +func (server *server) listDatabases() ([]databaseInfo, error) { + result, err := server.executeQuery(queryOptions{SQL: "SHOW DATABASES", MaxRows: metadataQueryLimit}) + if err == nil { + return databaseInfoFromQueryRows(result.Rows), nil + } + metadataResult, metadataErr := server.hiveMetadata(func(ctx context.Context, provider gohive.MetadataProvider) (gohive.MetadataResult, error) { + return provider.GetHiveSchemas(ctx, "%") + }) + if metadataErr != nil { + return nil, fmt.Errorf("SHOW DATABASES failed (%v); HiveServer2 metadata fallback failed: %w", err, metadataErr) + } + rows := newHiveMetadataRows(metadataResult) + values := make([]databaseInfo, 0, len(rows.rows)) + seen := map[string]bool{} + for _, row := range rows.rows { + name := metadataString(rows.value(row, "TABLE_SCHEM", "SCHEMA_NAME")) + if name == "" || seen[name] { + continue + } + seen[name] = true + values = append(values, databaseInfo{Name: name}) + } + sort.Slice(values, func(first, second int) bool { return values[first].Name < values[second].Name }) + return values, nil +} + +func databaseInfoFromQueryRows(rows [][]any) []databaseInfo { + values := make([]databaseInfo, 0, len(rows)) + seen := map[string]bool{} + for _, row := range rows { + name := firstRowValue(row) + if name == "" || seen[name] { + continue + } + seen[name] = true + values = append(values, databaseInfo{Name: name}) + } + sort.Slice(values, func(first, second int) bool { return values[first].Name < values[second].Name }) + return values +} + +func (server *server) listSchemas(visibleSchemas []string) ([]string, error) { + if visibleSchemas != nil && len(visibleSchemas) == 0 { + return []string{}, nil + } + databases, err := server.listDatabases() + if err != nil { + return nil, err + } + visible := map[string]bool{} + for _, schema := range visibleSchemas { + visible[schema] = true + } + values := make([]string, 0, len(databases)) + for _, database := range databases { + if visibleSchemas != nil && !visible[database.Name] { + continue + } + values = append(values, database.Name) + } + return values, nil +} + +func (server *server) listTables(schema string, constraints metadataListConstraints) ([]tableInfo, error) { + schema = firstNonEmpty(schema, server.config.Database) + requestedTypes := hiveTableTypes(constraints.ObjectTypes) + if len(constraints.ObjectTypes) > 0 && len(requestedTypes) == 0 { + return []tableInfo{}, nil + } + metadataResult, metadataErr := server.hiveMetadata(func(ctx context.Context, provider gohive.MetadataProvider) (gohive.MetadataResult, error) { + return provider.GetHiveTables(ctx, schema, "%", requestedTypes) + }) + if metadataErr == nil { + rows := newHiveMetadataRows(metadataResult) + values := make([]tableInfo, 0, len(rows.rows)) + for _, row := range rows.rows { + name := metadataString(rows.value(row, "TABLE_NAME")) + if name == "" || !metadataNameMatches(name, constraints.Filter) { + continue + } + values = append(values, tableInfo{ + Name: name, + TableType: normalizeHiveTableType(metadataString(rows.value(row, "TABLE_TYPE"))), + Comment: optionalString(metadataString(rows.value(row, "REMARKS", "COMMENT"))), + }) + } + sort.Slice(values, func(first, second int) bool { return values[first].Name < values[second].Name }) + return applyMetadataWindow(values, constraints.Offset, constraints.Limit), nil + } + fallbackStatement := "SHOW TABLES IN " + quoteHiveIdentifier(schema) + if len(requestedTypes) > 0 && !containsString(requestedTypes, "TABLE") { + fallbackStatement = "SHOW VIEWS IN " + quoteHiveIdentifier(schema) + } + result, err := server.executeQuery(queryOptions{ + SQL: fallbackStatement, + MaxRows: metadataQueryLimit, + }) + if err != nil { + return nil, fmt.Errorf("HiveServer2 metadata failed (%v); SHOW TABLES fallback failed: %w", metadataErr, err) + } + values := make([]tableInfo, 0, len(result.Rows)) + for _, row := range result.Rows { + name := showTablesRowName(result.Columns, row) + if name == "" || !metadataNameMatches(name, constraints.Filter) { + continue + } + values = append(values, tableInfo{Name: name, TableType: "TABLE", Comment: nil}) + } + sort.Slice(values, func(first, second int) bool { return values[first].Name < values[second].Name }) + return applyMetadataWindow(values, constraints.Offset, constraints.Limit), nil +} + +func (server *server) listObjects(schema string, constraints metadataListConstraints) ([]objectInfo, error) { + if !acceptsHiveTable(constraints.ObjectTypes) { + return []objectInfo{}, nil + } + tables, err := server.listTables(schema, constraints) + if err != nil { + return nil, err + } + schema = firstNonEmpty(schema, server.config.Database) + values := make([]objectInfo, 0, len(tables)) + for _, table := range tables { + values = append(values, objectInfo{Name: table.Name, ObjectType: table.TableType, Schema: schema, Comment: table.Comment}) + } + return values, nil +} + +func (server *server) getColumns(schema, table string) ([]columnInfo, error) { + if strings.TrimSpace(table) == "" { + return nil, errors.New("table is required") + } + schema = firstNonEmpty(schema, server.config.Database) + metadataResult, metadataErr := server.hiveMetadata(func(ctx context.Context, provider gohive.MetadataProvider) (gohive.MetadataResult, error) { + return provider.GetHiveColumns(ctx, schema, table, "%") + }) + if metadataErr == nil { + rows := newHiveMetadataRows(metadataResult) + values := make([]columnInfo, 0, len(rows.rows)) + for _, row := range rows.rows { + name := metadataString(rows.value(row, "COLUMN_NAME")) + if name == "" { + continue + } + dataType := metadataString(rows.value(row, "TYPE_NAME")) + columnSize := metadataIntPointer(rows.value(row, "COLUMN_SIZE")) + values = append(values, columnInfo{ + Name: name, + DataType: dataType, + IsNullable: metadataNullable(rows.value(row, "NULLABLE", "IS_NULLABLE")), + ColumnDefault: optionalString(metadataString(rows.value(row, "COLUMN_DEF"))), + Comment: optionalString(metadataString(rows.value(row, "REMARKS", "COMMENT"))), + NumericPrecision: columnSize, + NumericScale: metadataIntPointer(rows.value(row, "DECIMAL_DIGITS")), + CharacterMaximumLength: characterLengthForType(dataType, columnSize), + }) + } + return values, nil + } + qualified := qualifiedHiveName(schema, table) + result, err := server.executeQuery(queryOptions{SQL: "DESCRIBE " + qualified, MaxRows: metadataQueryLimit}) + if err != nil { + return nil, fmt.Errorf("HiveServer2 metadata failed (%v); DESCRIBE fallback failed: %w", metadataErr, err) + } + values := make([]columnInfo, 0, len(result.Rows)) + for _, row := range result.Rows { + name := rowString(row, 0) + if name == "" || strings.HasPrefix(name, "#") { + continue + } + dataType := rowString(row, 1) + comment := optionalString(rowString(row, 2)) + values = append(values, columnInfo{ + Name: name, + DataType: dataType, + IsNullable: true, + Comment: comment, + }) + } + return values, nil +} + +func (server *server) getTableComment(schema, table string) (*string, error) { + if strings.TrimSpace(table) == "" { + return nil, errors.New("table is required") + } + schema = firstNonEmpty(schema, server.config.Database) + metadataResult, err := server.hiveMetadata(func(ctx context.Context, provider gohive.MetadataProvider) (gohive.MetadataResult, error) { + return provider.GetHiveTables(ctx, schema, table, nil) + }) + if err != nil { + tables, listErr := server.listTables(schema, metadataListConstraints{Filter: table}) + if listErr != nil { + return nil, fmt.Errorf("HiveServer2 table comment metadata failed (%v); table listing fallback failed: %w", err, listErr) + } + for _, candidate := range tables { + if strings.EqualFold(candidate.Name, table) { + return candidate.Comment, nil + } + } + return nil, nil + } + rows := newHiveMetadataRows(metadataResult) + for _, row := range rows.rows { + if strings.EqualFold(metadataString(rows.value(row, "TABLE_NAME")), table) { + return optionalString(metadataString(rows.value(row, "REMARKS", "COMMENT"))), nil + } + } + return nil, nil +} + +func (server *server) listDataTypes() ([]string, error) { + metadataResult, err := server.hiveMetadata(func(ctx context.Context, provider gohive.MetadataProvider) (gohive.MetadataResult, error) { + return provider.GetHiveTypeInfo(ctx) + }) + if err != nil { + return hiveDataTypes(), nil + } + rows := newHiveMetadataRows(metadataResult) + values := make([]string, 0, len(rows.rows)) + seen := map[string]bool{} + for _, row := range rows.rows { + name := strings.ToLower(metadataString(rows.value(row, "TYPE_NAME"))) + if name == "" || seen[name] { + continue + } + seen[name] = true + values = append(values, name) + } + if len(values) == 0 { + return hiveDataTypes(), nil + } + sort.Strings(values) + return values, nil +} + +func (server *server) getTableDDL(schema, table string) (string, error) { + if strings.TrimSpace(table) == "" { + return "", errors.New("table is required") + } + result, err := server.executeQuery(queryOptions{ + SQL: "SHOW CREATE TABLE " + qualifiedHiveName(schema, table), + MaxRows: metadataQueryLimit, + }) + if err != nil { + return "", err + } + lines := make([]string, 0, len(result.Rows)) + for _, row := range result.Rows { + if line := firstRowValue(row); line != "" { + lines = append(lines, line) + } + } + if len(lines) == 0 { + return "", nil + } + return strings.Join(lines, "\n") + "\n", nil +} + +func (server *server) getExplainInfo(sqlText string) (string, error) { + sqlText = trimStatementSQL(sqlText) + if sqlText == "" { + return "", errors.New("SQL is required") + } + result, err := server.executeQuery(queryOptions{SQL: "EXPLAIN " + sqlText, MaxRows: metadataQueryLimit}) + if err != nil { + return "", err + } + lines := make([]string, 0, len(result.Rows)) + for _, row := range result.Rows { + lines = append(lines, firstRowValue(row)) + } + return strings.Join(lines, "\n"), nil +} + +func (server *server) completionAssistantSearch(input completionAssistantRequest) (completionAssistantResponse, error) { + maxResults := input.MaxResults + if maxResults <= 0 { + maxResults = 200 + } + schemas := []string{firstNonEmpty(input.Schema, input.Database, server.config.Database)} + if input.GlobalSearch { + listed, err := server.listSchemas(nil) + if err != nil { + return completionAssistantResponse{}, err + } + schemas = listed + } + values := make([]completionAssistantCandidate, 0, maxResults) + incomplete := false + for _, schema := range schemas { + tables, err := server.listTables(schema, metadataListConstraints{Limit: maxResults}) + if err != nil { + return completionAssistantResponse{}, err + } + for _, table := range tables { + if !completionNameMatches(table.Name, input) { + continue + } + schemaCopy := schema + values = append(values, completionAssistantCandidate{ + Name: table.Name, Kind: "table", Database: &schemaCopy, Schema: &schemaCopy, Comment: table.Comment, + }) + if len(values) >= maxResults { + incomplete = true + break + } + } + if len(values) >= maxResults { + break + } + } + return completionAssistantResponse{Candidates: values, Incomplete: incomplete, FallbackUsed: false}, nil +} + +func metadataListConstraintsFromParams(params map[string]json.RawMessage) metadataListConstraints { + return metadataListConstraints{ + Filter: stringParam(params, "filter"), + Limit: intParam(params, "limit"), + Offset: intParam(params, "offset"), + ObjectTypes: stringSliceParam(params, "objectTypes"), + } +} + +func qualifiedHiveName(schema, table string) string { + if strings.TrimSpace(schema) == "" { + return quoteHiveIdentifier(table) + } + return quoteHiveIdentifier(schema) + "." + quoteHiveIdentifier(table) +} + +func metadataNameMatches(name, filter string) bool { + return filter == "" || strings.Contains(strings.ToLower(name), strings.ToLower(filter)) +} + +func acceptsHiveTable(objectTypes []string) bool { + if len(objectTypes) == 0 { + return true + } + for _, objectType := range objectTypes { + if strings.EqualFold(objectType, "table") || strings.EqualFold(objectType, "view") || strings.EqualFold(objectType, "materialized view") { + return true + } + } + return false +} + +func containsString(values []string, expected string) bool { + for _, value := range values { + if value == expected { + return true + } + } + return false +} + +func hiveTableTypes(objectTypes []string) []string { + if len(objectTypes) == 0 { + return []string{"TABLE", "VIEW", "MATERIALIZED VIEW"} + } + values := make([]string, 0, len(objectTypes)) + seen := map[string]bool{} + for _, objectType := range objectTypes { + normalized := strings.ToUpper(strings.TrimSpace(objectType)) + switch normalized { + case "TABLE", "EXTERNAL TABLE", "MANAGED TABLE": + normalized = "TABLE" + case "VIEW": + normalized = "VIEW" + case "MATERIALIZED VIEW", "MATERIALIZED_VIEW": + normalized = "MATERIALIZED VIEW" + default: + continue + } + if !seen[normalized] { + seen[normalized] = true + values = append(values, normalized) + } + } + return values +} + +func normalizeHiveTableType(value string) string { + if strings.Contains(strings.ToUpper(value), "VIEW") { + return "VIEW" + } + return "TABLE" +} + +func metadataString(value any) string { + switch typed := value.(type) { + case nil: + return "" + case string: + return strings.TrimSpace(typed) + case []byte: + return strings.TrimSpace(string(typed)) + default: + return strings.TrimSpace(fmt.Sprint(typed)) + } +} + +func metadataIntPointer(value any) *int { + var parsed int64 + switch typed := value.(type) { + case nil: + return nil + case int: + parsed = int64(typed) + case int8: + parsed = int64(typed) + case int16: + parsed = int64(typed) + case int32: + parsed = int64(typed) + case int64: + parsed = typed + case float32: + parsed = int64(typed) + case float64: + parsed = int64(typed) + default: + value, err := strconv.ParseInt(metadataString(value), 10, 64) + if err != nil { + return nil + } + parsed = value + } + if parsed < 0 || parsed > int64(^uint(0)>>1) { + return nil + } + converted := int(parsed) + return &converted +} + +func metadataNullable(value any) bool { + if parsed := metadataIntPointer(value); parsed != nil { + return *parsed != 0 + } + switch strings.ToUpper(metadataString(value)) { + case "NO", "FALSE", "NOT NULL": + return false + default: + return true + } +} + +func characterLengthForType(dataType string, size *int) *int { + normalized := strings.ToLower(dataType) + if strings.Contains(normalized, "char") || strings.Contains(normalized, "text") || strings.Contains(normalized, "string") { + return size + } + return nil +} + +func applyMetadataWindow[T any](values []T, offset, limit int) []T { + if offset < 0 { + offset = 0 + } + if offset >= len(values) { + return []T{} + } + end := len(values) + if limit > 0 && offset+limit < end { + end = offset + limit + } + return values[offset:end] +} + +func completionNameMatches(name string, input completionAssistantRequest) bool { + mask := input.Mask + if mask == "" { + return true + } + if !input.CaseSensitive { + name = strings.ToLower(name) + mask = strings.ToLower(mask) + } + switch strings.ToLower(input.MatchMode) { + case "exact": + return name == mask + case "prefix": + return strings.HasPrefix(name, mask) + default: + return strings.Contains(name, mask) + } +} + +func firstRowValue(row []any) string { + for _, value := range row { + if text := stringValue(value); text != "" { + return text + } + } + return "" +} + +func showTablesRowName(columns []string, row []any) string { + for index, column := range columns { + normalized := strings.NewReplacer("_", "", "-", "", " ", "").Replace(strings.ToLower(column)) + if normalized == "tablename" || normalized == "tabname" { + if value := rowString(row, index); value != "" { + return value + } + } + } + if len(row) > 1 { + if value := rowString(row, 1); value != "" { + return value + } + } + return firstRowValue(row) +} + +func rowString(row []any, index int) string { + if index < 0 || index >= len(row) { + return "" + } + return strings.TrimSpace(stringValue(row[index])) +} + +func stringValue(value any) string { + if value == nil { + return "" + } + return fmt.Sprint(value) +} + +func optionalString(value string) *string { + value = strings.TrimSpace(value) + if value == "" { + return nil + } + return &value +} diff --git a/agents/drivers/hive-go/metadata_test.go b/agents/drivers/hive-go/metadata_test.go new file mode 100644 index 0000000000..6eb52b07a1 --- /dev/null +++ b/agents/drivers/hive-go/metadata_test.go @@ -0,0 +1,212 @@ +package main + +import ( + "context" + "database/sql/driver" + "errors" + "reflect" + "testing" + + "github.com/t8y2/dbx/agents/go-common/gohive" +) + +func TestShowTablesRowName(t *testing.T) { + if value := showTablesRowName([]string{"database", "tableName", "isTemporary"}, []any{"default", "events", false}); value != "events" { + t.Fatalf("unexpected table name: %q", value) + } + if value := showTablesRowName([]string{"tab_name"}, []any{"fallback"}); value != "fallback" { + t.Fatalf("unexpected fallback table name: %q", value) + } +} + +func TestListDatabasesUsesShowDatabasesBeforeHiveServerMetadata(t *testing.T) { + behavior := &scriptedBehavior{ + query: func(ctx context.Context, sql string) (driver.Rows, error) { + if sql != "SHOW DATABASES" { + t.Fatalf("unexpected query: %q", sql) + } + return newScriptedRows(ctx, []string{"database_name"}, []string{"STRING"}, [][]driver.Value{{"warehouse"}, {"default"}, {"default"}}), nil + }, + getSchemas: func(_ context.Context, pattern string) (gohive.MetadataResult, error) { + t.Fatalf("metadata fallback must not run after SHOW DATABASES succeeds: %q", pattern) + return gohive.MetadataResult{}, nil + }, + } + server := newScriptedServer(t, behavior) + defer server.disconnect() + values, err := server.listDatabases() + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(values, []databaseInfo{{Name: "default"}, {Name: "warehouse"}}) { + t.Fatalf("unexpected databases: %#v", values) + } +} + +func TestListDatabasesFallsBackToHiveServerMetadata(t *testing.T) { + behavior := &scriptedBehavior{ + query: func(context.Context, string) (driver.Rows, error) { + return nil, errors.New("SHOW DATABASES unsupported") + }, + getSchemas: func(_ context.Context, pattern string) (gohive.MetadataResult, error) { + if pattern != "%" { + t.Fatalf("unexpected schema pattern: %q", pattern) + } + return metadataResult([]string{"TABLE_SCHEM", "TABLE_CATALOG"}, []driver.Value{"warehouse", ""}, []driver.Value{"default", ""}, []driver.Value{"default", ""}), nil + }, + } + server := newScriptedServer(t, behavior) + defer server.disconnect() + values, err := server.listDatabases() + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(values, []databaseInfo{{Name: "default"}, {Name: "warehouse"}}) { + t.Fatalf("unexpected databases: %#v", values) + } +} + +func TestListSchemasHonorsVisibleSchemaFilter(t *testing.T) { + behavior := &scriptedBehavior{ + query: func(ctx context.Context, sql string) (driver.Rows, error) { + if sql != "SHOW DATABASES" { + t.Fatalf("unexpected query: %q", sql) + } + return newScriptedRows( + ctx, + []string{"database_name"}, + []string{"STRING"}, + [][]driver.Value{{"default"}, {"analytics"}, {"system"}}, + ), nil + }, + } + server := newScriptedServer(t, behavior) + defer server.disconnect() + + values, err := server.listSchemas([]string{"analytics", "missing"}) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(values, []string{"analytics"}) { + t.Fatalf("unexpected visible schemas: %#v", values) + } + + values, err = server.listSchemas([]string{}) + if err != nil { + t.Fatal(err) + } + if len(values) != 0 { + t.Fatalf("explicit empty visible schema filter must hide all schemas: %#v", values) + } +} + +func TestListTablesPreservesViewTypeCommentAndWindow(t *testing.T) { + behavior := &scriptedBehavior{ + getTables: func(_ context.Context, schema, table string, tableTypes []string) (gohive.MetadataResult, error) { + if schema != "analytics" || table != "%" || !reflect.DeepEqual(tableTypes, []string{"TABLE", "VIEW", "MATERIALIZED VIEW"}) { + t.Fatalf("unexpected GetTables request: schema=%q table=%q types=%#v", schema, table, tableTypes) + } + return metadataResult( + []string{"TABLE_CAT", "TABLE_SCHEM", "TABLE_NAME", "TABLE_TYPE", "REMARKS"}, + []driver.Value{"", "analytics", "events", "TABLE", "event data"}, + []driver.Value{"", "analytics", "events_view", "VIEW", "view data"}, + []driver.Value{"", "analytics", "other", "TABLE", nil}, + ), nil + }, + } + server := newScriptedServer(t, behavior) + defer server.disconnect() + values, err := server.listTables("analytics", metadataListConstraints{Filter: "events", Offset: 1, Limit: 1}) + if err != nil { + t.Fatal(err) + } + comment := "view data" + expected := []tableInfo{{Name: "events_view", TableType: "VIEW", Comment: &comment}} + if !reflect.DeepEqual(values, expected) { + t.Fatalf("unexpected tables: %#v", values) + } +} + +func TestGetColumnsUsesHiveServerMetadataFields(t *testing.T) { + behavior := &scriptedBehavior{ + getColumns: func(_ context.Context, schema, table, column string) (gohive.MetadataResult, error) { + if schema != "analytics" || table != "events" || column != "%" { + t.Fatalf("unexpected GetColumns request: %q %q %q", schema, table, column) + } + return metadataResult( + []string{"COLUMN_NAME", "TYPE_NAME", "COLUMN_SIZE", "DECIMAL_DIGITS", "NULLABLE", "REMARKS", "COLUMN_DEF"}, + []driver.Value{"name", "string", int64(255), nil, int64(1), "display name", "unknown"}, + []driver.Value{"amount", "decimal(18,2)", int64(18), int64(2), int64(0), nil, nil}, + ), nil + }, + } + server := newScriptedServer(t, behavior) + defer server.disconnect() + values, err := server.getColumns("analytics", "events") + if err != nil { + t.Fatal(err) + } + if len(values) != 2 { + t.Fatalf("unexpected columns: %#v", values) + } + if values[0].Name != "name" || !values[0].IsNullable || values[0].CharacterMaximumLength == nil || *values[0].CharacterMaximumLength != 255 || values[0].ColumnDefault == nil || *values[0].ColumnDefault != "unknown" { + t.Fatalf("unexpected string column: %#v", values[0]) + } + if values[1].Name != "amount" || values[1].IsNullable || values[1].NumericPrecision == nil || *values[1].NumericPrecision != 18 || values[1].NumericScale == nil || *values[1].NumericScale != 2 || values[1].CharacterMaximumLength != nil { + t.Fatalf("unexpected decimal column: %#v", values[1]) + } +} + +func TestTableCommentAndTypeInfoUseHiveServerMetadata(t *testing.T) { + behavior := &scriptedBehavior{ + getTables: func(_ context.Context, schema, table string, tableTypes []string) (gohive.MetadataResult, error) { + return metadataResult( + []string{"TABLE_SCHEM", "TABLE_NAME", "TABLE_TYPE", "REMARKS"}, + []driver.Value{schema, table, "TABLE", "table comment"}, + ), nil + }, + getTypeInfo: func(context.Context) (gohive.MetadataResult, error) { + return metadataResult([]string{"TYPE_NAME"}, []driver.Value{"STRING"}, []driver.Value{"decimal"}, []driver.Value{"STRING"}), nil + }, + } + server := newScriptedServer(t, behavior) + defer server.disconnect() + comment, err := server.getTableComment("analytics", "events") + if err != nil || comment == nil || *comment != "table comment" { + t.Fatalf("unexpected table comment: %v, %v", comment, err) + } + types, err := server.listDataTypes() + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(types, []string{"decimal", "string"}) { + t.Fatalf("unexpected data types: %#v", types) + } +} + +func TestListTablesFallsBackToShowTables(t *testing.T) { + behavior := &scriptedBehavior{ + getTables: func(context.Context, string, string, []string) (gohive.MetadataResult, error) { + return gohive.MetadataResult{}, errors.New("metadata unsupported") + }, + query: func(ctx context.Context, query string) (driver.Rows, error) { + if query != "SHOW TABLES IN `analytics`" { + t.Fatalf("unexpected fallback query: %q", query) + } + return newScriptedRows(ctx, []string{"tab_name"}, []string{"STRING"}, [][]driver.Value{{"events"}}), nil + }, + } + server := newScriptedServer(t, behavior) + defer server.disconnect() + values, err := server.listTables("analytics", metadataListConstraints{}) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(values, []tableInfo{{Name: "events", TableType: "TABLE"}}) { + t.Fatalf("unexpected fallback tables: %#v", values) + } +} + +func metadataResult(columns []string, rows ...[]driver.Value) gohive.MetadataResult { + return gohive.MetadataResult{Columns: columns, Rows: rows} +} diff --git a/agents/drivers/hive-go/protocol_error.go b/agents/drivers/hive-go/protocol_error.go new file mode 100644 index 0000000000..76cfb36e1a --- /dev/null +++ b/agents/drivers/hive-go/protocol_error.go @@ -0,0 +1,145 @@ +package main + +import ( + "context" + "database/sql" + "errors" + "fmt" + "io" + "net" + "strings" + + "github.com/t8y2/dbx/agents/go-common/gohive" +) + +type rpcError struct { + Code int `json:"code"` + Message string `json:"message"` + Data *rpcErrorData `json:"data,omitempty"` +} + +type rpcErrorData struct { + Category string `json:"category"` + Retryable bool `json:"retryable"` + SessionDisposition string `json:"sessionDisposition"` + Stage string `json:"stage"` + ContractVersion int `json:"contractVersion"` + OperationOutcome string `json:"operationOutcome"` + SQLState string `json:"sqlState,omitempty"` + VendorCode int32 `json:"vendorCode,omitempty"` + ExceptionClass string `json:"exceptionClass,omitempty"` + AgentSessionID string `json:"agentSessionId,omitempty"` +} + +func classifyRPCError(method, agentSessionID string, err error) *rpcError { + stage := rpcErrorStage(method) + data := &rpcErrorData{ + Category: "protocol", + Retryable: false, + SessionDisposition: "keep", + Stage: stage, + ContractVersion: 1, + OperationOutcome: rpcOperationOutcome(stage), + ExceptionClass: safeRPCDiagnostic(fmt.Sprintf("%T", err), 160), + AgentSessionID: strings.TrimSpace(agentSessionID), + } + + var hiveError *gohive.Error + if errors.As(err, &hiveError) { + data.Category = "sql" + data.SQLState = safeRPCDiagnostic(hiveError.SQLState, 160) + data.VendorCode = int32(hiveError.ErrorCode) + data.ExceptionClass = "gohive.Error" + } else if errors.Is(err, context.Canceled) { + data.Category = "canceled" + data.SessionDisposition = "quarantine" + } else if errors.Is(err, context.DeadlineExceeded) || isTimeoutError(err) { + data.Category = "timeout" + data.Retryable = stage == "connect" || stage == "validate" + data.SessionDisposition = "quarantine" + } else if isConnectionError(err) { + data.Category = "connection" + data.Retryable = stage == "connect" || stage == "validate" + if stage != "connect" { + data.SessionDisposition = "quarantine" + } + } else if isHiveSQLError(err) { + data.Category = "sql" + } + return &rpcError{Code: -1, Message: err.Error(), Data: data} +} + +func rpcErrorStage(method string) string { + switch method { + case "connect", "open_session", "test_connection": + return "connect" + case "validate_connection", "validate_session": + return "validate" + case "cancel_session": + return "cancel" + case "close_session", "disconnect", "close_query_session", "close_table_read_session", "shutdown": + return "close" + case "fetch_query_page", "fetch_table_read_page": + return "fetch" + case "handshake", "": + return "request" + default: + return "execute" + } +} + +func rpcOperationOutcome(stage string) string { + if stage == "request" || stage == "connect" || stage == "validate" { + return "not_started" + } + return "unknown" +} + +func isTimeoutError(err error) bool { + var timeout interface{ Timeout() bool } + return errors.As(err, &timeout) && timeout.Timeout() +} + +func isConnectionError(err error) bool { + if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) || errors.Is(err, sql.ErrConnDone) { + return true + } + var networkError *net.OpError + if errors.As(err, &networkError) { + return true + } + lower := strings.ToLower(err.Error()) + for _, marker := range []string{ + "connection refused", "connection reset", "broken pipe", "connection closed", "connection lost", + "unexpected eof", "no route to host", "all hiveserver2 endpoints failed", "transport is not open", + "socket is closed", "not open", + } { + if strings.Contains(lower, marker) { + return true + } + } + return false +} + +func isHiveSQLError(err error) bool { + lower := strings.ToLower(err.Error()) + for _, marker := range []string{"semanticexception", "parseexception", "hive error", "sqlstate", "error while compiling statement"} { + if strings.Contains(lower, marker) { + return true + } + } + return false +} + +func safeRPCDiagnostic(value string, maxLength int) string { + var result strings.Builder + for _, char := range value { + if result.Len() >= maxLength { + break + } + if char >= 0x21 && char <= 0x7e { + result.WriteRune(char) + } + } + return result.String() +} diff --git a/agents/drivers/hive-go/query.go b/agents/drivers/hive-go/query.go new file mode 100644 index 0000000000..734d4fa967 --- /dev/null +++ b/agents/drivers/hive-go/query.go @@ -0,0 +1,529 @@ +package main + +import ( + "context" + "database/sql" + "database/sql/driver" + "encoding/json" + "errors" + "fmt" + "strings" + "time" + + "github.com/t8y2/dbx/agents/go-common/gohive" +) + +func (server *server) validateConnection() error { + connection, err := server.requireConnection() + if err != nil { + return err + } + ctx, cancel := context.WithTimeout(context.Background(), server.config.ConnectTimeout) + defer cancel() + return connection.PingContext(ctx) +} + +func (server *server) executeQuery(options queryOptions) (queryResult, error) { + started := time.Now() + if options.FetchSize <= 0 { + options.FetchSize = server.effectiveFetchSize() + } + sqlText := trimStatementSQL(options.SQL) + if sqlText == "" { + return queryResult{}, errors.New("SQL is required") + } + maxRows := options.MaxRows + if maxRows <= 0 { + maxRows = defaultMaxRows + } + connection, err := server.requireConnection() + if err != nil { + return queryResult{}, err + } + + ctx, cancel := queryContext(options.TimeoutSecs) + server.setActiveOperation(cancel) + defer server.clearActiveOperation(cancel) + if err := server.applySchemaContext(ctx, connection, effectiveSchema(options)); err != nil { + return queryResult{}, err + } + rows, affected, hasResultSet, err := executeHiveStatement(ctx, connection, sqlText, options.FetchSize) + if err != nil { + return queryResult{}, err + } + if !hasResultSet { + return queryResult{ + Columns: []string{}, + ColumnTypes: []string{}, + Rows: [][]any{}, + AffectedRows: affected, + ExecutionTimeMS: time.Since(started).Milliseconds(), + Truncated: false, + }, nil + } + defer rows.Close() + columns, columnTypes, err := queryColumnMetadata(rows) + if err != nil { + return queryResult{}, err + } + values, truncated, err := readSQLRows(rows, columnTypes, maxRows) + if err != nil { + return queryResult{}, err + } + return queryResult{ + Columns: columns, + ColumnTypes: columnTypes, + Rows: values, + AffectedRows: 0, + ExecutionTimeMS: time.Since(started).Milliseconds(), + Truncated: truncated, + }, nil +} + +func (server *server) executeQueryPage(options queryOptions, requestedPageSize int) (queryPageResult, error) { + started := time.Now() + if options.FetchSize <= 0 { + options.FetchSize = server.effectiveFetchSize() + } + server.expireIdleQuerySessions(started) + sqlText := trimStatementSQL(options.SQL) + if sqlText == "" { + return queryPageResult{}, errors.New("SQL is required") + } + pageSize := requestedPageSize + if pageSize <= 0 { + pageSize = defaultPageSize + } + maxRows := options.MaxRows + if maxRows <= 0 { + maxRows = defaultMaxRows + } + connection, err := server.requireConnection() + if err != nil { + return queryPageResult{}, err + } + ctx, cancel := queryContext(options.TimeoutSecs) + server.setActiveOperation(cancel) + if err := server.applySchemaContext(ctx, connection, effectiveSchema(options)); err != nil { + server.clearActiveOperation(cancel) + return queryPageResult{}, err + } + rows, affected, hasResultSet, err := executeHiveStatement(ctx, connection, sqlText, options.FetchSize) + if err != nil { + server.clearActiveOperation(cancel) + return queryPageResult{}, err + } + if !hasResultSet { + server.clearActiveOperation(cancel) + return queryPageResult{ + Columns: []string{}, + ColumnTypes: []string{}, + Rows: [][]any{}, + AffectedRows: affected, + ExecutionTimeMS: time.Since(started).Milliseconds(), + Truncated: false, + SessionID: nil, + HasMore: false, + }, nil + } + + columns, columnTypes, err := queryColumnMetadata(rows) + if err != nil { + _ = rows.Close() + server.clearActiveOperation(cancel) + return queryPageResult{}, err + } + server.nextSessionID++ + sessionID := fmt.Sprintf("hive-%d", server.nextSessionID) + state := &querySession{ + rows: rows, + columns: columns, + columnTypes: columnTypes, + remaining: maxRows, + cancel: cancel, + lastAccessed: started, + } + server.querySessions[sessionID] = state + page, hasMore, truncated, err := server.readQuerySessionPage(ctx, state, pageSize) + server.activeMu.Lock() + if server.activeCancel != nil { + server.activeCancel = nil + } + server.activeMu.Unlock() + if err != nil { + server.closeQuerySession(sessionID) + return queryPageResult{}, err + } + if !hasMore { + server.closeQuerySession(sessionID) + return queryPageResult{ + Columns: columns, + ColumnTypes: columnTypes, + Rows: page, + ExecutionTimeMS: time.Since(started).Milliseconds(), + Truncated: truncated, + SessionID: nil, + HasMore: false, + }, nil + } + return queryPageResult{ + Columns: columns, + ColumnTypes: columnTypes, + Rows: page, + ExecutionTimeMS: time.Since(started).Milliseconds(), + Truncated: false, + SessionID: &sessionID, + HasMore: true, + }, nil +} + +func (server *server) fetchQueryPage(sessionID string, requestedPageSize int) (queryPageResult, error) { + server.expireIdleQuerySessions(time.Now()) + state := server.querySessions[sessionID] + if state == nil { + return queryPageResult{}, errors.New("query session not found") + } + pageSize := requestedPageSize + if pageSize <= 0 { + pageSize = defaultPageSize + } + ctx := context.Background() + server.setActiveOperation(state.cancel) + page, hasMore, truncated, err := server.readQuerySessionPage(ctx, state, pageSize) + server.activeMu.Lock() + server.activeCancel = nil + server.activeMu.Unlock() + if err != nil { + server.closeQuerySession(sessionID) + return queryPageResult{}, err + } + var resultSessionID *string + if hasMore { + resultSessionID = &sessionID + } else { + server.closeQuerySession(sessionID) + } + return queryPageResult{ + Columns: state.columns, + ColumnTypes: state.columnTypes, + Rows: page, + Truncated: truncated, + SessionID: resultSessionID, + HasMore: hasMore, + }, nil +} + +func (server *server) readQuerySessionPage(ctx context.Context, state *querySession, pageSize int) ([][]any, bool, bool, error) { + state.lastAccessed = time.Now() + values := make([][]any, 0, min(pageSize, state.remaining)) + if state.pending != nil && state.remaining > 0 { + values = append(values, state.pending) + state.pending = nil + state.remaining-- + } + for len(values) < pageSize && state.remaining > 0 { + row, ok, err := nextSQLRow(state.rows, state.columnTypes) + if err != nil { + return nil, false, false, err + } + if !ok { + return values, false, false, nil + } + values = append(values, row) + state.remaining-- + select { + case <-ctx.Done(): + return nil, false, false, ctx.Err() + default: + } + } + if state.remaining == 0 { + _, ok, err := nextSQLRow(state.rows, state.columnTypes) + if err != nil { + return nil, false, false, err + } + return values, false, ok, nil + } + row, ok, err := nextSQLRow(state.rows, state.columnTypes) + if err != nil { + return nil, false, false, err + } + if !ok { + return values, false, false, nil + } + state.pending = row + return values, true, false, nil +} + +func (server *server) closeQuerySession(sessionID string) bool { + state := server.querySessions[sessionID] + if state == nil { + return false + } + delete(server.querySessions, sessionID) + state.cancel() + _ = state.rows.Close() + return true +} + +func (server *server) closeAllQuerySessions() error { + var failures []string + for sessionID, state := range server.querySessions { + delete(server.querySessions, sessionID) + state.cancel() + if err := state.rows.Close(); err != nil { + failures = append(failures, fmt.Sprintf("%s: %v", sessionID, err)) + } + } + if len(failures) > 0 { + return errors.New(strings.Join(failures, "; ")) + } + return nil +} + +func (server *server) expireIdleQuerySessions(now time.Time) int { + expired := make([]string, 0) + for sessionID, state := range server.querySessions { + if !state.lastAccessed.IsZero() && now.Sub(state.lastAccessed) >= querySessionIdleTime { + expired = append(expired, sessionID) + } + } + for _, sessionID := range expired { + server.closeQuerySession(sessionID) + } + return len(expired) +} + +func (server *server) executeStatements(params map[string]json.RawMessage, transaction bool) (queryResult, error) { + started := time.Now() + statements := stringSliceParam(params, "statements") + if len(statements) == 0 { + return queryResult{}, errors.New("statements are required") + } + connection, err := server.requireConnection() + if err != nil { + return queryResult{}, err + } + ctx, cancel := queryContext(intParam(params, "timeoutSecs")) + server.setActiveOperation(cancel) + defer server.clearActiveOperation(cancel) + if err := server.applySchemaContext(ctx, connection, firstNonEmpty(stringParam(params, "schema"), stringParam(params, "database"))); err != nil { + return queryResult{}, err + } + + var affected int64 + if transaction { + tx, beginErr := connection.BeginTx(ctx, nil) + if beginErr == nil { + for _, statement := range statements { + trimmed := trimStatementSQL(statement) + if trimmed == "" { + continue + } + result, execErr := tx.ExecContext(ctx, trimmed) + if execErr != nil { + _ = tx.Rollback() + return queryResult{}, execErr + } + count, _ := result.RowsAffected() + affected += max(count, 0) + } + if err := tx.Commit(); err != nil { + return queryResult{}, err + } + return emptyQueryResult(affected, started), nil + } + if !transactionUnsupported(beginErr) { + return queryResult{}, beginErr + } + } + for _, statement := range statements { + trimmed := trimStatementSQL(statement) + if trimmed == "" { + continue + } + result, err := connection.ExecContext(ctx, trimmed) + if err != nil { + return queryResult{}, err + } + count, _ := result.RowsAffected() + affected += max(count, 0) + } + return emptyQueryResult(affected, started), nil +} + +func transactionUnsupported(err error) bool { + if err == nil { + return false + } + if errors.Is(err, sql.ErrTxDone) { + return false + } + message := strings.ToLower(err.Error()) + return errors.Is(err, driver.ErrSkip) || + strings.Contains(message, "transactions are not supported") || + strings.Contains(message, "transaction is not supported") || + strings.Contains(message, "unsupported transaction") || + strings.Contains(message, "driver: skip fast-path") +} + +func emptyQueryResult(affected int64, started time.Time) queryResult { + return queryResult{ + Columns: []string{}, + ColumnTypes: []string{}, + Rows: [][]any{}, + AffectedRows: affected, + ExecutionTimeMS: time.Since(started).Milliseconds(), + Truncated: false, + } +} + +func (server *server) applySchemaContext(ctx context.Context, connection *sql.Conn, schema string) error { + schema = strings.TrimSpace(schema) + if schema == "" || strings.EqualFold(schema, server.config.Database) { + return nil + } + _, err := connection.ExecContext(ctx, "USE "+quoteHiveIdentifier(schema)) + return err +} + +func queryColumnMetadata(rows *sql.Rows) ([]string, []string, error) { + columns, err := rows.Columns() + if err != nil { + return nil, nil, err + } + types, err := rows.ColumnTypes() + if err != nil { + return nil, nil, err + } + columnTypes := make([]string, len(columns)) + for index := range columns { + if index < len(types) { + columnTypes[index] = strings.ToLower(strings.TrimSpace(types[index].DatabaseTypeName())) + } + } + return columns, columnTypes, nil +} + +func readSQLRows(rows *sql.Rows, columnTypes []string, limit int) ([][]any, bool, error) { + values := make([][]any, 0, min(limit, defaultFetchSize)) + for len(values) < limit { + row, ok, err := nextSQLRow(rows, columnTypes) + if err != nil { + return nil, false, err + } + if !ok { + return values, false, nil + } + values = append(values, row) + } + _, ok, err := nextSQLRow(rows, columnTypes) + return values, ok, err +} + +func nextSQLRow(rows *sql.Rows, columnTypes []string) ([]any, bool, error) { + if !rows.Next() { + if err := rows.Err(); err != nil { + return nil, false, err + } + return nil, false, nil + } + values := make([]any, len(columnTypes)) + targets := make([]any, len(values)) + for index := range values { + targets[index] = &values[index] + } + if err := rows.Scan(targets...); err != nil { + return nil, false, err + } + for index, value := range values { + values[index] = normalizeHiveValue(value, columnTypes[index]) + } + return values, true, nil +} + +func normalizeHiveValue(value any, columnType string) any { + if value == nil { + return nil + } + switch typed := value.(type) { + case []byte: + return bytesToHex(typed) + case time.Time: + return formatHiveJDBCDateTime(typed, columnType) + case fmt.Stringer: + return typed.String() + case string: + return typed + default: + return fmt.Sprint(value) + } +} + +func formatHiveJDBCDateTime(value time.Time, columnType string) string { + if strings.EqualFold(strings.TrimSpace(columnType), "DATE") { + return value.Format("2006-01-02") + } + base := value.Format("2006-01-02 15:04:05") + if value.Nanosecond() == 0 { + return base + ".0" + } + fraction := strings.TrimRight(fmt.Sprintf("%09d", value.Nanosecond()), "0") + return base + "." + fraction +} + +func executeHiveStatement( + ctx context.Context, + connection *sql.Conn, + sqlText string, + fetchSize int, +) (*sql.Rows, int64, bool, error) { + ctx = gohive.WithFetchSize(ctx, fetchSize) + rows, err := connection.QueryContext(ctx, sqlText) + if err == nil { + return rows, 0, true, nil + } + var nonQuery *gohive.NonQueryResult + if errors.As(err, &nonQuery) { + return nil, max(nonQuery.AffectedRows, 0), false, nil + } + return nil, 0, false, err +} + +func bytesToHex(value []byte) string { + const digits = "0123456789abcdef" + result := make([]byte, 2+len(value)*2) + result[0] = '0' + result[1] = 'x' + for index, current := range value { + result[2+index*2] = digits[current>>4] + result[3+index*2] = digits[current&0x0f] + } + return string(result) +} + +func queryContext(timeoutSecs int) (context.Context, context.CancelFunc) { + if timeoutSecs > 0 { + return context.WithTimeout(context.Background(), time.Duration(timeoutSecs)*time.Second) + } + return context.WithCancel(context.Background()) +} + +func (server *server) effectiveFetchSize() int { + if server.config.FetchSize > 0 { + return server.config.FetchSize + } + return defaultFetchSize +} + +func effectiveSchema(options queryOptions) string { + return firstNonEmpty(options.Schema, options.Database) +} + +func trimStatementSQL(sqlText string) string { + return strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(sqlText), ";")) +} + +func quoteHiveIdentifier(value string) string { + return "`" + strings.ReplaceAll(value, "`", "``") + "`" +} diff --git a/agents/drivers/hive-go/query_test.go b/agents/drivers/hive-go/query_test.go new file mode 100644 index 0000000000..c446bfdc75 --- /dev/null +++ b/agents/drivers/hive-go/query_test.go @@ -0,0 +1,557 @@ +package main + +import ( + "context" + "database/sql" + "database/sql/driver" + "encoding/json" + "errors" + "fmt" + "io" + "reflect" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/t8y2/dbx/agents/go-common/gohive" +) + +var scriptedDriverSequence atomic.Uint64 + +type scriptedBehavior struct { + mu sync.Mutex + query func(context.Context, string) (driver.Rows, error) + exec func(context.Context, string) (driver.Result, error) + getSchemas func(context.Context, string) (gohive.MetadataResult, error) + getTables func(context.Context, string, string, []string) (gohive.MetadataResult, error) + getColumns func(context.Context, string, string, string) (gohive.MetadataResult, error) + getTypeInfo func(context.Context) (gohive.MetadataResult, error) + beginErr error + queries []string + executions []string + beginCalls int + closeCalls int +} + +func (behavior *scriptedBehavior) queryContext(ctx context.Context, query string) (driver.Rows, error) { + behavior.mu.Lock() + behavior.queries = append(behavior.queries, query) + operation := behavior.query + behavior.mu.Unlock() + if operation == nil { + return nil, fmt.Errorf("unexpected query: %s", query) + } + return operation(ctx, query) +} + +func (behavior *scriptedBehavior) execContext(ctx context.Context, query string) (driver.Result, error) { + behavior.mu.Lock() + behavior.executions = append(behavior.executions, query) + operation := behavior.exec + behavior.mu.Unlock() + if operation == nil { + return nil, fmt.Errorf("unexpected execution: %s", query) + } + return operation(ctx, query) +} + +func (behavior *scriptedBehavior) snapshot() (queries, executions []string, beginCalls, closeCalls int) { + behavior.mu.Lock() + defer behavior.mu.Unlock() + return append([]string(nil), behavior.queries...), append([]string(nil), behavior.executions...), behavior.beginCalls, behavior.closeCalls +} + +type scriptedDriver struct { + behavior *scriptedBehavior +} + +func (driverValue *scriptedDriver) Open(string) (driver.Conn, error) { + return &scriptedConnection{behavior: driverValue.behavior}, nil +} + +type scriptedConnection struct { + behavior *scriptedBehavior +} + +func (connection *scriptedConnection) Prepare(string) (driver.Stmt, error) { + return nil, errors.New("prepared statements are not supported") +} + +func (connection *scriptedConnection) Close() error { + connection.behavior.mu.Lock() + connection.behavior.closeCalls++ + connection.behavior.mu.Unlock() + return nil +} + +func (connection *scriptedConnection) Begin() (driver.Tx, error) { + connection.behavior.mu.Lock() + connection.behavior.beginCalls++ + connection.behavior.mu.Unlock() + if connection.behavior.beginErr != nil { + return nil, connection.behavior.beginErr + } + return nil, driver.ErrSkip +} + +func (connection *scriptedConnection) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) { + connection.behavior.mu.Lock() + connection.behavior.beginCalls++ + connection.behavior.mu.Unlock() + if connection.behavior.beginErr != nil { + return nil, connection.behavior.beginErr + } + return nil, driver.ErrSkip +} + +func (connection *scriptedConnection) Ping(context.Context) error { + return nil +} + +func (connection *scriptedConnection) QueryContext( + ctx context.Context, + query string, + _ []driver.NamedValue, +) (driver.Rows, error) { + return connection.behavior.queryContext(ctx, query) +} + +func (connection *scriptedConnection) ExecContext( + ctx context.Context, + query string, + _ []driver.NamedValue, +) (driver.Result, error) { + return connection.behavior.execContext(ctx, query) +} + +func (connection *scriptedConnection) GetHiveSchemas(ctx context.Context, pattern string) (gohive.MetadataResult, error) { + if connection.behavior.getSchemas == nil { + return gohive.MetadataResult{}, errors.New("GetSchemas unavailable") + } + return connection.behavior.getSchemas(ctx, pattern) +} + +func (connection *scriptedConnection) GetHiveTables(ctx context.Context, schema, table string, tableTypes []string) (gohive.MetadataResult, error) { + if connection.behavior.getTables == nil { + return gohive.MetadataResult{}, errors.New("GetTables unavailable") + } + return connection.behavior.getTables(ctx, schema, table, tableTypes) +} + +func (connection *scriptedConnection) GetHiveColumns(ctx context.Context, schema, table, column string) (gohive.MetadataResult, error) { + if connection.behavior.getColumns == nil { + return gohive.MetadataResult{}, errors.New("GetColumns unavailable") + } + return connection.behavior.getColumns(ctx, schema, table, column) +} + +func (connection *scriptedConnection) GetHiveTypeInfo(ctx context.Context) (gohive.MetadataResult, error) { + if connection.behavior.getTypeInfo == nil { + return gohive.MetadataResult{}, errors.New("GetTypeInfo unavailable") + } + return connection.behavior.getTypeInfo(ctx) +} + +type scriptedRows struct { + ctx context.Context + columns []string + types []string + values [][]driver.Value + blockAfter int + blocked chan struct{} + blockOnce sync.Once + + mu sync.Mutex + index int + closed bool +} + +func newScriptedRows(ctx context.Context, columns, types []string, values [][]driver.Value) *scriptedRows { + return &scriptedRows{ + ctx: ctx, + columns: columns, + types: types, + values: values, + blockAfter: -1, + } +} + +func (rows *scriptedRows) Columns() []string { + return append([]string(nil), rows.columns...) +} + +func (rows *scriptedRows) Close() error { + rows.mu.Lock() + rows.closed = true + rows.mu.Unlock() + return nil +} + +func (rows *scriptedRows) Next(destination []driver.Value) error { + rows.mu.Lock() + if rows.closed { + rows.mu.Unlock() + return io.EOF + } + index := rows.index + if rows.blockAfter >= 0 && index >= rows.blockAfter { + blocked := rows.blocked + ctx := rows.ctx + rows.mu.Unlock() + if blocked != nil { + rows.blockOnce.Do(func() { close(blocked) }) + } + <-ctx.Done() + return ctx.Err() + } + if index >= len(rows.values) { + rows.mu.Unlock() + return io.EOF + } + current := rows.values[index] + rows.index++ + rows.mu.Unlock() + copy(destination, current) + return nil +} + +func (rows *scriptedRows) ColumnTypeDatabaseTypeName(index int) string { + if index < 0 || index >= len(rows.types) { + return "" + } + return rows.types[index] +} + +func (rows *scriptedRows) isClosed() bool { + rows.mu.Lock() + defer rows.mu.Unlock() + return rows.closed +} + +func newScriptedServer(t *testing.T, behavior *scriptedBehavior) *server { + t.Helper() + driverName := fmt.Sprintf("dbx-hive-scripted-%d", scriptedDriverSequence.Add(1)) + sql.Register(driverName, &scriptedDriver{behavior: behavior}) + database, err := sql.Open(driverName, "") + if err != nil { + t.Fatal(err) + } + connection, err := database.Conn(context.Background()) + if err != nil { + database.Close() + t.Fatal(err) + } + server := &server{ + config: connectionConfig{Database: defaultHiveDatabase, ConnectTimeout: time.Second}, + database: database, + connection: connection, + querySessions: map[string]*querySession{}, + } + t.Cleanup(func() { _ = server.disconnect() }) + return server +} + +func rawParams(values map[string]any) map[string]json.RawMessage { + result := make(map[string]json.RawMessage, len(values)) + for key, value := range values { + encoded, err := json.Marshal(value) + if err != nil { + panic(err) + } + result[key] = encoded + } + return result +} + +func TestExecuteQueryUsesHiveServerResultSetSignal(t *testing.T) { + behavior := &scriptedBehavior{} + behavior.query = func(ctx context.Context, query string) (driver.Rows, error) { + switch { + case strings.HasPrefix(query, "WITH source AS"): + return nil, &gohive.NonQueryResult{AffectedRows: 4} + case strings.HasPrefix(query, "SET "): + return newScriptedRows(ctx, []string{"set"}, []string{"STRING"}, [][]driver.Value{{"hive.exec.dynamic.partition=true"}}), nil + default: + return nil, fmt.Errorf("unexpected SQL: %s", query) + } + } + server := newScriptedServer(t, behavior) + + insertResult, err := server.executeQuery(queryOptions{ + SQL: "WITH source AS (SELECT 1) INSERT INTO target SELECT * FROM source", + }) + if err != nil { + t.Fatal(err) + } + if insertResult.AffectedRows != 4 || len(insertResult.Columns) != 0 { + t.Fatalf("unexpected non-query result: %#v", insertResult) + } + + setResult, err := server.executeQuery(queryOptions{SQL: "SET hive.exec.dynamic.partition"}) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(setResult.Rows, [][]any{{"hive.exec.dynamic.partition=true"}}) { + t.Fatalf("unexpected SET result: %#v", setResult.Rows) + } + queries, executions, _, _ := behavior.snapshot() + if len(queries) != 2 || len(executions) != 0 { + t.Fatalf("statements must use HS2 result-set signaling, queries=%v executions=%v", queries, executions) + } +} + +func TestPagedQueryTruncatesAndPreservesLegacyJDBCValueSemantics(t *testing.T) { + largeValue := strings.Repeat("x", 256*1024) + createdAt := time.Date(2026, time.August, 11, 10, 11, 12, 345000000, time.UTC) + var sourceRows *scriptedRows + behavior := &scriptedBehavior{} + behavior.query = func(ctx context.Context, _ string) (driver.Rows, error) { + sourceRows = newScriptedRows( + ctx, + []string{"id", "enabled", "payload", "created_at", "complex_value"}, + []string{"BIGINT", "BOOLEAN", "BINARY", "TIMESTAMP", "ARRAY"}, + [][]driver.Value{ + {int64(1), true, []byte{0x00, 0xff}, createdAt, largeValue}, + {int64(2), false, []byte{0x10}, createdAt, "[1,2]"}, + {int64(3), true, []byte{}, createdAt, "map('a',1)"}, + {int64(4), true, []byte{0x01}, createdAt, "extra"}, + }, + ) + return sourceRows, nil + } + server := newScriptedServer(t, behavior) + + first, err := server.executeQueryPage(queryOptions{SQL: "SELECT * FROM values", MaxRows: 3, FetchSize: 2}, 2) + if err != nil { + t.Fatal(err) + } + if !first.HasMore || first.SessionID == nil || len(first.Rows) != 2 { + t.Fatalf("unexpected first page: %#v", first) + } + if !reflect.DeepEqual(first.ColumnTypes, []string{"bigint", "boolean", "binary", "timestamp", "array"}) { + t.Fatalf("unexpected column types: %#v", first.ColumnTypes) + } + if first.Rows[0][0] != "1" || first.Rows[0][1] != "true" || first.Rows[0][2] != "0x00ff" { + t.Fatalf("primitive value types changed: %#v", first.Rows[0]) + } + if first.Rows[0][3] != "2026-08-11 10:11:12.345" || first.Rows[0][4] != largeValue { + t.Fatalf("timestamp or large value changed: %#v", first.Rows[0]) + } + + second, err := server.fetchQueryPage(*first.SessionID, 2) + if err != nil { + t.Fatal(err) + } + if second.HasMore || second.SessionID != nil || !second.Truncated || len(second.Rows) != 1 { + t.Fatalf("unexpected final page: %#v", second) + } + if len(server.querySessions) != 0 || sourceRows == nil || !sourceRows.isClosed() { + t.Fatalf("query session was not closed: sessions=%d rows=%#v", len(server.querySessions), sourceRows) + } +} + +func TestCancelPagedFetchQuarantinesWithoutReplayingSQL(t *testing.T) { + blocked := make(chan struct{}) + behavior := &scriptedBehavior{} + behavior.query = func(ctx context.Context, _ string) (driver.Rows, error) { + rows := newScriptedRows(ctx, []string{"id"}, []string{"BIGINT"}, [][]driver.Value{{int64(1)}, {int64(2)}, {int64(3)}}) + rows.blockAfter = 2 + rows.blocked = blocked + return rows, nil + } + server := newScriptedServer(t, behavior) + first, err := server.executeQueryPage(queryOptions{SQL: "SELECT id FROM slow_table", MaxRows: 10}, 1) + if err != nil { + t.Fatal(err) + } + if first.SessionID == nil { + t.Fatal("expected a paged query session") + } + + result := make(chan error, 1) + go func() { + _, fetchErr := server.fetchQueryPage(*first.SessionID, 1) + result <- fetchErr + }() + select { + case <-blocked: + case <-time.After(2 * time.Second): + t.Fatal("fetch did not reach the blocking row") + } + server.cancelActiveQuery() + fetchErr := <-result + if !errors.Is(fetchErr, context.Canceled) { + t.Fatalf("expected cancellation, got %v", fetchErr) + } + rpcErr := classifyRPCError("fetch_query_page", "session-a", fetchErr) + if rpcErr.Data.Category != "canceled" || rpcErr.Data.SessionDisposition != "quarantine" { + t.Fatalf("unexpected cancellation classification: %#v", rpcErr) + } + queries, _, _, _ := behavior.snapshot() + if len(queries) != 1 { + t.Fatalf("SQL must never be replayed after cancellation: %v", queries) + } + if len(server.querySessions) != 0 { + t.Fatalf("canceled query session was retained: %#v", server.querySessions) + } +} + +func TestPagedFetchHonorsOriginalStatementTimeout(t *testing.T) { + blocked := make(chan struct{}) + behavior := &scriptedBehavior{} + behavior.query = func(ctx context.Context, _ string) (driver.Rows, error) { + rows := newScriptedRows(ctx, []string{"id"}, []string{"BIGINT"}, [][]driver.Value{{int64(1)}, {int64(2)}, {int64(3)}}) + rows.blockAfter = 2 + rows.blocked = blocked + return rows, nil + } + server := newScriptedServer(t, behavior) + first, err := server.executeQueryPage(queryOptions{SQL: "SELECT id FROM slow_table", MaxRows: 10, TimeoutSecs: 1}, 1) + if err != nil { + t.Fatal(err) + } + if first.SessionID == nil { + t.Fatal("expected a paged query session") + } + + started := time.Now() + _, fetchErr := server.fetchQueryPage(*first.SessionID, 1) + if !errors.Is(fetchErr, context.DeadlineExceeded) { + t.Fatalf("expected deadline exceeded, got %v", fetchErr) + } + if elapsed := time.Since(started); elapsed > 3*time.Second { + t.Fatalf("statement timeout was not enforced promptly: %s", elapsed) + } + if len(server.querySessions) != 0 { + t.Fatalf("timed out query session was retained: %#v", server.querySessions) + } +} + +func TestRuntimeSessionsRemainIsolated(t *testing.T) { + behaviorA := &scriptedBehavior{query: func(ctx context.Context, _ string) (driver.Rows, error) { + return newScriptedRows(ctx, []string{"owner"}, []string{"STRING"}, [][]driver.Value{{"a"}}), nil + }} + behaviorB := &scriptedBehavior{query: func(ctx context.Context, _ string) (driver.Rows, error) { + return newScriptedRows(ctx, []string{"owner"}, []string{"STRING"}, [][]driver.Value{{"b"}}), nil + }} + serverA := newScriptedServer(t, behaviorA) + serverB := newScriptedServer(t, behaviorB) + runtimeServer := newRuntimeServer() + runtimeServer.sessions["a"] = &agentSession{server: serverA} + runtimeServer.sessions["b"] = &agentSession{server: serverB} + + resultA, _, err := runtimeServer.dispatch("execute_query", rawParams(map[string]any{ + "agentSessionId": "a", + "sql": "SELECT owner", + })) + if err != nil { + t.Fatal(err) + } + resultB, _, err := runtimeServer.dispatch("execute_query", rawParams(map[string]any{ + "agentSessionId": "b", + "sql": "SELECT owner", + })) + if err != nil { + t.Fatal(err) + } + if resultA.(queryResult).Rows[0][0] != "a" || resultB.(queryResult).Rows[0][0] != "b" { + t.Fatalf("session results crossed: a=%#v b=%#v", resultA, resultB) + } + if err := runtimeServer.closeSession("a"); err != nil { + t.Fatal(err) + } + if _, err := runtimeServer.session("a"); err == nil { + t.Fatal("closed session a is still registered") + } + if _, err := runtimeServer.session("b"); err != nil { + t.Fatalf("closing session a affected session b: %v", err) + } +} + +func TestTransactionFallbackDoesNotReplayFailedStatements(t *testing.T) { + behavior := &scriptedBehavior{} + behavior.exec = func(_ context.Context, query string) (driver.Result, error) { + if strings.Contains(query, "second") { + return nil, io.EOF + } + return driver.RowsAffected(1), nil + } + server := newScriptedServer(t, behavior) + _, err := server.executeStatements(rawParams(map[string]any{ + "statements": []string{"INSERT first", "INSERT second", "INSERT third"}, + }), true) + if !errors.Is(err, io.EOF) { + t.Fatalf("expected connection failure, got %v", err) + } + _, executions, beginCalls, _ := behavior.snapshot() + if !reflect.DeepEqual(executions, []string{"INSERT first", "INSERT second"}) { + t.Fatalf("failed transaction was replayed or continued: %v", executions) + } + if beginCalls == 0 { + t.Fatal("transaction capability was not attempted before fallback") + } + rpcErr := classifyRPCError("execute_transaction", "session-a", err) + if rpcErr.Data.Category != "connection" || rpcErr.Data.SessionDisposition != "quarantine" { + t.Fatalf("unexpected connection failure classification: %#v", rpcErr) + } +} + +func TestTransactionDoesNotFallbackAfterOperationalBeginFailure(t *testing.T) { + behavior := &scriptedBehavior{ + beginErr: errors.New("connection reset while beginning transaction"), + exec: func(context.Context, string) (driver.Result, error) { + return driver.RowsAffected(1), nil + }, + } + server := newScriptedServer(t, behavior) + defer server.disconnect() + params := map[string]json.RawMessage{ + "statements": json.RawMessage(`["INSERT INTO sample VALUES (1)"]`), + } + if _, err := server.executeStatements(params, true); err == nil || !strings.Contains(err.Error(), "connection reset") { + t.Fatalf("expected begin failure, got %v", err) + } + _, executions, _, _ := behavior.snapshot() + if len(executions) != 0 { + t.Fatalf("statements must not execute after begin failure: %#v", executions) + } +} + +func TestExpireIdleQuerySessions(t *testing.T) { + behavior := &scriptedBehavior{query: func(ctx context.Context, _ string) (driver.Rows, error) { + return newScriptedRows(ctx, []string{"id"}, []string{"BIGINT"}, [][]driver.Value{{int64(1)}, {int64(2)}}), nil + }} + server := newScriptedServer(t, behavior) + first, err := server.executeQueryPage(queryOptions{SQL: "SELECT id", MaxRows: 10}, 1) + if err != nil { + t.Fatal(err) + } + if first.SessionID == nil { + t.Fatal("expected a query session") + } + server.querySessions[*first.SessionID].lastAccessed = time.Now().Add(-querySessionIdleTime - time.Second) + if expired := server.expireIdleQuerySessions(time.Now()); expired != 1 { + t.Fatalf("unexpected expired session count: %d", expired) + } + if len(server.querySessions) != 0 { + t.Fatalf("idle session was retained: %#v", server.querySessions) + } +} + +func TestStructuredHiveErrorIncludesServerDiagnostics(t *testing.T) { + err := &gohive.Error{ + Err: errors.New("compile failed"), + Message: "SemanticException", + ErrorCode: 40000, + SQLState: "42000", + } + rpcErr := classifyRPCError("execute_query", "session-a", err) + if rpcErr.Data.Category != "sql" || rpcErr.Data.SQLState != "42000" || rpcErr.Data.VendorCode != 40000 { + t.Fatalf("unexpected Hive diagnostics: %#v", rpcErr) + } + if rpcErr.Data.SessionDisposition != "keep" || rpcErr.Data.OperationOutcome != "unknown" { + t.Fatalf("unexpected SQL failure recovery hints: %#v", rpcErr) + } +} diff --git a/agents/drivers/hive-go/zookeeper_protocol.go b/agents/drivers/hive-go/zookeeper_protocol.go new file mode 100644 index 0000000000..66982d67af --- /dev/null +++ b/agents/drivers/hive-go/zookeeper_protocol.go @@ -0,0 +1,568 @@ +package main + +import ( + "bytes" + "crypto/tls" + "encoding/binary" + "errors" + "fmt" + "io" + "math" + "math/rand/v2" + "net" + "strings" + "sync" + "time" + + "github.com/beltran/gosasl" + "github.com/go-zookeeper/zk" +) + +const ( + zooKeeperProtocolVersion = int32(0) + zooKeeperOpGetData = int32(4) + zooKeeperOpGetChildren2 = int32(12) + zooKeeperOpClose = int32(-11) + zooKeeperOpSetAuth = int32(100) + zooKeeperOpSASL = int32(102) + zooKeeperMaxFrameSize = 16 << 20 + zooKeeperMaxSASLRounds = 8 +) + +var errZooKeeperSessionClosedRequiresSASL = errors.New("ZooKeeper session closed because SASL authentication is required") + +type zooKeeperSASLClient interface { + Start() ([]byte, error) + Step([]byte) ([]byte, error) + Complete() bool + Dispose() +} + +var newZooKeeperSASLClient = func(host string, config connectionConfig) (zooKeeperSASLClient, error) { + service, options := zooKeeperGSSAPIOptions(config) + mechanism, err := gosasl.NewGSSAPIMechanismWithOptions(service, options) + if err != nil { + return nil, err + } + return gosasl.NewSaslClient(host, mechanism), nil +} + +var dialZooKeeperConnection = func(address string, timeout time.Duration, tlsConfig *tls.Config) (net.Conn, error) { + dialer := &net.Dialer{Timeout: timeout} + if tlsConfig == nil { + return dialer.Dial("tcp", address) + } + config := tlsConfig.Clone() + if config.ServerName == "" { + host, _, err := net.SplitHostPort(address) + if err != nil { + return nil, fmt.Errorf("parse ZooKeeper TLS address %q: %w", address, err) + } + config.ServerName = host + } + return tls.DialWithDialer(dialer, "tcp", address, config) +} + +var shuffleZooKeeperServers = func(servers []string) { + rand.Shuffle(len(servers), func(first, second int) { + servers[first], servers[second] = servers[second], servers[first] + }) +} + +func zooKeeperGSSAPIOptions(config connectionConfig) (string, gosasl.GSSAPIOptions) { + service := firstNonEmpty(config.ZooKeeperKerberos.Service, "zookeeper") + options := gssapiOptionsFromKerberos(config.Kerberos) + options.QOP = "auth" + options.AuthorizationID = "" + options.ServiceHost = "" + options.CanonicalizeHost = config.ZooKeeperKerberos.CanonicalHostname + options.ServerName = config.ZooKeeperKerberos.ServerPrincipal + if options.ServerName == "" && config.ZooKeeperKerberos.Realm != "" { + options.ServerName = service + "/_HOST@" + config.ZooKeeperKerberos.Realm + } + return service, options +} + +func connectKerberosZooKeeper( + servers []string, + timeout time.Duration, + tlsConfig *tls.Config, + config connectionConfig, +) (zooKeeperClient, <-chan zk.Event, error) { + if len(servers) == 0 { + return nil, nil, errors.New("ZooKeeper server list is empty") + } + if !config.Kerberos.Enabled { + return nil, nil, errors.New("ZooKeeper Kerberos SASL requires Hive Kerberos credentials") + } + ordered := append([]string(nil), servers...) + shuffleZooKeeperServers(ordered) + var failures []string + for _, address := range ordered { + host, _, err := net.SplitHostPort(address) + if err != nil { + failures = append(failures, fmt.Sprintf("%s: %v", address, err)) + continue + } + connection, err := dialZooKeeperConnection(address, timeout, tlsConfig) + if err != nil { + failures = append(failures, fmt.Sprintf("%s: %v", address, err)) + continue + } + client, err := newProtocolZooKeeperClient(connection, timeout) + if err == nil { + var saslClient zooKeeperSASLClient + saslClient, err = newZooKeeperSASLClient(host, config) + if err == nil { + err = client.authenticateSASL(saslClient) + } + } + if err != nil { + connection.Close() + failures = append(failures, fmt.Sprintf("%s: %v", address, err)) + continue + } + events := make(chan zk.Event, 1) + events <- zk.Event{State: zk.StateHasSession, Server: address} + close(events) + return client, events, nil + } + return nil, nil, fmt.Errorf("connect and authenticate to ZooKeeper: %s", strings.Join(failures, "; ")) +} + +type protocolZooKeeperClient struct { + connection net.Conn + timeout time.Duration + xid int32 + mutex sync.Mutex + closed bool +} + +func newProtocolZooKeeperClient(connection net.Conn, timeout time.Duration) (*protocolZooKeeperClient, error) { + if connection == nil { + return nil, errors.New("ZooKeeper connection is nil") + } + if timeout <= 0 { + timeout = defaultConnectTimeout + } + client := &protocolZooKeeperClient{connection: connection, timeout: timeout} + request := &zooKeeperEncoder{} + request.int32(zooKeeperProtocolVersion) + request.int64(0) + request.int32(zooKeeperTimeoutMillis(timeout)) + request.int64(0) + request.bytes(make([]byte, 16)) + if err := client.writeFrame(request.data()); err != nil { + return nil, fmt.Errorf("send ZooKeeper connect request: %w", err) + } + response, err := client.readFrame() + if err != nil { + return nil, fmt.Errorf("read ZooKeeper connect response: %w", err) + } + decoder := newZooKeeperDecoder(response) + if _, err := decoder.int32(); err != nil { + return nil, fmt.Errorf("decode ZooKeeper protocol version: %w", err) + } + if _, err := decoder.int32(); err != nil { + return nil, fmt.Errorf("decode ZooKeeper session timeout: %w", err) + } + sessionID, err := decoder.int64() + if err != nil { + return nil, fmt.Errorf("decode ZooKeeper session ID: %w", err) + } + if _, err := decoder.bytes(); err != nil { + return nil, fmt.Errorf("decode ZooKeeper session password: %w", err) + } + if sessionID == 0 { + return nil, zk.ErrSessionExpired + } + return client, nil +} + +func zooKeeperTimeoutMillis(timeout time.Duration) int32 { + milliseconds := timeout.Milliseconds() + if milliseconds < 1 { + return 1 + } + if milliseconds > math.MaxInt32 { + return math.MaxInt32 + } + return int32(milliseconds) +} + +func (client *protocolZooKeeperClient) authenticateSASL(saslClient zooKeeperSASLClient) error { + if saslClient == nil { + return errors.New("ZooKeeper SASL client is nil") + } + defer saslClient.Dispose() + token, err := saslClient.Start() + if err != nil { + return fmt.Errorf("start ZooKeeper GSSAPI negotiation: %w", err) + } + for round := 0; round < zooKeeperMaxSASLRounds; round++ { + response, requestErr := client.request(zooKeeperOpSASL, func(encoder *zooKeeperEncoder) { + if token == nil { + encoder.bytes([]byte{}) + return + } + encoder.bytes(token) + }) + if requestErr != nil { + return fmt.Errorf("ZooKeeper SASL round %d: %w", round+1, requestErr) + } + decoder := newZooKeeperDecoder(response) + challenge, decodeErr := decoder.bytes() + if decodeErr != nil { + return fmt.Errorf("decode ZooKeeper SASL round %d: %w", round+1, decodeErr) + } + if saslClient.Complete() { + if len(challenge) != 0 { + return errors.New("ZooKeeper sent an unexpected token after GSSAPI completion") + } + return nil + } + token, err = saslClient.Step(challenge) + if err != nil { + return fmt.Errorf("continue ZooKeeper GSSAPI negotiation at round %d: %w", round+1, err) + } + } + return fmt.Errorf("ZooKeeper GSSAPI negotiation exceeded %d rounds", zooKeeperMaxSASLRounds) +} + +func (client *protocolZooKeeperClient) AddAuth(scheme string, auth []byte) error { + _, err := client.request(zooKeeperOpSetAuth, func(encoder *zooKeeperEncoder) { + encoder.int32(0) + encoder.string(scheme) + encoder.bytes(auth) + }) + return err +} + +func (client *protocolZooKeeperClient) Children(path string) ([]string, *zk.Stat, error) { + response, err := client.request(zooKeeperOpGetChildren2, func(encoder *zooKeeperEncoder) { + encoder.string(path) + encoder.boolean(false) + }) + if err != nil { + return nil, nil, err + } + decoder := newZooKeeperDecoder(response) + children, err := decoder.strings() + if err != nil { + return nil, nil, err + } + stat, err := decoder.stat() + if err != nil { + return nil, nil, err + } + return children, stat, nil +} + +func (client *protocolZooKeeperClient) Get(path string) ([]byte, *zk.Stat, error) { + response, err := client.request(zooKeeperOpGetData, func(encoder *zooKeeperEncoder) { + encoder.string(path) + encoder.boolean(false) + }) + if err != nil { + return nil, nil, err + } + decoder := newZooKeeperDecoder(response) + data, err := decoder.bytes() + if err != nil { + return nil, nil, err + } + stat, err := decoder.stat() + if err != nil { + return nil, nil, err + } + return data, stat, nil +} + +func (client *protocolZooKeeperClient) Close() { + client.mutex.Lock() + defer client.mutex.Unlock() + if client.closed { + return + } + client.closed = true + _ = client.connection.SetDeadline(time.Now().Add(client.timeout)) + client.xid++ + request := &zooKeeperEncoder{} + request.int32(client.xid) + request.int32(zooKeeperOpClose) + _ = client.writeFrame(request.data()) + _ = client.connection.Close() +} + +func (client *protocolZooKeeperClient) request(opcode int32, encodeBody func(*zooKeeperEncoder)) ([]byte, error) { + client.mutex.Lock() + defer client.mutex.Unlock() + if client.closed { + return nil, zk.ErrConnectionClosed + } + client.xid++ + request := &zooKeeperEncoder{} + request.int32(client.xid) + request.int32(opcode) + if encodeBody != nil { + encodeBody(request) + } + if err := client.writeFrame(request.data()); err != nil { + return nil, err + } + response, err := client.readFrame() + if err != nil { + return nil, err + } + decoder := newZooKeeperDecoder(response) + xid, err := decoder.int32() + if err != nil { + return nil, err + } + if xid != client.xid { + return nil, fmt.Errorf("ZooKeeper response XID %d does not match request XID %d", xid, client.xid) + } + if _, err := decoder.int64(); err != nil { + return nil, err + } + code, err := decoder.int32() + if err != nil { + return nil, err + } + if err := zooKeeperError(code); err != nil { + return nil, err + } + return decoder.remaining(), nil +} + +func (client *protocolZooKeeperClient) writeFrame(payload []byte) error { + if len(payload) > zooKeeperMaxFrameSize { + return fmt.Errorf("ZooKeeper request frame is %d bytes, maximum is %d", len(payload), zooKeeperMaxFrameSize) + } + if err := client.connection.SetWriteDeadline(time.Now().Add(client.timeout)); err != nil { + return err + } + header := make([]byte, 4) + binary.BigEndian.PutUint32(header, uint32(len(payload))) + if err := writeAll(client.connection, header); err != nil { + return err + } + return writeAll(client.connection, payload) +} + +func (client *protocolZooKeeperClient) readFrame() ([]byte, error) { + if err := client.connection.SetReadDeadline(time.Now().Add(client.timeout)); err != nil { + return nil, err + } + header := make([]byte, 4) + if _, err := io.ReadFull(client.connection, header); err != nil { + return nil, err + } + length := int(binary.BigEndian.Uint32(header)) + if length < 0 || length > zooKeeperMaxFrameSize { + return nil, fmt.Errorf("ZooKeeper response frame is %d bytes, maximum is %d", length, zooKeeperMaxFrameSize) + } + payload := make([]byte, length) + if _, err := io.ReadFull(client.connection, payload); err != nil { + return nil, err + } + return payload, nil +} + +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 +} + +func zooKeeperError(code int32) error { + switch code { + case 0: + return nil + case -4: + return zk.ErrConnectionClosed + case -101: + return zk.ErrNoNode + case -102: + return zk.ErrNoAuth + case -112: + return zk.ErrSessionExpired + case -115: + return zk.ErrAuthFailed + case -124: + return errZooKeeperSessionClosedRequiresSASL + default: + return fmt.Errorf("ZooKeeper request failed with error code %d", code) + } +} + +type zooKeeperEncoder struct { + buffer bytes.Buffer +} + +func (encoder *zooKeeperEncoder) int32(value int32) { + var data [4]byte + binary.BigEndian.PutUint32(data[:], uint32(value)) + encoder.buffer.Write(data[:]) +} + +func (encoder *zooKeeperEncoder) int64(value int64) { + var data [8]byte + binary.BigEndian.PutUint64(data[:], uint64(value)) + encoder.buffer.Write(data[:]) +} + +func (encoder *zooKeeperEncoder) boolean(value bool) { + if value { + encoder.buffer.WriteByte(1) + return + } + encoder.buffer.WriteByte(0) +} + +func (encoder *zooKeeperEncoder) string(value string) { + encoder.bytes([]byte(value)) +} + +func (encoder *zooKeeperEncoder) bytes(value []byte) { + if value == nil { + encoder.int32(-1) + return + } + encoder.int32(int32(len(value))) + encoder.buffer.Write(value) +} + +func (encoder *zooKeeperEncoder) data() []byte { + return encoder.buffer.Bytes() +} + +type zooKeeperDecoder struct { + data []byte + offset int +} + +func newZooKeeperDecoder(data []byte) *zooKeeperDecoder { + return &zooKeeperDecoder{data: data} +} + +func (decoder *zooKeeperDecoder) take(length int) ([]byte, error) { + if length < 0 || decoder.offset > len(decoder.data)-length { + return nil, io.ErrUnexpectedEOF + } + value := decoder.data[decoder.offset : decoder.offset+length] + decoder.offset += length + return value, nil +} + +func (decoder *zooKeeperDecoder) int32() (int32, error) { + value, err := decoder.take(4) + if err != nil { + return 0, err + } + return int32(binary.BigEndian.Uint32(value)), nil +} + +func (decoder *zooKeeperDecoder) int64() (int64, error) { + value, err := decoder.take(8) + if err != nil { + return 0, err + } + return int64(binary.BigEndian.Uint64(value)), nil +} + +func (decoder *zooKeeperDecoder) bytes() ([]byte, error) { + length, err := decoder.int32() + if err != nil { + return nil, err + } + if length == -1 { + return nil, nil + } + if length < -1 { + return nil, fmt.Errorf("invalid ZooKeeper buffer length %d", length) + } + value, err := decoder.take(int(length)) + if err != nil { + return nil, err + } + return append([]byte(nil), value...), nil +} + +func (decoder *zooKeeperDecoder) string() (string, error) { + value, err := decoder.bytes() + return string(value), err +} + +func (decoder *zooKeeperDecoder) strings() ([]string, error) { + length, err := decoder.int32() + if err != nil { + return nil, err + } + if length == -1 { + return nil, nil + } + if length < -1 || length > zooKeeperMaxFrameSize/4 { + return nil, fmt.Errorf("invalid ZooKeeper string vector length %d", length) + } + values := make([]string, 0, length) + for index := int32(0); index < length; index++ { + value, valueErr := decoder.string() + if valueErr != nil { + return nil, valueErr + } + values = append(values, value) + } + return values, nil +} + +func (decoder *zooKeeperDecoder) stat() (*zk.Stat, error) { + stat := &zk.Stat{} + var err error + if stat.Czxid, err = decoder.int64(); err != nil { + return nil, err + } + if stat.Mzxid, err = decoder.int64(); err != nil { + return nil, err + } + if stat.Ctime, err = decoder.int64(); err != nil { + return nil, err + } + if stat.Mtime, err = decoder.int64(); err != nil { + return nil, err + } + if stat.Version, err = decoder.int32(); err != nil { + return nil, err + } + if stat.Cversion, err = decoder.int32(); err != nil { + return nil, err + } + if stat.Aversion, err = decoder.int32(); err != nil { + return nil, err + } + if stat.EphemeralOwner, err = decoder.int64(); err != nil { + return nil, err + } + if stat.DataLength, err = decoder.int32(); err != nil { + return nil, err + } + if stat.NumChildren, err = decoder.int32(); err != nil { + return nil, err + } + if stat.Pzxid, err = decoder.int64(); err != nil { + return nil, err + } + return stat, nil +} + +func (decoder *zooKeeperDecoder) remaining() []byte { + return decoder.data[decoder.offset:] +} diff --git a/agents/drivers/hive-go/zookeeper_protocol_test.go b/agents/drivers/hive-go/zookeeper_protocol_test.go new file mode 100644 index 0000000000..e228a0d98c --- /dev/null +++ b/agents/drivers/hive-go/zookeeper_protocol_test.go @@ -0,0 +1,562 @@ +package main + +import ( + "bytes" + "crypto/tls" + "errors" + "fmt" + "io" + "log" + "net" + "os" + "path/filepath" + "testing" + "time" + + "github.com/go-zookeeper/zk" + gsskrb5 "github.com/golang-auth/go-gssapi/v2/krb5" + "github.com/jcmturner/krb5test" +) + +type scriptedZooKeeperSASLClient struct { + complete bool + disposed bool +} + +func (client *scriptedZooKeeperSASLClient) Start() ([]byte, error) { + return []byte("client-initial"), nil +} + +func (client *scriptedZooKeeperSASLClient) Step(challenge []byte) ([]byte, error) { + if string(challenge) != "server-challenge" { + return nil, fmt.Errorf("unexpected challenge %q", challenge) + } + client.complete = true + return []byte("client-final"), nil +} + +func (client *scriptedZooKeeperSASLClient) Complete() bool { return client.complete } +func (client *scriptedZooKeeperSASLClient) Dispose() { client.disposed = true } + +func TestProtocolZooKeeperClientAuthenticatesAndReadsDiscoveryData(t *testing.T) { + clientConnection, serverConnection := net.Pipe() + serverErrors := make(chan error, 1) + go func() { + defer serverConnection.Close() + serverErrors <- serveZooKeeperProtocolTest(serverConnection) + }() + + client, err := newProtocolZooKeeperClient(clientConnection, time.Second) + if err != nil { + t.Fatal(err) + } + saslClient := &scriptedZooKeeperSASLClient{} + if err := client.authenticateSASL(saslClient); err != nil { + t.Fatal(err) + } + if !saslClient.disposed { + t.Fatal("SASL credentials were not disposed") + } + if err := client.AddAuth("digest", []byte("user:password")); err != nil { + t.Fatal(err) + } + children, stat, err := client.Children("/hiveserver2") + if err != nil { + t.Fatal(err) + } + if len(children) != 1 || children[0] != "server-1" || stat.NumChildren != 1 { + t.Fatalf("children=%#v stat=%#v", children, stat) + } + data, stat, err := client.Get("/hiveserver2/server-1") + if err != nil { + t.Fatal(err) + } + if string(data) != "serverUri=hs2.example.com:10000" || stat.DataLength != int32(len(data)) { + t.Fatalf("data=%q stat=%#v", data, stat) + } + client.Close() + if err := <-serverErrors; err != nil { + t.Fatal(err) + } +} + +func TestProtocolZooKeeperClientMapsRequiredSASLError(t *testing.T) { + clientConnection, serverConnection := net.Pipe() + serverErrors := make(chan error, 1) + go func() { + defer serverConnection.Close() + if err := acceptZooKeeperSession(serverConnection); err != nil { + serverErrors <- err + return + } + request, err := readZooKeeperTestFrame(serverConnection) + if err != nil { + serverErrors <- err + return + } + decoder := newZooKeeperDecoder(request) + xid, _ := decoder.int32() + opcode, _ := decoder.int32() + if opcode != zooKeeperOpSASL { + serverErrors <- fmt.Errorf("opcode=%d", opcode) + return + } + serverErrors <- writeZooKeeperTestResponse(serverConnection, xid, -124, nil) + }() + + client, err := newProtocolZooKeeperClient(clientConnection, time.Second) + if err != nil { + t.Fatal(err) + } + saslClient := &scriptedZooKeeperSASLClient{} + err = client.authenticateSASL(saslClient) + client.Close() + if !errors.Is(err, errZooKeeperSessionClosedRequiresSASL) { + t.Fatalf("unexpected error %v", err) + } + if serverErr := <-serverErrors; serverErr != nil { + t.Fatal(serverErr) + } +} + +func TestZooKeeperGSSAPIOptionsDoNotReuseHiveServerIdentity(t *testing.T) { + config := connectionConfig{ + Kerberos: kerberosConfig{ + ServerPrincipal: "hive/_HOST@HIVE.EXAMPLE.COM", + ClientPrincipal: "alice@CLIENT.EXAMPLE.COM", + AuthorizationID: "hive-proxy", + QOP: "auth-conf", + ConfigPath: "/etc/krb5.conf", + CCachePath: "/tmp/alice.ccache", + UseTicketCache: true, + CanonicalHostname: true, + }, + ZooKeeperKerberos: zooKeeperKerberosConfig{ + Enabled: true, + Service: "zookeeper", + Realm: "ZK.EXAMPLE.COM", + CanonicalHostname: false, + }, + } + service, options := zooKeeperGSSAPIOptions(config) + if service != "zookeeper" || options.ServerName != "zookeeper/_HOST@ZK.EXAMPLE.COM" { + t.Fatalf("service=%q options=%#v", service, options) + } + if options.QOP != "auth" || options.AuthorizationID != "" || options.CanonicalizeHost || options.Principal != "alice@CLIENT.EXAMPLE.COM" || options.CCachePath != "/tmp/alice.ccache" { + t.Fatalf("unexpected ZooKeeper GSSAPI options: %#v", options) + } +} + +func TestConnectKerberosZooKeeperFailsOverAndUsesTargetHost(t *testing.T) { + previousDial := dialZooKeeperConnection + previousFactory := newZooKeeperSASLClient + previousShuffle := shuffleZooKeeperServers + t.Cleanup(func() { + dialZooKeeperConnection = previousDial + newZooKeeperSASLClient = previousFactory + shuffleZooKeeperServers = previousShuffle + }) + shuffleZooKeeperServers = func([]string) {} + + clientConnection, serverConnection := net.Pipe() + serverErrors := make(chan error, 1) + go func() { + defer serverConnection.Close() + if err := acceptZooKeeperSession(serverConnection); err != nil { + serverErrors <- err + return + } + if err := expectZooKeeperSASLRound(serverConnection, "client-initial", []byte("server-challenge")); err != nil { + serverErrors <- err + return + } + if err := expectZooKeeperSASLRound(serverConnection, "client-final", nil); err != nil { + serverErrors <- err + return + } + request, err := readZooKeeperTestFrame(serverConnection) + if err != nil { + serverErrors <- err + return + } + decoder := newZooKeeperDecoder(request) + _, _ = decoder.int32() + opcode, _ := decoder.int32() + if opcode != zooKeeperOpClose { + serverErrors <- fmt.Errorf("unexpected close opcode %d", opcode) + return + } + serverErrors <- nil + }() + + var dialed []string + dialZooKeeperConnection = func(address string, _ time.Duration, _ *tls.Config) (net.Conn, error) { + dialed = append(dialed, address) + if address == "first.example.com:2181" { + return nil, errors.New("first unavailable") + } + return clientConnection, nil + } + var targetHost string + newZooKeeperSASLClient = func(host string, _ connectionConfig) (zooKeeperSASLClient, error) { + targetHost = host + return &scriptedZooKeeperSASLClient{}, nil + } + + config := connectionConfig{ + Kerberos: kerberosConfig{Enabled: true}, + ZooKeeperKerberos: zooKeeperKerberosConfig{Enabled: true}, + } + client, events, err := connectKerberosZooKeeper( + []string{"first.example.com:2181", "second.example.com:2181"}, + time.Second, + nil, + config, + ) + if err != nil { + t.Fatal(err) + } + if len(dialed) != 2 || dialed[0] != "first.example.com:2181" || dialed[1] != "second.example.com:2181" { + t.Fatalf("dialed = %#v", dialed) + } + if targetHost != "second.example.com" { + t.Fatalf("target host = %q", targetHost) + } + event := <-events + if event.State != zk.StateHasSession || event.Server != "second.example.com:2181" { + t.Fatalf("event = %#v", event) + } + client.Close() + if err := <-serverErrors; err != nil { + t.Fatal(err) + } +} + +func TestZooKeeperKerberosSASLWithMiniKDC(t *testing.T) { + logger := log.New(io.Discard, "", 0) + kdc, err := krb5test.NewKDC(map[string][]string{ + "alice": nil, + "zookeeper/localhost": nil, + }, logger) + if err != nil { + t.Fatal(err) + } + kdc.KRB5Conf.LibDefaults.UDPPreferenceLimit = 1 + kdc.Start() + defer kdc.Close() + + directory := t.TempDir() + configPath := filepath.Join(directory, "krb5.conf") + keytabPath := filepath.Join(directory, "zookeeper.keytab") + configContents := fmt.Sprintf(`[libdefaults] + default_realm = %s + dns_lookup_realm = false + dns_lookup_kdc = false + rdns = false + udp_preference_limit = 1 + default_tgs_enctypes = aes256-cts-hmac-sha1-96 + default_tkt_enctypes = aes256-cts-hmac-sha1-96 + permitted_enctypes = aes256-cts-hmac-sha1-96 + +[realms] + %s = { + kdc = %s + } +`, kdc.Realm, kdc.Realm, kdc.TCPListener.Addr().String()) + if err := os.WriteFile(configPath, []byte(configContents), 0o600); err != nil { + t.Fatal(err) + } + keytabBytes, err := kdc.Keytab.Marshal() + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(keytabPath, keytabBytes, 0o600); err != nil { + t.Fatal(err) + } + t.Setenv("KRB5_KTNAME", keytabPath) + t.Setenv("KRB5_CLIENT_KTNAME", keytabPath) + + clientConnection, serverConnection := net.Pipe() + serverErrors := make(chan error, 1) + servicePrincipal := "zookeeper/localhost@" + kdc.Realm + go func() { + defer serverConnection.Close() + serverErrors <- serveKerberosZooKeeperSASL(serverConnection, servicePrincipal) + }() + + protocolClient, err := newProtocolZooKeeperClient(clientConnection, 5*time.Second) + if err != nil { + t.Fatal(err) + } + config := connectionConfig{ + Kerberos: kerberosConfig{ + Enabled: true, + ClientPrincipal: "alice@" + kdc.Realm, + ConfigPath: configPath, + Password: kdc.Principals["alice"].Password, + DisablePAFXFAST: true, + }, + ZooKeeperKerberos: zooKeeperKerberosConfig{ + Enabled: true, + Service: "zookeeper", + ServerPrincipal: servicePrincipal, + CanonicalHostname: false, + }, + } + saslClient, err := newZooKeeperSASLClient("localhost", config) + if err != nil { + t.Fatal(err) + } + if err := protocolClient.authenticateSASL(saslClient); err != nil { + t.Fatal(err) + } + protocolClient.Close() + if err := <-serverErrors; err != nil { + t.Fatal(err) + } +} + +func serveZooKeeperProtocolTest(connection net.Conn) error { + if err := acceptZooKeeperSession(connection); err != nil { + return err + } + if err := expectZooKeeperSASLRound(connection, "client-initial", []byte("server-challenge")); err != nil { + return err + } + if err := expectZooKeeperSASLRound(connection, "client-final", nil); err != nil { + return err + } + request, err := readZooKeeperTestFrame(connection) + if err != nil { + return err + } + decoder := newZooKeeperDecoder(request) + xid, _ := decoder.int32() + opcode, _ := decoder.int32() + authType, _ := decoder.int32() + scheme, _ := decoder.string() + auth, _ := decoder.bytes() + if opcode != zooKeeperOpSetAuth || authType != 0 || scheme != "digest" || string(auth) != "user:password" { + return fmt.Errorf("unexpected auth request opcode=%d type=%d scheme=%q auth=%q", opcode, authType, scheme, auth) + } + if err := writeZooKeeperTestResponse(connection, xid, 0, nil); err != nil { + return err + } + request, err = readZooKeeperTestFrame(connection) + if err != nil { + return err + } + decoder = newZooKeeperDecoder(request) + xid, _ = decoder.int32() + opcode, _ = decoder.int32() + path, _ := decoder.string() + watch, _ := decoder.take(1) + if opcode != zooKeeperOpGetChildren2 || path != "/hiveserver2" || !bytes.Equal(watch, []byte{0}) { + return fmt.Errorf("unexpected children request opcode=%d path=%q watch=%v", opcode, path, watch) + } + body := &zooKeeperEncoder{} + body.int32(1) + body.string("server-1") + encodeZooKeeperTestStat(body, 0, 1) + if err := writeZooKeeperTestResponse(connection, xid, 0, body.data()); err != nil { + return err + } + request, err = readZooKeeperTestFrame(connection) + if err != nil { + return err + } + decoder = newZooKeeperDecoder(request) + xid, _ = decoder.int32() + opcode, _ = decoder.int32() + path, _ = decoder.string() + watch, _ = decoder.take(1) + if opcode != zooKeeperOpGetData || path != "/hiveserver2/server-1" || !bytes.Equal(watch, []byte{0}) { + return fmt.Errorf("unexpected get request opcode=%d path=%q watch=%v", opcode, path, watch) + } + data := []byte("serverUri=hs2.example.com:10000") + body = &zooKeeperEncoder{} + body.bytes(data) + encodeZooKeeperTestStat(body, int32(len(data)), 0) + if err := writeZooKeeperTestResponse(connection, xid, 0, body.data()); err != nil { + return err + } + request, err = readZooKeeperTestFrame(connection) + if err != nil { + return err + } + decoder = newZooKeeperDecoder(request) + _, _ = decoder.int32() + opcode, _ = decoder.int32() + if opcode != zooKeeperOpClose { + return fmt.Errorf("unexpected close opcode %d", opcode) + } + return nil +} + +func serveKerberosZooKeeperSASL(connection net.Conn, servicePrincipal string) error { + if err := acceptZooKeeperSession(connection); err != nil { + return err + } + acceptor := gsskrb5.NewKrb5Mech() + if err := acceptor.Accept(servicePrincipal); err != nil { + return err + } + xid, token, err := readZooKeeperSASLRequest(connection) + if err != nil { + return err + } + apReply, err := acceptor.Continue(token) + if err != nil { + return err + } + if !acceptor.IsEstablished() { + return errors.New("Kerberos acceptor context was not established") + } + if err := writeZooKeeperSASLResponse(connection, xid, apReply); err != nil { + return err + } + xid, token, err = readZooKeeperSASLRequest(connection) + if err != nil { + return err + } + if len(token) != 0 { + return fmt.Errorf("expected empty post-AP-REP token, got %x", token) + } + securityChallenge, err := acceptor.Wrap([]byte{1, 0, 0, 0}, false) + if err != nil { + return err + } + if err := writeZooKeeperSASLResponse(connection, xid, securityChallenge); err != nil { + return err + } + xid, token, err = readZooKeeperSASLRequest(connection) + if err != nil { + return err + } + securityResponse, sealed, err := acceptor.Unwrap(token) + if err != nil { + return err + } + if sealed || len(securityResponse) < 4 || !bytes.Equal(securityResponse[:4], []byte{1, 0, 0, 0}) { + return fmt.Errorf("invalid GSSAPI security-layer response sealed=%v payload=%x", sealed, securityResponse) + } + if len(securityResponse[4:]) == 0 { + return errors.New("GSSAPI authorization identity is empty") + } + if err := writeZooKeeperSASLResponse(connection, xid, nil); err != nil { + return err + } + request, err := readZooKeeperTestFrame(connection) + if err != nil { + return err + } + decoder := newZooKeeperDecoder(request) + _, _ = decoder.int32() + opcode, _ := decoder.int32() + if opcode != zooKeeperOpClose { + return fmt.Errorf("unexpected close opcode %d", opcode) + } + return nil +} + +func readZooKeeperSASLRequest(connection net.Conn) (int32, []byte, error) { + request, err := readZooKeeperTestFrame(connection) + if err != nil { + return 0, nil, err + } + decoder := newZooKeeperDecoder(request) + xid, err := decoder.int32() + if err != nil { + return 0, nil, err + } + opcode, err := decoder.int32() + if err != nil { + return 0, nil, err + } + if opcode != zooKeeperOpSASL { + return 0, nil, fmt.Errorf("unexpected SASL opcode %d", opcode) + } + token, err := decoder.bytes() + return xid, token, err +} + +func writeZooKeeperSASLResponse(connection net.Conn, xid int32, token []byte) error { + body := &zooKeeperEncoder{} + body.bytes(token) + return writeZooKeeperTestResponse(connection, xid, 0, body.data()) +} + +func acceptZooKeeperSession(connection net.Conn) error { + request, err := readZooKeeperTestFrame(connection) + if err != nil { + return err + } + decoder := newZooKeeperDecoder(request) + version, _ := decoder.int32() + _, _ = decoder.int64() + timeout, _ := decoder.int32() + sessionID, _ := decoder.int64() + password, _ := decoder.bytes() + if version != 0 || timeout <= 0 || sessionID != 0 || len(password) != 16 { + return fmt.Errorf("invalid connect request version=%d timeout=%d session=%d password=%d", version, timeout, sessionID, len(password)) + } + response := &zooKeeperEncoder{} + response.int32(0) + response.int32(timeout) + response.int64(42) + response.bytes(make([]byte, 16)) + return writeZooKeeperTestFrame(connection, response.data()) +} + +func expectZooKeeperSASLRound(connection net.Conn, wantToken string, responseToken []byte) error { + request, err := readZooKeeperTestFrame(connection) + if err != nil { + return err + } + decoder := newZooKeeperDecoder(request) + xid, _ := decoder.int32() + opcode, _ := decoder.int32() + token, _ := decoder.bytes() + if opcode != zooKeeperOpSASL || string(token) != wantToken { + return fmt.Errorf("unexpected SASL request opcode=%d token=%q", opcode, token) + } + body := &zooKeeperEncoder{} + body.bytes(responseToken) + return writeZooKeeperTestResponse(connection, xid, 0, body.data()) +} + +func encodeZooKeeperTestStat(encoder *zooKeeperEncoder, dataLength, children int32) { + encoder.int64(1) + encoder.int64(2) + encoder.int64(3) + encoder.int64(4) + encoder.int32(5) + encoder.int32(6) + encoder.int32(7) + encoder.int64(8) + encoder.int32(dataLength) + encoder.int32(children) + encoder.int64(9) +} + +func writeZooKeeperTestResponse(connection net.Conn, xid, code int32, body []byte) error { + response := &zooKeeperEncoder{} + response.int32(xid) + response.int64(0) + response.int32(code) + response.buffer.Write(body) + return writeZooKeeperTestFrame(connection, response.data()) +} + +func readZooKeeperTestFrame(connection net.Conn) ([]byte, error) { + client := &protocolZooKeeperClient{connection: connection, timeout: time.Second} + return client.readFrame() +} + +func writeZooKeeperTestFrame(connection net.Conn, payload []byte) error { + client := &protocolZooKeeperClient{connection: connection, timeout: time.Second} + return client.writeFrame(payload) +} + +var _ zooKeeperClient = (*protocolZooKeeperClient)(nil) +var _ = zk.StateSaslAuthenticated diff --git a/agents/drivers/hive-go/zookeeper_tls.go b/agents/drivers/hive-go/zookeeper_tls.go new file mode 100644 index 0000000000..4ad6265a45 --- /dev/null +++ b/agents/drivers/hive-go/zookeeper_tls.go @@ -0,0 +1,240 @@ +package main + +import ( + "bytes" + "crypto/tls" + "crypto/x509" + "encoding/pem" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + keystore "github.com/pavlo-v-chernykh/keystore-go/v4" + pkcs12 "software.sslmate.com/src/go-pkcs12" +) + +func buildZooKeeperTLSConfig(values map[string]string) (*tls.Config, error) { + if !parameterBool(values, "zookeepersslenable") { + return nil, nil + } + config := &tls.Config{ + MinVersion: tls.VersionTLS12, + ServerName: parameter(values, "zookeeperservername"), + } + trustStoreLocation := parameter(values, "zookeepertruststorelocation") + if trustStoreLocation != "" { + certificates, err := loadTrustStore( + trustStoreLocation, + parameter(values, "zookeepertruststorepassword"), + parameter(values, "zookeepertruststoretype"), + ) + if err != nil { + return nil, fmt.Errorf("load ZooKeeper truststore: %w", err) + } + pool := x509.NewCertPool() + for _, certificate := range certificates { + pool.AddCert(certificate) + } + config.RootCAs = pool + } + keyStoreLocation := parameter(values, "zookeeperkeystorelocation") + if keyStoreLocation != "" { + certificate, err := loadClientKeyStore( + keyStoreLocation, + parameter(values, "zookeeperkeystorepassword"), + parameter(values, "zookeeperkeystoretype"), + ) + if err != nil { + return nil, fmt.Errorf("load ZooKeeper keystore: %w", err) + } + config.Certificates = []tls.Certificate{certificate} + } + if parameterBool(values, "zookeepersslinsecureskipverify") { + config.InsecureSkipVerify = true + } + return config, nil +} + +func loadTrustStore(path, password, storeType string) ([]*x509.Certificate, error) { + contents, err := os.ReadFile(path) + if err != nil { + return nil, err + } + switch normalizedStoreType(storeType, path) { + case "PEM": + return parsePEMCertificates(contents) + case "PKCS12": + certificates, err := pkcs12.DecodeTrustStore(contents, password) + if err == nil { + return certificates, nil + } + _, certificate, chain, chainErr := pkcs12.DecodeChain(contents, password) + if chainErr != nil { + return nil, err + } + return append([]*x509.Certificate{certificate}, chain...), nil + case "JKS": + store, err := loadJKS(contents, password) + if err != nil { + return nil, err + } + var certificates []*x509.Certificate + for _, alias := range store.Aliases() { + switch { + case store.IsTrustedCertificateEntry(alias): + entry, getErr := store.GetTrustedCertificateEntry(alias) + if getErr != nil { + return nil, getErr + } + certificate, parseErr := x509.ParseCertificate(entry.Certificate.Content) + if parseErr != nil { + return nil, parseErr + } + certificates = append(certificates, certificate) + case store.IsPrivateKeyEntry(alias): + chain, getErr := store.GetPrivateKeyEntryCertificateChain(alias) + if getErr != nil { + return nil, getErr + } + for _, entry := range chain { + certificate, parseErr := x509.ParseCertificate(entry.Content) + if parseErr != nil { + return nil, parseErr + } + certificates = append(certificates, certificate) + } + } + } + if len(certificates) == 0 { + return nil, errors.New("JKS truststore contains no certificates") + } + return certificates, nil + default: + return nil, fmt.Errorf("unsupported store type %q", storeType) + } +} + +func loadClientKeyStore(path, password, storeType string) (tls.Certificate, error) { + contents, err := os.ReadFile(path) + if err != nil { + return tls.Certificate{}, err + } + switch normalizedStoreType(storeType, path) { + case "PEM": + return tls.X509KeyPair(contents, contents) + case "PKCS12": + privateKey, certificate, chain, err := pkcs12.DecodeChain(contents, password) + if err != nil { + return tls.Certificate{}, err + } + result := tls.Certificate{PrivateKey: privateKey, Leaf: certificate} + result.Certificate = append(result.Certificate, certificate.Raw) + for _, entry := range chain { + result.Certificate = append(result.Certificate, entry.Raw) + } + return result, nil + case "JKS": + store, err := loadJKS(contents, password) + if err != nil { + return tls.Certificate{}, err + } + passwordBytes := []byte(password) + defer clear(passwordBytes) + for _, alias := range store.Aliases() { + if !store.IsPrivateKeyEntry(alias) { + continue + } + entry, getErr := store.GetPrivateKeyEntry(alias, passwordBytes) + if getErr != nil { + return tls.Certificate{}, getErr + } + privateKey, parseErr := parsePrivateKey(entry.PrivateKey) + if parseErr != nil { + return tls.Certificate{}, parseErr + } + result := tls.Certificate{PrivateKey: privateKey} + for index, certificate := range entry.CertificateChain { + result.Certificate = append(result.Certificate, certificate.Content) + if index == 0 { + result.Leaf, _ = x509.ParseCertificate(certificate.Content) + } + } + if len(result.Certificate) == 0 { + return tls.Certificate{}, errors.New("JKS private key entry has no certificate chain") + } + return result, nil + } + return tls.Certificate{}, errors.New("JKS keystore contains no private key entry") + default: + return tls.Certificate{}, fmt.Errorf("unsupported store type %q", storeType) + } +} + +func normalizedStoreType(storeType, path string) string { + value := strings.ToUpper(strings.TrimSpace(storeType)) + switch value { + case "P12", "PFX", "PKCS#12": + return "PKCS12" + case "X509", "X.509": + return "PEM" + case "": + switch strings.ToLower(filepath.Ext(path)) { + case ".jks": + return "JKS" + case ".p12", ".pfx", ".pkcs12": + return "PKCS12" + default: + return "PEM" + } + default: + return value + } +} + +func loadJKS(contents []byte, password string) (keystore.KeyStore, error) { + store := keystore.New() + passwordBytes := []byte(password) + defer clear(passwordBytes) + if err := store.Load(bytes.NewReader(contents), passwordBytes); err != nil { + return keystore.KeyStore{}, err + } + return store, nil +} + +func parsePEMCertificates(contents []byte) ([]*x509.Certificate, error) { + var certificates []*x509.Certificate + for len(contents) > 0 { + block, rest := pem.Decode(contents) + if block == nil { + break + } + contents = rest + if block.Type != "CERTIFICATE" { + continue + } + certificate, err := x509.ParseCertificate(block.Bytes) + if err != nil { + return nil, err + } + certificates = append(certificates, certificate) + } + if len(certificates) == 0 { + return nil, errors.New("PEM truststore contains no certificates") + } + return certificates, nil +} + +func parsePrivateKey(contents []byte) (any, error) { + if value, err := x509.ParsePKCS8PrivateKey(contents); err == nil { + return value, nil + } + if value, err := x509.ParsePKCS1PrivateKey(contents); err == nil { + return value, nil + } + if value, err := x509.ParseECPrivateKey(contents); err == nil { + return value, nil + } + return nil, errors.New("unsupported private key encoding") +} diff --git a/agents/drivers/hive-go/zookeeper_tls_test.go b/agents/drivers/hive-go/zookeeper_tls_test.go new file mode 100644 index 0000000000..07e65d002f --- /dev/null +++ b/agents/drivers/hive-go/zookeeper_tls_test.go @@ -0,0 +1,162 @@ +package main + +import ( + "crypto/rand" + "crypto/rsa" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "math/big" + "os" + "path/filepath" + "testing" + "time" + + keystore "github.com/pavlo-v-chernykh/keystore-go/v4" + pkcs12 "software.sslmate.com/src/go-pkcs12" +) + +func testZooKeeperCertificate(t *testing.T) (*rsa.PrivateKey, *x509.Certificate, []byte) { + t.Helper() + privateKey, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + t.Fatal(err) + } + template := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "zk.example.com"}, + DNSNames: []string{"zk.example.com"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + KeyUsage: x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment | x509.KeyUsageCertSign, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth}, + IsCA: true, + } + raw, err := x509.CreateCertificate(rand.Reader, template, template, &privateKey.PublicKey, privateKey) + if err != nil { + t.Fatal(err) + } + certificate, err := x509.ParseCertificate(raw) + if err != nil { + t.Fatal(err) + } + return privateKey, certificate, raw +} + +func TestBuildZooKeeperTLSConfigFromPEM(t *testing.T) { + _, _, raw := testZooKeeperCertificate(t) + path := filepath.Join(t.TempDir(), "trust.pem") + if err := os.WriteFile(path, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: raw}), 0o600); err != nil { + t.Fatal(err) + } + config, err := buildZooKeeperTLSConfig(map[string]string{ + "zookeepersslenable": "true", + "zookeepertruststorelocation": path, + "zookeepertruststoretype": "PEM", + "zookeeperservername": "zk.example.com", + }) + if err != nil { + t.Fatal(err) + } + if config == nil || config.RootCAs == nil || config.ServerName != "zk.example.com" { + t.Fatalf("unexpected TLS config: %#v", config) + } +} + +func TestBuildZooKeeperTLSConfigFromJKS(t *testing.T) { + privateKey, certificate, raw := testZooKeeperCertificate(t) + password := []byte("changeit") + store := keystore.New() + if err := store.SetTrustedCertificateEntry("ca", keystore.TrustedCertificateEntry{ + CreationTime: time.Now(), + Certificate: keystore.Certificate{Type: "X509", Content: raw}, + }); err != nil { + t.Fatal(err) + } + encodedKey, err := x509.MarshalPKCS8PrivateKey(privateKey) + if err != nil { + t.Fatal(err) + } + if err := store.SetPrivateKeyEntry("client", keystore.PrivateKeyEntry{ + CreationTime: time.Now(), + PrivateKey: encodedKey, + CertificateChain: []keystore.Certificate{ + {Type: "X509", Content: certificate.Raw}, + }, + }, password); err != nil { + t.Fatal(err) + } + path := filepath.Join(t.TempDir(), "client.jks") + file, err := os.Create(path) + if err != nil { + t.Fatal(err) + } + if err := store.Store(file, password); err != nil { + file.Close() + t.Fatal(err) + } + if err := file.Close(); err != nil { + t.Fatal(err) + } + config, err := buildZooKeeperTLSConfig(map[string]string{ + "zookeepersslenable": "true", + "zookeepertruststorelocation": path, + "zookeepertruststorepassword": string(password), + "zookeepertruststoretype": "JKS", + "zookeeperkeystorelocation": path, + "zookeeperkeystorepassword": string(password), + "zookeeperkeystoretype": "JKS", + }) + if err != nil { + t.Fatal(err) + } + if config.RootCAs == nil || len(config.Certificates) != 1 || config.Certificates[0].PrivateKey == nil { + t.Fatalf("unexpected JKS TLS config: %#v", config) + } +} + +func TestBuildZooKeeperTLSConfigFromPKCS12(t *testing.T) { + privateKey, certificate, _ := testZooKeeperCertificate(t) + password := "changeit" + keyStore, err := pkcs12.Modern.Encode(privateKey, certificate, nil, password) + if err != nil { + t.Fatal(err) + } + trustStore, err := pkcs12.Modern.EncodeTrustStore([]*x509.Certificate{certificate}, password) + if err != nil { + t.Fatal(err) + } + directory := t.TempDir() + keyPath := filepath.Join(directory, "client.p12") + trustPath := filepath.Join(directory, "trust.p12") + if err := os.WriteFile(keyPath, keyStore, 0o600); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(trustPath, trustStore, 0o600); err != nil { + t.Fatal(err) + } + config, err := buildZooKeeperTLSConfig(map[string]string{ + "zookeepersslenable": "true", + "zookeepertruststorelocation": trustPath, + "zookeepertruststorepassword": password, + "zookeepertruststoretype": "PKCS12", + "zookeeperkeystorelocation": keyPath, + "zookeeperkeystorepassword": password, + "zookeeperkeystoretype": "PKCS12", + }) + if err != nil { + t.Fatal(err) + } + if config.RootCAs == nil || len(config.Certificates) != 1 { + t.Fatalf("unexpected PKCS12 TLS config: %#v", config) + } +} + +func TestZooKeeperTLSRequiresExistingStore(t *testing.T) { + if _, err := buildZooKeeperTLSConfig(map[string]string{ + "zookeepersslenable": "true", + "zookeepertruststorelocation": filepath.Join(t.TempDir(), "missing.jks"), + }); err == nil { + t.Fatal("expected missing truststore error") + } +} diff --git a/agents/go-common/go-gssapi/LICENSE b/agents/go-common/go-gssapi/LICENSE new file mode 100644 index 0000000000..6e9eb7a10c --- /dev/null +++ b/agents/go-common/go-gssapi/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [2021] [Jacob Scott] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/agents/go-common/go-gssapi/README.md b/agents/go-common/go-gssapi/README.md new file mode 100644 index 0000000000..8d24f8134c --- /dev/null +++ b/agents/go-common/go-gssapi/README.md @@ -0,0 +1,228 @@ +# go-gssapi: pure Go GSS-API implementation + + +![GitHub tag (latest SemVer)](https://img.shields.io/github/v/tag/golang-auth/go-gssapi) +[![Git Workflow](https://img.shields.io/github/workflow/status/golang-auth/go-gssapi/unit-tests/v2)](https://img.shields.io/github/workflow/status/golang-auth/go-gssapi/unit-tests/v2) +[![Go Version](https://img.shields.io/badge/go%20version-%3E=1.13-61CFDD.svg?style=flat-square)](https://golang.org/) +[![PkgGoDev](https://pkg.go.dev/badge/mod/github.com/golang-auth/go-gssapi/v2)](https://pkg.go.dev/mod/github.com/golang-auth/go-gssapi/v2) + + +go-gssapi is a pure Golang implementation of the GSS-API version 2 specification ([RFC 2743](https://tools.ietf.org/html/rfc2743)). It uses the pure [Golang Kerberos implementation](https://github.com/jcmturner/gokrb5/tree/master/v8) to obtain Kerberos tickets and perform cryptographic operations. + +Documentation at https://pkg.go.dev/github.com/golang-auth/go-gssapi/v2 + +## Implemented functionality + +The current version implements the Kerberos V5 authentication mechanism ([RFC 4121](https://tools.ietf.org/html/rfc4121)) We plan +to support SPNEGO in a future release. + +The following features are currently available: + + * Initiator (client) and Acceptor (server) + * Mutual authentication + * Message Integrity and Confidentiality + * GSS MIC and Wrap tokens + * Basic support for detecting out-of-sequence and duplicate messages + + +The following functionality is currently not available: + + * GSS-API v1 message tokens ([RFC1964](https://tools.ietf.org/html/rfc1964)) + * Delegation + * Channel binding + + +## Platforms + +Currently Unix platforms. Support for the Windows SSPI will be added in future. + +## Project status + +go-gssapi is not yet considered to be production ready, and the interface is not yet considered stable. We started +at v2 to support the new Go versioning standrd. A v2.0.0 tag will be created after which we will follow normal [semantic versioning conventions](https://golang.org/ref/mod#versions). + +## Configuration + +The Kerberos implementiaton behing the GSS-API functionality is hidden form the caller. The name of the krb5.conf file, credential cache (by an Initiator) and keytab file (by an Acceptor) are identified from the environment or using default paths if not specified in the environment : + +| Configuration item | Environment variable | Default | +| ------------------ | -------------------- | ------- | +| Kerberos client configuartion file | `KRB5_CONFIG` | `/etc/krb5.conf` | +| Credential cache file | `KRB5CCNAME` | `/etc/krb5cc_%{UID}` | +| Keytab file | `KRB5_KTNAME` | `/var/kerberos/krb5/user/%{UID}/client.keytab` | + +.. where `%{UID}` is the current username as returned by `id -u` + +> :droplet: A future releaes will choose defaults based on the platform we're running on, and may provide more complete control of the gokrb client. Please file a new Github issue if you have a use-case. + +## Import paths and versions + +Import the base `go-gssapi` package where you will use GSS-API, and the mechanism specific package somewhere (usually the main package) to register the mechanism : + +```go + package myclient + + import "github.com/golang-auth/go-gssapi/v2" +``` + +```go + package main + + // register the Kerberos GSS-API mechsnism + import _ "github.com/golang-auth/go-gssapi/v2/krb5" +``` + +We will maintain major versions of the library in separate Git branches. To use the latest version in a branch, use `go get` with the branch name, eg : + + go get github.com/golang-auth/go-gssapi/v2@v2.0.0-alpha + +This will result in a concrete version being added to `go.mod`: + +``` +module test + +go 1.15 + +require github.com/golang-auth/go-gssapi/v2 v2.0.0-alpha //indirect +``` + +## Initialization + +Obtain an instance (_context_) of a mechanism-specif implementation: +```go + ctx := gssapi.NewMech("kerberos_v5) +``` + +then configure the context as an Initiator (client): + +```go + service := "ldap/ldap.example.com" + var flags gssapi.ContextFlag = gssapi.ContextFlagInteg | + gssapi.ContextFlagConf | + gssapi.ContextFlagReplay | + gssapi.ContextFlagSequence | + gssapi.ContextFlagMutual + err := ctx.Initiate(service, flags) +``` + + + or Acceptor (server): + ```go + err := ctx.Accept("") + ``` + + > Note: go-gssapi uses the Kebreros V5 principal name format (primary/instance), not the NT_HOSTBASED_SERVICE format (service@host) used by other GSS-API implementations + + ## Negotiation loop + + Once the context is initialized as an Initiator or Acceptor, the client and server enter a loop until `IsEstablished` returns true. Each pass through the loop, the `Continue` method is called on the context. `Continue` may return an opaque token that should be sent to the peer. The peer passes the token to the `Continue` method and the process is repeated. + + ### Initiator loop + ```go +var inToken, outToken []byte + +for !ctx.IsEstablished() { + outToken, err = ctx.Continue(inToken) + if err != nil { + break + } + + if len(outToken) > 0 { + if sendErr := sendToken(conn, outToken); sendErr != nil { + err = sendErr + break + } + } + + if !ctx.IsEstablished() { + inToken, err = recvToken(conn) + if err != nil { + break + } + } +} + +if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) +} + ``` + +### Acceptor loop +```go +var inToken, outToken []byte + +for !ctx.IsEstablished() { + inToken, err = recvToken(conn) + if err != nil { + fmt.Fprintln(os.Stderr, err) + break + } + + outToken, err = ctx.Continue(inToken) + if len(outToken) > 0 { + if sendErr := sendToken(conn, outToken); sendErr != nil { + err = sendErr + break + } + } +} + +if err != nil { + fmt.Fprintln(os.Stderr, err) + return +} +``` + +## Exchanging messages + +GSS-API defines two message types : + * Wrap messages encapsulate a message payload plus either a signature, or with the payload encrypted (if confidentially is requested) + * Message Integrity Code (MIC) messages convey a signature of a payload but do not encapsulate the actual payload + +The context `Wrap` and `MakeSignature` methods are used to generate serialized GSS-API message tokens for a given payload. MIC tokens can be communicated separately from the payload, for example as a part of an acknowledgement or via a separate channel from the data. + +The receiver of a message passes the token to the context `Unwrap` and `VerifySignature` methods. Unwrap verifies the message payload signature or decrypts the message payload (if confidentially is in use). VerifySignature verifies the signature of a payload passed to the method. + +### sender +```go + +msg := "Hello go-gssapi!" +seal := true // ask for confidentiality + +// Wrap the message +outToken, err = ctx.Wrap([]byte(msg), seal) +if err != nil { + return err +} + +// send it to the GSS-API peer +if err = sendToken(conn, outToken); err != nil { + retirm err +} + +``` + + +### receiver +```go +inToken, err := recvToken(conn) +if err != nil { + return err +} + +msg, isSealed, err := ctx.Unwrap(inToken) +if err != nil { + return err +} + +protStr := "signed" +if isSealed { + protStr = "sealed" +} +fmt.Printf(`Received %s message: "%s"`+"\n", protStr, msg) +``` + +## Sample code + +Examples implemented in C (for use with the MIT or Heimdal GSS-API library) and Go are included in the `examples` directory. The C samples are provided to demonstrate go-gssapi interoperatbility. diff --git a/agents/go-common/go-gssapi/common/channel_binding.go b/agents/go-common/go-gssapi/common/channel_binding.go new file mode 100644 index 0000000000..a58ddb9cf2 --- /dev/null +++ b/agents/go-common/go-gssapi/common/channel_binding.go @@ -0,0 +1,17 @@ +package common + +import "net" + +type AddressFamily uint32 + +const ( + GssAddrFamilyUnspecified AddressFamily = 0 + GssAddrFamilyLOCAL AddressFamily = 1 + GssAddrFamilyINET AddressFamily = 2 +) + +type ChannelBinding struct { + InitiatorAddr net.Addr + AcceptorAddr net.Addr + Data []byte +} diff --git a/agents/go-common/go-gssapi/flags.go b/agents/go-common/go-gssapi/flags.go new file mode 100644 index 0000000000..e4abc3ed5e --- /dev/null +++ b/agents/go-common/go-gssapi/flags.go @@ -0,0 +1,59 @@ +package gssapi + +import "strings" + +type ContextFlag uint32 + +// GSS-API context flags assigned numbers. +const ( + ContextFlagDeleg ContextFlag = 1 << iota // delegate credentials, not currently supported + ContextFlagMutual // request remote peer authenticates itself + ContextFlagReplay // enable replay detection for signed/sealed messages + ContextFlagSequence // enable detection of out of sequence signed/sealed messages + ContextFlagConf // confidentiality available + ContextFlagInteg // integrity available +) + +// FlagList returns a slice of individual flags derived from the +// composite value f +func FlagList(f ContextFlag) (fl []ContextFlag) { + t := ContextFlag(1) + for i := 0; i < 32; i++ { + if f&t != 0 { + fl = append(fl, t) + } + + t <<= 1 + } + + return +} + +// FlagName returns a human-readable description of a context flag value +func FlagName(f ContextFlag) string { + switch f { + case ContextFlagDeleg: + return "Delegation" + case ContextFlagMutual: + return "Mutual authentication" + case ContextFlagReplay: + return "Message replay detection" + case ContextFlagSequence: + return "Out of sequence message detection" + case ContextFlagConf: + return "Confidentiality" + case ContextFlagInteg: + return "Integrity" + } + + return "Unknown" +} + +func (f ContextFlag) String() string { + var names []string + for _, flag := range FlagList(f) { + names = append(names, FlagName(flag)) + } + + return strings.Join(names, ", ") +} diff --git a/agents/go-common/go-gssapi/flags_test.go b/agents/go-common/go-gssapi/flags_test.go new file mode 100644 index 0000000000..c248d68f80 --- /dev/null +++ b/agents/go-common/go-gssapi/flags_test.go @@ -0,0 +1,29 @@ +package gssapi + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestFlagList(t *testing.T) { + flags := ContextFlagConf | ContextFlagMutual | ContextFlagDeleg + flaglist := FlagList(flags) + + assert.ElementsMatch(t, []ContextFlag{ContextFlagConf, ContextFlagMutual, ContextFlagDeleg}, flaglist) +} + +func TestFlagName(t *testing.T) { + assert.Equal(t, "Delegation", FlagName(ContextFlagDeleg)) + assert.Equal(t, "Integrity", FlagName(ContextFlagInteg)) +} + +func TestFlagString(t *testing.T) { + flags := ContextFlagConf | ContextFlagMutual | ContextFlagDeleg + str := flags.String() + + assert.Contains(t, str, "Delegation") + assert.Contains(t, str, "Mutual") + assert.Contains(t, str, "Confidentiality") + assert.NotContains(t, str, "sequence") +} diff --git a/agents/go-common/go-gssapi/go.mod b/agents/go-common/go-gssapi/go.mod new file mode 100644 index 0000000000..b63d8cfeb6 --- /dev/null +++ b/agents/go-common/go-gssapi/go.mod @@ -0,0 +1,11 @@ +module github.com/golang-auth/go-gssapi/v2 + +go 1.15 + +require ( + github.com/jcmturner/gofork v1.0.0 + github.com/jcmturner/gokrb5 v8.4.2+incompatible + github.com/jcmturner/gokrb5/v8 v8.4.2 + github.com/stretchr/testify v1.7.0 + golang.org/x/net v0.0.0-20210326220855-61e056675ecf // indirect +) diff --git a/agents/go-common/go-gssapi/go.sum b/agents/go-common/go-gssapi/go.sum new file mode 100644 index 0000000000..1f60a22d80 --- /dev/null +++ b/agents/go-common/go-gssapi/go.sum @@ -0,0 +1,49 @@ +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/gorilla/securecookie v1.1.1 h1:miw7JPhV+b/lAHSXz4qd/nN9jRiAFV5FwjeKyCS8BvQ= +github.com/gorilla/securecookie v1.1.1/go.mod h1:ra0sb63/xPlUeL+yeDciTfxMRAA+MP+HVt/4epWDjd4= +github.com/gorilla/sessions v1.2.1 h1:DHd3rPN5lE3Ts3D8rKkQ8x/0kqfeNmBAaiSi+o7FsgI= +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 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +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/go-common/go-gssapi/interface.go b/agents/go-common/go-gssapi/interface.go new file mode 100644 index 0000000000..2322e5cf14 --- /dev/null +++ b/agents/go-common/go-gssapi/interface.go @@ -0,0 +1,95 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +/* +Package gssapi provides a Go interface to the Generic Security +Services Application Program Interface. + +The package defines an interface that GSS-API mechanism specific +code should conform to. + +An Initiator (ie. client) uses the Initiate method to start the +authentiation process. An Acceptor (ie. server) uses the Accpet +method instead. After that, both sides call Continue in a loop, +transferring token between themselves using a suitable communication +protocol. When IsEstablished returns true, the security context +can be used to securely transfer messages or message signatures using +Wrap/Unwrap or MakeSignature/VerifySignature. +*/ +package gssapi + +import "github.com/golang-auth/go-gssapi/v2/common" + +// Mech defines the interface to a GSS-API mechanism +type Mech interface { + // IsEstablished can be used to determine whether the security + // context between an Initiator and Acceptor is complete and + // is ready to transfer messages between the peers. + IsEstablished() bool + + // ContextFlags returns the security flags negotiated between + // the initiator and acceptor. The flags *SHOULD* be checked + // before using the context to verify that desired security + // requirements have been met. + ContextFlags() ContextFlag + + // PeerName returns a string representing the peer's identity + PeerName() string + + // SSF returns the Security Strength Factor of the channel established + // by the security context + SSF() uint + + // WrapSizeLimit returns the maximum possible message size that can + // be presented to Wrap() to produce to output token no longer than + // requestedOutputSize bytes + WrapSizeLimit(requestedOutputSize uint32, confidentiality bool) uint32 + + // Initiate is used by a GSS-API Initiator to start the + // context negotiation process with a remote Acceptor. + // serverName is the mechanism specific name of the remote + // Acceptor, and flags represent the desired security + // properties of the context. + Initiate(serviceName string, flags ContextFlag, cb *common.ChannelBinding) (err error) + + // Accept is used by a GSS-API Acceptor to begin context + // negotiation with a remote Initiator. + // If provided, serviceName is the mechanism specific identifier + // of the local Acceptor + Accept(serviceName string) (err error) + + // Continue is called in a loop by Initiators and Acceptors after + // first calling one of Initiate or Accept. + // tokenIn represents a token received from the peer + // If tokenOut is non-zero, it should be send to the peer + Continue(tokenIn []byte) (tokenOut []byte, err error) + + // Wrap is called by either peer after the context is establighed + // to create a token that encapsulates a payload. If confidentially + // is required, the payload is encrypted (*sealed*) using a key + // negotiated during context establishment. Otherwise, the key + // is used to sign the payload which is encapsulated in the clear. + // tokenOut should be communicated to the peer which should use Unwrap + // on the token. + Wrap(tokenIn []byte, confidentiality bool) (tokenOut []byte, err error) + + // Unwrap is passed a wrap token received from a peer. If the token + // provides confidentially, the key negotiated during context establishment + // is used to decrypt (*unseal*) the payload. Otherwise, the key is used + // to verify the signature that the remote Wrap call calculated for the + // payload. + // tokenOut is the original payload + // isSealed conveys whether the payload was encrypted or not + Unwrap(tokenIn []byte) (tokenOut []byte, isSealed bool, err error) + + // MakeSignature creates a token that includes the signature of the + // provided payload but does not include the payload itself. The + // output token should be sent to the peer, which should use its copy of + // the payload (communicated separately) to verify the signature. + MakeSignature(payload []byte) (tokenOut []byte, err error) + + // VerifySignature is used to check the signature received from a peer + // using a local copy of the payloads. + VerifySignature(payload []byte, tokenIn []byte) (err error) +} diff --git a/agents/go-common/go-gssapi/krb5/APRep.go b/agents/go-common/go-gssapi/krb5/APRep.go new file mode 100644 index 0000000000..ab49e45de9 --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/APRep.go @@ -0,0 +1,138 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +package krb5 + +/* + * Derived from github.com/jcmturner/gokrb5/v8/messages/APRep.go + * + * The modified version adds marshalling functionality. Remove if/when the changes + * make it upstream. + */ + +import ( + "fmt" + "time" + + "github.com/jcmturner/gofork/encoding/asn1" + "github.com/jcmturner/gokrb5/v8/asn1tools" + "github.com/jcmturner/gokrb5/v8/crypto" + "github.com/jcmturner/gokrb5/v8/iana" + "github.com/jcmturner/gokrb5/v8/iana/asnAppTag" + "github.com/jcmturner/gokrb5/v8/iana/keyusage" + "github.com/jcmturner/gokrb5/v8/iana/msgtype" + "github.com/jcmturner/gokrb5/v8/krberror" + + "github.com/jcmturner/gokrb5/v8/messages" + "github.com/jcmturner/gokrb5/v8/types" +) + +// aPRep implements RFC 4120 KRB_AP_REP: https://tools.ietf.org/html/rfc4120#section-5.5.2. +type aPRep struct { + PVNO int `asn1:"explicit,tag:0"` + MsgType int `asn1:"explicit,tag:1"` + EncPart types.EncryptedData `asn1:"explicit,tag:2"` +} + +// encAPRepPart is the encrypted part of KRB_AP_REP. +type encAPRepPart struct { + CTime time.Time `asn1:"generalized,explicit,tag:0"` + Cusec int `asn1:"explicit,tag:1"` + Subkey types.EncryptionKey `asn1:"optional,explicit,tag:2"` + SequenceNumber int64 `asn1:"optional,explicit,tag:3"` +} + +// unmarshal bytes b into the APRep struct. +func (a *aPRep) unmarshal(b []byte) error { + _, err := asn1.UnmarshalWithParams(b, a, fmt.Sprintf("application,explicit,tag:%v", asnAppTag.APREP)) + if err != nil { + return processUnmarshalReplyError(b, err) + } + expectedMsgType := msgtype.KRB_AP_REP + if a.MsgType != expectedMsgType { + return krberror.NewErrorf(krberror.KRBMsgError, "message ID does not indicate a KRB_AP_REP. Expected: %v; Actual: %v", expectedMsgType, a.MsgType) + } + return nil +} + +// marshal the AP-REP message to a byte slice +func (a *aPRep) marshal() (b []byte, err error) { + b, err = asn1.Marshal(*a) + if err != nil { + return + } + + b = asn1tools.AddASNAppTag(b, asnAppTag.APREP) + return +} + +func (a *aPRep) decryptEncPart(sessionKey types.EncryptionKey) (encpart encAPRepPart, err error) { + decrypted, err := crypto.DecryptEncPart(a.EncPart, sessionKey, uint32(keyusage.AP_REP_ENCPART)) + if err != nil { + err = krberror.Errorf(err, krberror.DecryptingError, "error decrypting AP-REP enc-part") + return + } + + err = encpart.unmarshal(decrypted) + if err != nil { + err = krberror.Errorf(err, krberror.EncodingError, "error unmarshalling decrypted AP-REP enc-part") + return + } + + return +} + +// unmarshal bytes b into the APRep encrypted part struct. +func (a *encAPRepPart) unmarshal(b []byte) error { + _, err := asn1.UnmarshalWithParams(b, a, fmt.Sprintf("application,explicit,tag:%v", asnAppTag.EncAPRepPart)) + if err != nil { + return krberror.Errorf(err, krberror.EncodingError, "AP_REP unmarshal error") + } + return nil +} + +func (a *encAPRepPart) marshal() (b []byte, err error) { + b, err = asn1.Marshal(*a) + if err != nil { + return + } + + b = asn1tools.AddASNAppTag(b, asnAppTag.EncAPRepPart) + return +} + +func newAPRep(tkt messages.Ticket, sessionKey types.EncryptionKey, encPart encAPRepPart) (a aPRep, err error) { + m, err := encPart.marshal() + if err != nil { + err = krberror.Errorf(err, krberror.EncodingError, "marshaling error of AP-REP enc-part") + return + } + + ed, err := crypto.GetEncryptedData(m, sessionKey, uint32(keyusage.AP_REP_ENCPART), tkt.EncPart.KVNO) + if err != nil { + err = krberror.Errorf(err, krberror.EncryptingError, "error encrypting AP-REP enc-part") + return + } + + a = aPRep{ + PVNO: iana.PVNO, + MsgType: msgtype.KRB_AP_REP, + EncPart: ed, + } + return +} + +func processUnmarshalReplyError(b []byte, err error) error { + switch err.(type) { + case asn1.StructuralError: + var krberr messages.KRBError + tmperr := krberr.Unmarshal(b) + if tmperr != nil { + return krberror.Errorf(err, krberror.EncodingError, "failed to unmarshal message") + } + return krberr + default: + return krberror.Errorf(err, krberror.EncodingError, "failed to unmarshal message") + } +} diff --git a/agents/go-common/go-gssapi/krb5/APRep_test.go b/agents/go-common/go-gssapi/krb5/APRep_test.go new file mode 100644 index 0000000000..844e4073f1 --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/APRep_test.go @@ -0,0 +1,118 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +package krb5 + +import ( + "encoding/hex" + "testing" + "time" + + "github.com/jcmturner/gokrb5/v8/iana" + "github.com/jcmturner/gokrb5/v8/iana/msgtype" + "github.com/jcmturner/gokrb5/v8/test/testdata" + "github.com/stretchr/testify/assert" + + "github.com/jcmturner/gokrb5/v8/types" +) + +func TestUnmarshalAPRep(t *testing.T) { + t.Parallel() + var a aPRep + b, err := hex.DecodeString(testdata.MarshaledKRB5ap_rep) + if err != nil { + t.Fatalf("Test vector read error: %v", err) + } + err = a.unmarshal(b) + if err != nil { + t.Fatalf("Unmarshal error: %v", err) + } + assert.Equal(t, iana.PVNO, a.PVNO, "PVNO not as expected") + assert.Equal(t, msgtype.KRB_AP_REP, a.MsgType, "MsgType is not as expected") + assert.Equal(t, testdata.TEST_ETYPE, a.EncPart.EType, "Ticket encPart etype not as expected") + assert.Equal(t, iana.PVNO, a.EncPart.KVNO, "Ticket encPart KVNO not as expected") + assert.Equal(t, []byte(testdata.TEST_CIPHERTEXT), a.EncPart.Cipher, "Ticket encPart cipher not as expected") +} + +func TestUnmarshalEncAPRepPart(t *testing.T) { + t.Parallel() + var a encAPRepPart + b, err := hex.DecodeString(testdata.MarshaledKRB5ap_rep_enc_part) + if err != nil { + t.Fatalf("Test vector read error: %v", err) + } + err = a.unmarshal(b) + if err != nil { + t.Fatalf("Unmarshal error: %v", err) + } + // Parse the test time value into a time.Time type + tt, _ := time.Parse(testdata.TEST_TIME_FORMAT, testdata.TEST_TIME) + + assert.Equal(t, tt, a.CTime, "CTime not as expected") + assert.Equal(t, 123456, a.Cusec, "Client microseconds not as expected") + assert.Equal(t, int32(1), a.Subkey.KeyType, "Subkey type not as expected") + assert.Equal(t, []byte("12345678"), a.Subkey.KeyValue, "Subkey value not as expected") + assert.Equal(t, int64(17), a.SequenceNumber, "Sequence number not as expected") +} + +func TestUnmarshalEncAPRepPart_optionalsNULL(t *testing.T) { + t.Parallel() + var a encAPRepPart + b, err := hex.DecodeString(testdata.MarshaledKRB5ap_rep_enc_partOptionalsNULL) + if err != nil { + t.Fatalf("Test vector read error: %v", err) + } + err = a.unmarshal(b) + if err != nil { + t.Fatalf("Unmarshal error: %v", err) + } + // Parse the test time value into a time.Time type + tt, _ := time.Parse(testdata.TEST_TIME_FORMAT, testdata.TEST_TIME) + + assert.Equal(t, tt, a.CTime, "CTime not as expected") + assert.Equal(t, 123456, a.Cusec, "Client microseconds not as expected") +} + +// test with all fields populated +func TestAPRepEncPartMarshall(t *testing.T) { + t.Parallel() + + want, err := hex.DecodeString(testdata.MarshaledKRB5ap_rep_enc_part) + assert.Nil(t, err, "error not expected decoding test data") + + encpart := ktestMakeSampleAPReqEncPart() + + b, err := encpart.marshal() + assert.Nil(t, err, "enc part marshal error not expected") + assert.Equal(t, want, b) +} + +// test with the optionals not present +func TestAPRepEncPartMarshall_optionalsNULL(t *testing.T) { + t.Parallel() + + want, err := hex.DecodeString(testdata.MarshaledKRB5ap_rep_enc_partOptionalsNULL) + assert.Nil(t, err, "error not expected decoding test data") + + encpart := ktestMakeSampleAPReqEncPart() + encpart.SequenceNumber = 0 + encpart.Subkey = types.EncryptionKey{} + + b, err := encpart.marshal() + assert.Nil(t, err, "enc part marshal error not expected") + assert.Equal(t, want, b) +} + +func TestAprepMarshal(t *testing.T) { + t.Parallel() + + want, err := hex.DecodeString(testdata.MarshaledKRB5ap_rep) + assert.Nil(t, err, "error not expected decoding test data") + + aprep := ktestMakeSampleApRep() + + b, err := aprep.marshal() + assert.Nil(t, err, "enc part marshal error not expected") + assert.Equal(t, want, b) +} diff --git a/agents/go-common/go-gssapi/krb5/context_token.go b/agents/go-common/go-gssapi/krb5/context_token.go new file mode 100644 index 0000000000..838193eb52 --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/context_token.go @@ -0,0 +1,232 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +package krb5 + +/* + * Derived from github.com/jcmturner/gokrb5/v8/spnego/krb5Token.go + * + * The modified version adds functionality to marshal an APReq message + * to be used as part of a mutually-authenticated GSSAPI security + * context; verification is moved out. + */ + +import ( + "crypto/md5" + "encoding/binary" + "encoding/hex" + "fmt" + "net" + + "github.com/jcmturner/gofork/encoding/asn1" + "github.com/jcmturner/gokrb5/v8/asn1tools" + + "github.com/golang-auth/go-gssapi/v2" + "github.com/golang-auth/go-gssapi/v2/common" + "github.com/jcmturner/gokrb5/v8/messages" +) + +// GSSAPI KRB5 MechToken IDs. +const ( + tokenIDKrbAPReq = "0100" + tokenIDKrbAPRep = "0200" + tokenIDKrbError = "0300" +) + +// kRB5Token context token implementation for GSSAPI. +type kRB5Token struct { + oID asn1.ObjectIdentifier + tokID []byte + aPReq *messages.APReq + aPRep *aPRep + kRBError *messages.KRBError +} + +// marshal a KRB5Token into a slice of bytes. +func (m *kRB5Token) marshal() (outTok []byte, err error) { + // Create the header + b, _ := asn1.Marshal(m.oID) + b = append(b, m.tokID...) + var tb []byte + switch hex.EncodeToString(m.tokID) { + case tokenIDKrbAPReq: + tb, err = m.aPReq.Marshal() + if err != nil { + err = fmt.Errorf("gssapi: error marshalling AP-REQ for MechToken: %v", err) + } + case tokenIDKrbAPRep: + tb, err = m.aPRep.marshal() + if err != nil { + err = fmt.Errorf("gssapi: error marshalling AP-REP for MechToken: %v", err) + } + case tokenIDKrbError: + tb, err = m.kRBError.Marshal() + if err != nil { + err = fmt.Errorf("gssapi: error marshalling KRB-ERROR for MechToken: %v", err) + } + } + if err != nil { + return + } + b = append(b, tb...) + + outTok = asn1tools.AddASNAppTag(b, 0) + return +} + +// unmarshal a KRB5Token. +func (m *kRB5Token) unmarshal(b []byte) error { + m.aPReq = nil + m.aPRep = nil + m.kRBError = nil + + var oid asn1.ObjectIdentifier + r, err := asn1.UnmarshalWithParams(b, &oid, fmt.Sprintf("application,explicit,tag:%v", 0)) + if err != nil { + return fmt.Errorf("gssapi: error unmarshalling KRB5Token OID: %v", err) + } + if !oid.Equal(oID()) { + return fmt.Errorf("gssapi: error unmarshalling KRB5Token, OID is %s not %s", oid.String(), oID().String()) + } + m.oID = oid + if len(r) < 2 { + return fmt.Errorf("gssapi: krb5token too short") + } + m.tokID = r[0:2] + switch hex.EncodeToString(m.tokID) { + case tokenIDKrbAPReq: + var a messages.APReq + err = a.Unmarshal(r[2:]) + if err != nil { + return fmt.Errorf("gssapi: error unmarshalling KRB5Token AP_REQ: %v", err) + } + m.aPReq = &a + case tokenIDKrbAPRep: + var a aPRep + err = a.unmarshal(r[2:]) + if err != nil { + return fmt.Errorf("gssapi: error unmarshalling KRB5Token AP_REP: %v", err) + } + m.aPRep = &a + case tokenIDKrbError: + var a messages.KRBError + err = a.Unmarshal(r[2:]) + if err != nil { + return fmt.Errorf("gssapi: error unmarshalling KRB5Token KRBError: %v", err) + } + m.kRBError = &a + } + return nil +} + +// Create the GSSAPI checksum for the authenticator. This isn't really +// a checksum, it is a way to carry GSSAPI level context information in +// the Kerberos AP-RREQ message. See RFC 4121 § 4.1.1 +func newAuthenticatorChksum(flags gssapi.ContextFlag, cb *common.ChannelBinding) []byte { + // 24 octet minimum length, up to and including context-establishment flags + a := make([]byte, 24) + + // 4-byte length of "channel binding" info, always 16 bytes + binary.LittleEndian.PutUint32(a[:4], 16) + + // Octets 4..19: Channel binding info + if cb != nil { + copy(a[4:20], cbChecksum(cb)) + } + + // Context-establishment flags + binary.LittleEndian.PutUint32(a[20:24], uint32(flags)) + + return a +} + +func cbChecksum(cb *common.ChannelBinding) []byte { + bufSz := 5*4 + len(cb.Data) // 5 x 32 bit length fields plus the data + + // .. plus the length of the address types, if not null + for _, addr := range []net.Addr{cb.InitiatorAddr, cb.AcceptorAddr} { + if addr == nil { + continue + } + + switch c := addr.(type) { + case *net.IPAddr: + bufSz += ipLength(c.IP) + case *net.TCPAddr: + bufSz += ipLength(c.IP) + case *net.UDPAddr: + bufSz += ipLength(c.IP) + case *net.UnixAddr: + bufSz += len(c.Name) + } + } + + buf := make([]byte, 0, bufSz) + + // write the address types and address data + for _, addr := range []net.Addr{cb.InitiatorAddr, cb.AcceptorAddr} { + addrType := 0 + addrData := []byte{} + + if addr != nil { + switch c := addr.(type) { + case *net.IPAddr: + addrType = int(common.GssAddrFamilyINET) + addrData = ipData(c.IP) + case *net.TCPAddr: + addrType = int(common.GssAddrFamilyINET) + addrData = ipData(c.IP) + case *net.UDPAddr: + addrType = int(common.GssAddrFamilyINET) + addrData = ipData(c.IP) + case *net.UnixAddr: + addrType = int(common.GssAddrFamilyLOCAL) + addrData = []byte(c.Name) + } + } + + // write little endian 32-bit address type and address size + bufTmp := [8]byte{} + binary.LittleEndian.PutUint32(bufTmp[:], uint32(addrType)) + binary.LittleEndian.PutUint32(bufTmp[4:], uint32(len(addrData))) + buf = append(buf, bufTmp[:]...) + + // write the address data + buf = append(buf, addrData...) + } + + // write the data + bufTmp := [4]byte{} + binary.LittleEndian.PutUint32(bufTmp[:], uint32(len(cb.Data))) + buf = append(buf, bufTmp[:]...) + buf = append(buf, cb.Data...) + + fmt.Printf("Channel binding data (%d bytes):\n% x\n", len(buf), buf) + + hashed := md5.Sum(buf) + return hashed[:] +} + +func ipLength(addr net.IP) int { + if addr.To4() != nil { + return 4 + } + if addr.To16() != nil { + return 16 + } + + return 0 +} + +func ipData(addr net.IP) (ret net.IP) { + if ret = addr.To4(); ret != nil { + return ret + } + + if ret = addr.To16(); ret != nil { + return ret + } + + return nil +} diff --git a/agents/go-common/go-gssapi/krb5/context_token_test.go b/agents/go-common/go-gssapi/krb5/context_token_test.go new file mode 100644 index 0000000000..dd63fa9b96 --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/context_token_test.go @@ -0,0 +1,162 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +package krb5 + +import ( + "encoding/hex" + "testing" + + "github.com/jcmturner/gokrb5/v8/iana/msgtype" + "github.com/stretchr/testify/assert" +) + +const ( + // GSSAPI tokens encapsulating the apreq/apreq/krberror test vectors from the MIT Kerberos V source + // see krb-1.19.1/src/tests/asn.1/reference_encode.out + KRB5TokenApreqHex = "6081AD06092a864886f71201020201006E819D30819AA003020105A10302010EA207030500FEDCBA98A35E615C305AA003020105A1101B0E415448454E412E4D49542E454455A21A3018A003020101A111300F1B066866747361691B056578747261A3253023A003020100A103020105A21704156B726241534E2E312074657374206D657373616765A4253023A003020100A103020105A21704156B726241534E2E312074657374206D657373616765" + KRB5TOKENAprepHex = "604206092a864886f71201020202006F333031A003020105A10302010FA2253023A003020100A103020105A21704156B726241534E2E312074657374206D657373616765" + KRB5TOKENKrberrorHex = "6081ca06092a864886f71201020203007E81BA3081B7A003020105A10302011EA211180F31393934303631303036303331375AA305020301E240A411180F31393934303631303036303331375AA505020301E240A60302013CA7101B0E415448454E412E4D49542E454455A81A3018A003020101A111300F1B066866747361691B056578747261A9101B0E415448454E412E4D49542E454455AA1A3018A003020101A111300F1B066866747361691B056578747261AB0A1B086B72623564617461AC0A04086B72623564617461" + AuthChksum = "100000000000000000000000000000000000000030000000" +) + +func TestKRB5TokenApreq_Unmarshal(t *testing.T) { + t.Parallel() + b, err := hex.DecodeString(KRB5TokenApreqHex) + if err != nil { + t.Fatalf("Error decoding KRB5Token hex: %v", err) + } + var mt kRB5Token + err = mt.unmarshal(b) + if err != nil { + t.Fatalf("Error unmarshalling KRB5Token: %v", err) + } + assert.Equal(t, oID(), mt.oID, "KRB5Token OID not as expected.") + assert.Equal(t, []byte{1, 0}, mt.tokID, "TokID not as expected") + assert.NotNil(t, mt.aPReq) + assert.Nil(t, mt.aPRep) + assert.Nil(t, mt.kRBError) + assert.Equal(t, msgtype.KRB_AP_REQ, mt.aPReq.MsgType, "KRB5Token AP_REQ does not have the right message type.") + assert.Equal(t, int32(0), mt.aPReq.EncryptedAuthenticator.EType, "Authenticator within AP_REQ does not have the etype expected.") + assert.Equal(t, 5, mt.aPReq.EncryptedAuthenticator.KVNO, "Authenticator within AP_REQ does not have the KVNO expected.") + assert.Equal(t, []byte("krbASN.1 test message"), mt.aPReq.EncryptedAuthenticator.Cipher, "Authenticator within AP_REQ does not have the ciphertext expected.") +} + +func TestKRB5TokenAprep_Unmarshal(t *testing.T) { + t.Parallel() + b, err := hex.DecodeString(KRB5TOKENAprepHex) + if err != nil { + t.Fatalf("Error decoding KRB5Token hex: %v", err) + } + var mt kRB5Token + err = mt.unmarshal(b) + if err != nil { + t.Fatalf("Error unmarshalling KRB5Token: %v", err) + } + assert.Equal(t, oID(), mt.oID, "KRB5Token OID not as expected.") + assert.Equal(t, []byte{2, 0}, mt.tokID, "TokID not as expected") + assert.Nil(t, mt.aPReq) + assert.NotNil(t, mt.aPRep) + assert.Nil(t, mt.kRBError) + assert.Equal(t, msgtype.KRB_AP_REP, mt.aPRep.MsgType, "KRB5Token AP_REP does not have the right message type.") + assert.NotNil(t, mt.aPRep.EncPart, "Nil AP_REP encrypted part") + assert.NotEmpty(t, mt.aPRep.EncPart, "Empty AP_REP encrypted part") + assert.Equal(t, int32(0), mt.aPRep.EncPart.EType, "Authenticator within AP_REP does not have the etype expected.") + assert.Equal(t, 5, mt.aPRep.EncPart.KVNO, "Authenticator within AP_REP does not have the KVNO expected.") + assert.Equal(t, []byte("krbASN.1 test message"), mt.aPRep.EncPart.Cipher, "Authenticator within AP_REP does not have the ciphertext expected.") +} + +func TestKRB5TokenKrberror_Unmarshal(t *testing.T) { + t.Parallel() + b, err := hex.DecodeString(KRB5TOKENKrberrorHex) + if err != nil { + t.Fatalf("Error decoding KRB5Token hex: %v", err) + } + var mt kRB5Token + err = mt.unmarshal(b) + if err != nil { + t.Fatalf("Error unmarshalling KRB5Token: %v", err) + } + assert.Equal(t, oID(), mt.oID, "KRB5Token OID not as expected.") + assert.Equal(t, []byte{3, 0}, mt.tokID, "TokID not as expected") + assert.Nil(t, mt.aPReq) + assert.Nil(t, mt.aPRep) + assert.NotNil(t, mt.kRBError) + assert.Equal(t, msgtype.KRB_ERROR, mt.kRBError.MsgType, "KRB5Token KRB_ERROR does not have the right message type.") + assert.Equal(t, int32(SampleError), mt.kRBError.ErrorCode, "KRB5Token KRB_ERROR has the wrong error code.") + assert.Equal(t, "ATHENA.MIT.EDU", mt.kRBError.Realm, "KRB5Token KRB_ERROR has the wrong realm.") + assert.Equal(t, SampleData, mt.kRBError.EText, "KRB5Token KRB_ERROR has the wrong error test.") +} + +func TestKrb5TokenApreq_Marshal(t *testing.T) { + t.Parallel() + + apreq := ktestMakeSampleApReq() + + mt := kRB5Token{ + oID: oID(), + tokID: []byte{1, 0}, + aPReq: &apreq, + } + + tok, err := mt.marshal() + if err != nil { + t.Fatalf("Error marshalling KRB5Token: %s", err) + } + + ref, err := hex.DecodeString(KRB5TokenApreqHex) + if err != nil { + t.Fatalf("Error decoding KRB5Token hex: %s", err) + } + + assert.Equal(t, ref, tok) +} + +func TestKrb5TokenAprep_Marshal(t *testing.T) { + t.Parallel() + + aprep := ktestMakeSampleApRep() + + mt := kRB5Token{ + oID: oID(), + tokID: []byte{2, 0}, + aPRep: &aprep, + } + + tok, err := mt.marshal() + if err != nil { + t.Fatalf("Error marshalling KRB5Token: %s", err) + } + + ref, err := hex.DecodeString(KRB5TOKENAprepHex) + if err != nil { + t.Fatalf("Error decoding KRB5Token hex: %s", err) + } + + assert.Equal(t, ref, tok) +} + +func TestKrb5TokenKrberror_Marshal(t *testing.T) { + t.Parallel() + + krberr := ktestMakeSampleError() + + mt := kRB5Token{ + oID: oID(), + tokID: []byte{3, 0}, + kRBError: &krberr, + } + + tok, err := mt.marshal() + if err != nil { + t.Fatalf("Error marshalling KRB5Token: %s", err) + } + + ref, err := hex.DecodeString(KRB5TOKENKrberrorHex) + if err != nil { + t.Fatalf("Error decoding KRB5Token hex: %s", err) + } + + assert.Equal(t, ref, tok) +} diff --git a/agents/go-common/go-gssapi/krb5/credentials_test.go b/agents/go-common/go-gssapi/krb5/credentials_test.go new file mode 100644 index 0000000000..60355c81be --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/credentials_test.go @@ -0,0 +1,46 @@ +package krb5 + +import ( + "testing" + + "github.com/jcmturner/gokrb5/v8/config" +) + +func TestSplitPrincipal(t *testing.T) { + username, realm, err := splitPrincipal("alice/admin@EXAMPLE.COM", "") + if err != nil { + t.Fatal(err) + } + if username != "alice/admin" || realm != "EXAMPLE.COM" { + t.Fatalf("unexpected principal: %q %q", username, realm) + } +} + +func TestSplitPrincipalUsesDefaultRealm(t *testing.T) { + username, realm, err := splitPrincipal("alice", "EXAMPLE.COM") + if err != nil { + t.Fatal(err) + } + if username != "alice" || realm != "EXAMPLE.COM" { + t.Fatalf("unexpected principal: %q %q", username, realm) + } +} + +func TestConfigureServiceRealmSeparatesExplicitRealm(t *testing.T) { + configuration := config.New() + service := configureServiceRealm(configuration, "zookeeper/zk.example.com@ZK.EXAMPLE.COM") + if service != "zookeeper/zk.example.com" { + t.Fatalf("service = %q", service) + } + if configuration.DomainRealm["zk.example.com"] != "ZK.EXAMPLE.COM" { + t.Fatalf("domain realm mapping = %#v", configuration.DomainRealm) + } +} + +func TestConfigureServiceRealmLeavesImplicitRealmResolution(t *testing.T) { + configuration := config.New() + service := configureServiceRealm(configuration, "hive/hs2.example.com") + if service != "hive/hs2.example.com" || len(configuration.DomainRealm) != 0 { + t.Fatalf("service=%q domain realms=%#v", service, configuration.DomainRealm) + } +} diff --git a/agents/go-common/go-gssapi/krb5/default_paths_unix.go b/agents/go-common/go-gssapi/krb5/default_paths_unix.go new file mode 100644 index 0000000000..f3722a7a96 --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/default_paths_unix.go @@ -0,0 +1,20 @@ +//go:build !windows + +package krb5 + +import ( + "fmt" + "os" +) + +func defaultKrbConfFile() string { + return "/etc/krb5.conf" +} + +func defaultKrbCCFile() string { + return fmt.Sprintf("/tmp/krb5cc_%d", os.Getuid()) +} + +func defaultKrbKTFile() string { + return fmt.Sprintf("/var/kerberos/krb5/user/%d/client.keytab", os.Getuid()) +} diff --git a/agents/go-common/go-gssapi/krb5/default_paths_windows.go b/agents/go-common/go-gssapi/krb5/default_paths_windows.go new file mode 100644 index 0000000000..803429aa4e --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/default_paths_windows.go @@ -0,0 +1,23 @@ +//go:build windows + +package krb5 + +import ( + "os" + "path/filepath" +) + +func defaultKrbConfFile() string { + if windowsDirectory := os.Getenv("WINDIR"); windowsDirectory != "" { + return filepath.Join(windowsDirectory, "krb5.ini") + } + return `C:\Windows\krb5.ini` +} + +func defaultKrbCCFile() string { + return "" +} + +func defaultKrbKTFile() string { + return "" +} diff --git a/agents/go-common/go-gssapi/krb5/keyinfo.go b/agents/go-common/go-gssapi/krb5/keyinfo.go new file mode 100644 index 0000000000..3376ba0e9a --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/keyinfo.go @@ -0,0 +1,117 @@ +package krb5 + +import ( + "github.com/jcmturner/gokrb5/v8/crypto" +) + +func keySSF(keyType int32) uint { + // From MIT Kerberos 1.16 (src/lib/crypto/krb/etypes.c) + /* + ENCTYPE_DES3_CBC_RAW 112 + ENCTYPE_DES3_CBC_SHA1 112 + ENCTYPE_ARCFOUR_HMAC 64 + ENCTYPE_ARCFOUR_HMAC_EXP 40 + ENCTYPE_AES128_CTS_HMAC_SHA1_96 128 + ENCTYPE_AES256_CTS_HMAC_SHA1_96 256 + ENCTYPE_CAMELLIA128_CTS_CMAC 128 + ENCTYPE_CAMELLIA256_CTS_CMAC 256 + ENCTYPE_AES128_CTS_HMAC_SHA256_128 128 + ENCTYPE_AES256_CTS_HMAC_SHA384_192 256 + */ + + key, _ := crypto.GetEtype(keyType) + + switch key.(type) { + case crypto.Des3CbcSha1Kd: + return 112 + case crypto.RC4HMAC: + return 64 + } + + // default to the key length in bits + return uint(key.GetKeyByteSize()) * 8 +} + +// port from MIT Kerberos 1.16 (krb5_c_encrypt_length) +func encryptedLength(keyType int32, plainTextSize uint32) uint32 { + paddingLen := paddingLength(keyType, plainTextSize) + return uint32(keyHeaderLength(keyType)) + + plainTextSize + + paddingLen + + uint32(keyTrailerLength(keyType)) +} + +// port from MIT Kerberos 1.16 (krb5int_c_padding_length) +func paddingLength(keyType int32, dataLength uint32) uint32 { + dataLength += uint32(keyHeaderLength(keyType)) + padding := uint32(keyPaddingLength(keyType)) + + if padding == 0 || (dataLength%padding) == 0 { + return 0 + } else { + return padding - (dataLength % padding) + } +} + +func keyHeaderLength(keyType int32) uint { + key, _ := crypto.GetEtype(keyType) + + switch key.(type) { + case crypto.Des3CbcSha1Kd: + return uint(key.GetCypherBlockBitLength()) / 8 + case crypto.RC4HMAC: + return uint(key.GetHMACBitLength()/8 + key.GetConfounderByteSize()) + case crypto.Aes128CtsHmacSha96: + return uint(key.GetCypherBlockBitLength()) / 8 + case crypto.Aes128CtsHmacSha256128: + return uint(key.GetCypherBlockBitLength()) / 8 + case crypto.Aes256CtsHmacSha96: + return uint(key.GetCypherBlockBitLength()) / 8 + case crypto.Aes256CtsHmacSha384192: + return uint(key.GetCypherBlockBitLength()) / 8 + } + + return 0 +} + +func keyPaddingLength(keyType int32) uint { + key, _ := crypto.GetEtype(keyType) + + switch key.(type) { + case crypto.Des3CbcSha1Kd: + return uint(key.GetCypherBlockBitLength()) / 8 + case crypto.RC4HMAC: + return 0 + case crypto.Aes128CtsHmacSha96: + return 0 + case crypto.Aes128CtsHmacSha256128: + return 0 + case crypto.Aes256CtsHmacSha96: + return 0 + case crypto.Aes256CtsHmacSha384192: + return 0 + } + + return 0 +} + +func keyTrailerLength(keyType int32) uint { + key, _ := crypto.GetEtype(keyType) + + switch key.(type) { + case crypto.Des3CbcSha1Kd: + return uint(key.GetHMACBitLength()) / 8 + case crypto.RC4HMAC: + return 0 + case crypto.Aes128CtsHmacSha96: + return uint(key.GetHMACBitLength()) / 8 + case crypto.Aes128CtsHmacSha256128: + return uint(key.GetHMACBitLength()) / 8 + case crypto.Aes256CtsHmacSha96: + return uint(key.GetHMACBitLength()) / 8 + case crypto.Aes256CtsHmacSha384192: + return uint(key.GetHMACBitLength()) / 8 + } + + return 0 +} diff --git a/agents/go-common/go-gssapi/krb5/keyinfo_test.go b/agents/go-common/go-gssapi/krb5/keyinfo_test.go new file mode 100644 index 0000000000..7b56d31917 --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/keyinfo_test.go @@ -0,0 +1,55 @@ +package krb5 + +import ( + "fmt" + "math/rand" + "testing" + + "github.com/jcmturner/gokrb5/v8/crypto" + "github.com/jcmturner/gokrb5/v8/iana/etypeID" + "github.com/jcmturner/gokrb5/v8/iana/keyusage" + "github.com/stretchr/testify/assert" +) + +func TestEncryptedLength(t *testing.T) { + var tests = []struct { + keyTypeName string + }{ + {"rc4-hmac"}, + {"des3-cbc-sha1-kd"}, + {"aes128-cts-hmac-sha1-96"}, + {"aes256-cts-hmac-sha1-96"}, + {"aes128-cts-hmac-sha256-128"}, + {"aes256-cts-hmac-sha384-192"}, + } + + var plainTextLengths = []int{0, 1, 2, 8, 16, 64, 128, 256, 1024, 3000, 4000} + + for _, tt := range tests { + for _, plainLength := range plainTextLengths { + plainData := make([]byte, plainLength) + _, _ = rand.Read(plainData) + + name := fmt.Sprintf("%s (%d byte plaintext)", tt.keyTypeName, plainLength) + + t.Run(name, func(t *testing.T) { + etypeID := etypeID.EtypeSupported(tt.keyTypeName) + assert.Positive(t, etypeID, "key type %s should be supported") + + etype, _ := crypto.GetEtype(etypeID) + key, err := GenerateBaseKey(etype) + assert.NoError(t, err, "failed to generate encryption key") + + _, cipherText, err := etype.EncryptMessage(key.KeyValue, plainData, keyusage.GSSAPI_INITIATOR_SEAL) + assert.NoError(t, err, "failed to encrypt message") + + expectedLength := encryptedLength(etypeID, uint32(plainLength)) + assert.Equal(t, expectedLength, uint32(len(cipherText))) + + _ = cipherText + }) + } + } + + assert.True(t, true) +} diff --git a/agents/go-common/go-gssapi/krb5/krb5.go b/agents/go-common/go-gssapi/krb5/krb5.go new file mode 100644 index 0000000000..1573676338 --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/krb5.go @@ -0,0 +1,1122 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +/* +Package krb5 provides the pure-Go implementation of the GSS-API interface +Kerberos mechanism (RFC 4121). + +Normally, this package would be imported by application code (eg. in its +main package) in order to register the Kerberos V mechanism. Application +code that uses GSS-API would import the generic github.com/golang-auth/go-gssapi/v2 +package instead and obtain a handle to this mechanism from the registry by +passing the name "kerberos_v5" or the OID "1.2.840.113554.1.2.2", eg : + +# Main Package + +A relatively high-level package should include the mechanisms that the +application is to use. The idea is that the mechanisms that are supported +in an application can be managed in one place, without changing any of +the lower level code that uses the GSS-API functionality: + + package main + import ( + _ "github.com/golang-auth/go-gssapi/v2/krb5" + "stuff" + ) + + stuff.doStuff("kerberos_v5") + +# Implementation package + +The package that uses GSS-API should accept the name of the mechanism to +use, and use that name to obtain an instance of that mechanism-specific +implementation: + + package stuff + import "github.com/golang-auth/go-gssapi/v2" + + func doStuff(mech) { + ctx := gssapi.NewMech(mech) + ... + } + +# See Also + +github.com/golang-auth/go-gssapi/v2 +*/ +package krb5 + +import ( + "crypto/rand" + "encoding/binary" + "encoding/hex" + "errors" + "fmt" + "math" + "math/big" + "os" + "strings" + "time" + + "github.com/jcmturner/gofork/encoding/asn1" + "github.com/jcmturner/gokrb5/crypto/etype" + "github.com/jcmturner/gokrb5/iana/etypeID" + "github.com/jcmturner/gokrb5/v8/client" + "github.com/jcmturner/gokrb5/v8/config" + "github.com/jcmturner/gokrb5/v8/credentials" + "github.com/jcmturner/gokrb5/v8/crypto" + "github.com/jcmturner/gokrb5/v8/iana/chksumtype" + ianaerrcode "github.com/jcmturner/gokrb5/v8/iana/errorcode" + ianaflags "github.com/jcmturner/gokrb5/v8/iana/flags" + "github.com/jcmturner/gokrb5/v8/keytab" + "github.com/jcmturner/gokrb5/v8/messages" + "github.com/jcmturner/gokrb5/v8/types" + + "github.com/golang-auth/go-gssapi/v2" + "github.com/golang-auth/go-gssapi/v2/common" +) + +func init() { + gssapi.Register("kerberos_v5", NewKrb5Mech) + gssapi.Register("1.2.840.113554.1.2.2}", NewKrb5Mech) +} + +// ClockSkew defines the maximum tolerable difference between the two peers +// of a GSS-API context, and defaults to 10 seconds. Increase this number if +// there is poor syncronisation between client and server nodes. Decrease +// the value to enhance security where there is good synchronisation. +var ClockSkew = time.Second * 10 + +type acceptorISN int + +// These constants define how the Acceptor initial sequence number is derived +// when the context does not use mutual authentication. In this case, the +// Acceptor does not have the opportunity to communicate its own sequence number +// to the Initiator. Two different schemes are in use: +// +// 1. Acceptor uses the Initiator's initial sequence number +// +// 2. The Acceptor ISN is zero +// +// The default is (1), but may be changed to (2) by setting AcceptorISN to +// the value DefaultAcceptorISNZero. +const ( + // DefaultAcceptorISNInitiator is the acceptor ISN policy that uses the Initiator's initial sequence number + // as the Acceptor ISN when not performing mutual authentication. Use this for compatibility with MIT. + DefaultAcceptorISNInitiator acceptorISN = iota + + // DefaultAcceptorISNZero is the acceptor ISN policy that uses zero as the Acceptor ISN when not + // performing mutual authentication. Use this for compatibility with Heimdal. + DefaultAcceptorISNZero +) + +// AcceptorISN holds the default Acceptor-Initial-Sequence derivation policy +// for contexts not using mutual authentication. The default provides +// compatibility with MIT Kerberos. +// Set this to DefaultAcceptorISNZero for compatibility with Heimdal Kerberos. +var AcceptorISN acceptorISN = DefaultAcceptorISNInitiator + +// krb5Mech is the implementation of the Mech interface for the +// Kerberos V mechanism +type Krb5Mech struct { + krbClient *client.Client + isInitiator bool + isEstablished bool + waitingForMutual bool + service string + channelBinding *common.ChannelBinding + ticket *messages.Ticket + sessionKey *types.EncryptionKey + clientCTime time.Time + clientCusec int + sessionFlags gssapi.ContextFlag + requestFlags gssapi.ContextFlag + ourSequenceNumber uint64 + theirSequenceNumber uint64 + initiatorSubKey *types.EncryptionKey + acceptorSubKey *types.EncryptionKey + peerName string + clientOptions *ClientOptions +} + +// ClientOptions configures initiator credentials without relying on process +// environment variables. +type ClientOptions struct { + ConfigPath string + CCachePath string + KeytabPath string + Principal string + Password string + UseCCache bool + UseKeytab bool + DisablePAFXFAST bool +} + +// NewMech returns a new Kerberos V mechanism context. This function is +// registered with the GSS-API registry and is used by gssapi.NewMech() +// when a caller requests an instance of the "kerberos_v5" mechanism. +func NewKrb5Mech() gssapi.Mech { + return &Krb5Mech{} +} + +// NewKrb5MechWithOptions returns a Kerberos mechanism with connection-scoped +// credential configuration. +func NewKrb5MechWithOptions(options ClientOptions) gssapi.Mech { + copy := options + return &Krb5Mech{clientOptions: ©} +} + +func oID() asn1.ObjectIdentifier { + return asn1.ObjectIdentifier{1, 2, 840, 113554, 1, 2, 2} +} + +// IsEstablished returns false until the Krb5Mech context has been negotiated +// and the context is ready to use for exchanging messages. +func (m Krb5Mech) IsEstablished() bool { + return m.isEstablished +} + +// ContextFlags returns the subset of requested context flags that are available +// and may change during establishmane of the context. The Initiator and +// Acceptor should examine the flags before using the context for message +// exchange, to verify that the state of the context matches the appliation +// security requirements. +func (m Krb5Mech) ContextFlags() (f gssapi.ContextFlag) { + return m.sessionFlags +} + +// SSF returns the Security Strength Factor of the channel established +// by the security context. For Kerberos V, this depends on the type of +// key being used to secure the channel. +func (m Krb5Mech) SSF() uint { + var key types.EncryptionKey + switch { + case m.acceptorSubKey != nil: + key = *m.acceptorSubKey + case m.initiatorSubKey != nil: + key = *m.initiatorSubKey + default: + key = *m.sessionKey + } + + return keySSF(key.KeyType) +} + +// From MIT Kerberos 1.16 (src/lib/gssapi/krb5/wrap_size_limit.c) +func (m Krb5Mech) WrapSizeLimit(requestedOutputSize uint32, confidentiality bool) uint32 { + var keyType int32 + switch { + case m.acceptorSubKey != nil: + keyType = m.acceptorSubKey.KeyType + case m.initiatorSubKey != nil: + keyType = m.initiatorSubKey.KeyType + default: + keyType = m.sessionKey.KeyType + } + + sz := requestedOutputSize + + if confidentiality { + // try decreasing message lengths until the encrypted length including the + // header will fit the requested size + for sz > 0 { + wrapSize := 16 + encryptedLength(keyType, sz) + if wrapSize <= requestedOutputSize { + break + } + + sz-- + } + + // account for the header + if sz > 16 { + sz -= 16 + } else { + sz = 0 + } + } else { + key, _ := crypto.GetEtype(keyType) + cksumSize := key.GetHMACBitLength() / 8 + + if sz < uint32(16+cksumSize) { + sz = 0 + } else { + sz -= uint32(16 + cksumSize) + } + } + + return sz +} + +// Accept is used by a GSS-API Acceptor to begin context +// negotiation with a remote Initiator. +// +// If provided, serviceName is the name of a service principal +// to use from the keytab. If not supplied, any principal in the +// keytab matching the request will be used. +// +// See: RFC 4121 § 4.1 +func (m *Krb5Mech) Accept(serviceName string) (err error) { + m.isEstablished = false + m.waitingForMutual = false + m.isInitiator = false + m.service = serviceName + + // Stash the subset of the request flags that we can support, except mutual + // which we won't know about until we receive a token + m.sessionFlags = gssapi.ContextFlagConf | gssapi.ContextFlagInteg | + gssapi.ContextFlagReplay | gssapi.ContextFlagSequence + + return +} + +// Initiate is used by a GSS-API Initiator to start the +// context negotiation process with a remote Acceptor. +// +// serverName is the name of the service principal to use when +// obtaining a Kerberos ticket. +// +// flags represent the desired security properties of the context +// +// cb is the channel binding data, or nil to disable +// +// It is highly recommended to make use of mutual authentication wherever +// possible and to include replay detection: +// +// gssapi.ContextFlagMutual | gssapi.ContextFlagInteg |gssapi.ContextFlagReplay +// +// Most users should also include gssapi.ContextFlagConf to enable the use +// of message sealing. +func (m *Krb5Mech) Initiate(serviceName string, requestFlags gssapi.ContextFlag, cb *common.ChannelBinding) (err error) { + m.isEstablished = false + m.waitingForMutual = false + m.isInitiator = true + m.channelBinding = cb + + // Obtain a Kerberos ticket for the service + if err = m.krbClientInit(serviceName); err != nil { + return + } + + // Stash the subset of the request flags that we can support minus mutual until that completes + m.sessionFlags = gssapi.ContextFlagConf | gssapi.ContextFlagInteg | + gssapi.ContextFlagReplay | gssapi.ContextFlagSequence + + // requuest flags is the subset that we support of the requested flags, used in the context + // negotiation. The set we will tell the caller that we actually support is the above, + // sessionFlags which may include more than the requested set + m.requestFlags = requestFlags & (gssapi.ContextFlagConf | gssapi.ContextFlagInteg | + gssapi.ContextFlagMutual | gssapi.ContextFlagReplay | gssapi.ContextFlagSequence) + + return +} + +// Continue is called in a loop by Initiators and Acceptors after +// first calling one of Initiate or Accept. +// +// tokenIn represents a token received from the peer +// If tokenOut is non-zero, it should be send to the peer +// The caller should check the result of m.IsEstablished() to determine +// then the loop should end. +func (m *Krb5Mech) Continue(tokenIn []byte) (tokenOut []byte, err error) { + if m.isEstablished { + err = nil + return + } + + if m.isInitiator { + return m.continueInitiator(tokenIn) + } else { + return m.continueAcceptor(tokenIn) + } +} + +func (m *Krb5Mech) continueInitiator(tokenIn []byte) (tokenOut []byte, err error) { + // first time, create the first context-establishment token + // + if len(tokenIn) == 0 { + // Create a Kerberos AP-REQ message with GSSAPI checksum + var apreq messages.APReq + apreq, err = m.getAPReqMessage() + if err != nil { + return + } + + // Create the GSSAPI token + tb, _ := hex.DecodeString(tokenIDKrbAPReq) + gssToken := kRB5Token{ + oID: oID(), + tokID: tb, + aPReq: &apreq, + } + + tokenOut, err = gssToken.marshal() + if err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + + // we need another round if we're doing mutual auth - we will receive an AP-REP from the server + if m.requestFlags&gssapi.ContextFlagMutual == 0 { + m.isEstablished = true + + // if there is no mutual auth, we can't tell the client what our initial sequence number is + // MIT and Microsoft use the client's ISN so let's do that, unless we're in Heimdal mode + // see https://bugs.openjdk.java.net/browse/JDK-8201814 + switch AcceptorISN { + case DefaultAcceptorISNInitiator: + m.theirSequenceNumber = m.ourSequenceNumber + case DefaultAcceptorISNZero: + m.theirSequenceNumber = 0 + default: + err = fmt.Errorf("gssapi: unknown acceptor-initial-sequence-number policy configured") + return + } + } else { + m.waitingForMutual = true + } + + return + } + + // called again due to a previous ContinueNeeded result ?.. + if !m.waitingForMutual { + err = fmt.Errorf("gssapi: context is not ready, call Start to initialize a new context") + return + } + + // unmarshal the GSSAPI token + gssToken := kRB5Token{} + if err = gssToken.unmarshal(tokenIn); err != nil { + return + } + + if gssToken.kRBError != nil { + err = fmt.Errorf("gssapi: %s", gssToken.kRBError.Error()) + return + } + + if gssToken.aPRep == nil { + err = errors.New("gssapi: GSSAPI token does not contain AP-REP message") + return + } + + // decrypt/verify the private part of the AP-REP message + msg, err := gssToken.aPRep.decryptEncPart(*m.sessionKey) + if err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + + // stash their sequence number and subkey for use in GSS Wrap/Unwrap + m.theirSequenceNumber = uint64(msg.SequenceNumber) + if msg.Subkey.KeyType != 0 { + m.acceptorSubKey = &msg.Subkey + } + + // check the response has the same time values as the request + // Note - we can't use time.Equal() as m.clientCTime has a monotomic clock value and + // which causes the equality to fail + if !(msg.CTime.Unix() == m.clientCTime.Unix() && msg.Cusec == m.clientCusec) { + err = fmt.Errorf("gssapi: mutual authentication failed") + return + } + + // we're done! + m.isEstablished = true + m.waitingForMutual = false + m.sessionFlags |= gssapi.ContextFlagMutual + + return tokenOut, nil +} + +func (m *Krb5Mech) continueAcceptor(tokenIn []byte) (tokenOut []byte, err error) { + // try to unmarshal the token + gssInToken := kRB5Token{} + if err = gssInToken.unmarshal(tokenIn); err != nil { + return + } + + if gssInToken.kRBError != nil { + err = fmt.Errorf("gssapi: %s", gssInToken.kRBError.Error()) + return + } + + // RFC says: must return a KRBError message to the client if the token ID was invalid + // Note sure other implementatios really do this + if gssInToken.kRBError == nil && gssInToken.aPReq == nil && gssInToken.aPRep == nil { + tokenOut, err = mkGssErrKrbCode(ianaerrcode.KRB_AP_ERR_MSG_TYPE, "gss accept failed") + return + } + + // avoid crash if GSSAPI token isn't an initial token + if gssInToken.aPReq == nil { + err = errors.New("gssapi: GSSAPI token does not contain AP-REQ message") + return + } + + ktFile := krbKtFile() + + err, krbErr := verifyAPReq(ktFile, gssInToken.aPReq, ClockSkew) + if err != nil { + tokenOut, err = mkGssErrFromKrbErr(krbErr.(messages.KRBError)) + return + } + + // stash the sequence number for use in GSS Wrap + // Authenticator.SeqNumber is actually a 32 bit number (in the protocol), so the cast here is safe + m.theirSequenceNumber = uint64(gssInToken.aPReq.Authenticator.SeqNumber) + + // stash the APReq time flags for use in mutual authentication + m.clientCTime = gssInToken.aPReq.Authenticator.CTime + m.clientCusec = gssInToken.aPReq.Authenticator.Cusec + + // stash the session key and ticket + m.ticket = &gssInToken.aPReq.Ticket + m.sessionKey = &gssInToken.aPReq.Ticket.DecryptedEncPart.Key + + // stash the initiator subkey if there is one + if gssInToken.aPReq.Authenticator.SubKey.KeyType != 0 { + m.initiatorSubKey = &gssInToken.aPReq.Authenticator.SubKey + } + + // get the context-establishment flags from the authenticator + requestedFlags := binary.LittleEndian.Uint32(gssInToken.aPReq.Authenticator.Cksum.Checksum[20:24]) + m.sessionFlags &= gssapi.ContextFlag(requestedFlags) + + // stash the client's principal name + m.peerName = fmt.Sprintf("%s@%s", + gssInToken.aPReq.Ticket.DecryptedEncPart.CName.PrincipalNameString(), + gssInToken.aPReq.Ticket.DecryptedEncPart.CRealm) + + // if the client requested mutual authentication, send them an AP-REP message + if types.IsFlagSet(&gssInToken.aPReq.APOptions, ianaflags.APOptionMutualRequired) { + tb, _ := hex.DecodeString(tokenIDKrbAPRep) + gssOutToken := kRB5Token{ + oID: oID(), + tokID: tb, + } + + var aprep aPRep + aprep, err = m.getAPRepMessage() + if err != nil { + return + } + gssOutToken.aPRep = &aprep + + tokenOut, err = gssOutToken.marshal() + if err != nil { + return + } + + m.sessionFlags |= gssapi.ContextFlagMutual + } else { + // if there is no mutual auth, we can't tell the client what our initial sequence number is + // MIT and Microsoft use the client's ISN so let's do that, unless we're in Heimdal mode + // see https://bugs.openjdk.java.net/browse/JDK-8201814 + switch AcceptorISN { + case DefaultAcceptorISNInitiator: + m.ourSequenceNumber = m.theirSequenceNumber + case DefaultAcceptorISNZero: + m.ourSequenceNumber = 0 + default: + err = fmt.Errorf("gssapi: unknown acceptor-initial-sequence-number policy configured") + return + } + } + + // we're done from an acceptor perspective + m.isEstablished = true + return tokenOut, nil +} + +// PeerName returns the name of the remote peer's Kerberos principal +func (m *Krb5Mech) PeerName() string { + return m.peerName +} + +// InitiatorName returns the authenticated initiator principal. +func (m *Krb5Mech) InitiatorName() string { + if m.krbClient == nil || m.krbClient.Credentials == nil { + return "" + } + username := m.krbClient.Credentials.UserName() + realm := m.krbClient.Credentials.Domain() + if username == "" || realm == "" { + return username + } + return username + "@" + realm +} + +// Dispose clears the Kerberos client and session state held by the mechanism. +func (m *Krb5Mech) Dispose() { + if m.krbClient != nil { + m.krbClient.Destroy() + } + m.krbClient = nil + m.ticket = nil + m.sessionKey = nil + m.initiatorSubKey = nil + m.acceptorSubKey = nil + m.isEstablished = false +} + +// Wrap encapsulates the payload in a GSS-API Wap oken that can be passed to the +// remote peer. The payload is sealed if confidentiality is requested, and +// signed if not. Note that the use of confidentially requires the +// gssapi.ContextFlagMutual flag to be enabled on the context. +func (m *Krb5Mech) Wrap(tokenIn []byte, confidentiality bool) (tokenOut []byte, err error) { + wt, err := m.newWrapToken(tokenIn, confidentiality) + if err != nil { + return + } + + tokenOut, err = wt.Marshal() + return +} + +// Unwrap is used to parse a token created with Wrap(). It returns the original +// payload after unsealing or verification of the signature. isSealed can be +// inspected to determine whether the payload was encrypted or only signed. +func (m *Krb5Mech) Unwrap(tokenIn []byte) (tokenOut []byte, isSealed bool, err error) { + // Unmarshall the token + wt := wrapToken{} + if err = wt.Unmarshal(tokenIn); err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + + var key types.EncryptionKey + switch { + case wt.Flags&gSSMessageTokenFlagAcceptorSubkey != 0: + if m.acceptorSubKey == nil { + err = errors.New("gssapi: acceptor subkey not negotiated, cannot unwrap message") + return + } + + key = *m.acceptorSubKey + case m.initiatorSubKey != nil: + key = *m.initiatorSubKey + default: + key = *m.sessionKey + } + + // Verify the token's integrity and get the unsealed / unsigned payload + if isSealed, err = wt.VerifyAndDecode(key, m.isInitiator); err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + + // Check the sequence number + if m.sessionFlags&gssapi.ContextFlagReplay != 0 || m.sessionFlags&gssapi.ContextFlagSequence != 0 { + if wt.SequenceNumber != m.theirSequenceNumber { + err = fmt.Errorf("gssapi: bad sequence number from peer, got %d, wanted %d", wt.SequenceNumber, m.theirSequenceNumber) + return + } + } + m.theirSequenceNumber++ + + tokenOut = wt.Payload + return tokenOut, isSealed, nil +} + +// MakeSignature creates a GSS-API MIC token, containing the signature of +// payload but not encapsulating any payload. The MIC token is passed to the +// peer separately to the payload and can be used by the peer to verify +// the integrity of that payload. +func (m *Krb5Mech) MakeSignature(payload []byte) (tokenOut []byte, err error) { + var flags gSSMessageTokenFlag + + if !m.isInitiator { + flags |= gSSMessageTokenFlagSentByAcceptor // send by acceptor + } + + // use the acceptor subkey if it was negotiated during auth + key := m.sessionKey + switch { + case m.acceptorSubKey != nil: + key = m.acceptorSubKey + flags |= gSSMessageTokenFlagAcceptorSubkey + case m.initiatorSubKey != nil: + key = m.initiatorSubKey + } + + mt := mICToken{ + Flags: flags, + SequenceNumber: m.ourSequenceNumber, + } + + if err = mt.Sign(payload, *key); err != nil { + return + } + + tokenOut, err = mt.Marshal() + return +} + +// VerifySignature checks the cryptographic signature created by a call +// to MakeSignature() on the supplied payload. +func (m *Krb5Mech) VerifySignature(payload []byte, tokenIn []byte) (err error) { + mt := mICToken{} + if err = mt.Unmarshal(tokenIn); err != nil { + return + } + + var key types.EncryptionKey + switch { + case mt.Flags&gSSMessageTokenFlagAcceptorSubkey != 0: + if m.acceptorSubKey == nil { + err = errors.New("gssapi: acceptor subkey not negotiated, cannot verify MIC") + return + } + + key = *m.acceptorSubKey + case m.initiatorSubKey != nil: + key = *m.initiatorSubKey + default: + key = *m.sessionKey + } + + if err = mt.Verify(payload, key, m.isInitiator); err != nil { + return + } + + // Check the sequence number + if m.sessionFlags&gssapi.ContextFlagReplay != 0 || m.sessionFlags&gssapi.ContextFlagSequence != 0 { + if mt.SequenceNumber != m.theirSequenceNumber { + err = fmt.Errorf("gssapi: bad sequence number from peer, got %d, wanted %d", mt.SequenceNumber, m.theirSequenceNumber) + return + } + } + m.theirSequenceNumber++ + + return nil +} + +func (m *Krb5Mech) getAPReqMessage() (apreq messages.APReq, err error) { + auth, err := types.NewAuthenticator(m.krbClient.Credentials.Domain(), m.krbClient.Credentials.CName()) + if err != nil { + err = fmt.Errorf("gssapi: generating new authenticator: %s", err) + return + } + + // MIT compatibility + auth.SeqNumber &= 0x3fffffff + + auth.Cksum = types.Checksum{ + CksumType: chksumtype.GSSAPI, + Checksum: newAuthenticatorChksum(m.requestFlags, m.channelBinding), + } + + apreq, err = messages.NewAPReq(*m.ticket, *m.sessionKey, auth) + if err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + + // set the Kerberos APREQ MUTUAL-REQUIRED option if we've been asked to perform mutual auth + if m.requestFlags&gssapi.ContextFlagMutual != 0 { + types.SetFlag(&apreq.APOptions, ianaflags.APOptionMutualRequired) + } + + // stash the sequence number for use in GSS Wrap + // Authenticator.SeqNumber is actually a 32 bit number (in the protocol), so the cast here is safe + m.ourSequenceNumber = uint64(auth.SeqNumber) + + // stash the APReq time flags for use in mutual authentication + m.clientCTime = auth.CTime + m.clientCusec = auth.Cusec + + return apreq, err +} + +func (m *Krb5Mech) getAPRepMessage() (aprep aPRep, err error) { + seq, err := rand.Int(rand.Reader, big.NewInt(math.MaxUint32)) + if err != nil { + return + } + + /* + * Work around implementation incompatibilities by not generating + * initial sequence numbers greater than 2^30. Previous MIT + * implementations use signed sequence numbers, so initial + * sequence numbers 2^31 to 2^32-1 inclusive will be rejected. + * Letting the maximum initial sequence number be 2^30-1 allows + * for about 2^30 messages to be sent before wrapping into + * "negative" numbers. + */ + seqNum := seq.Int64() & 0x3fffffff + + encPart := encAPRepPart{ + CTime: m.clientCTime, // copied from the APReq + Cusec: m.clientCusec, + SequenceNumber: seqNum, + } + + aprep, err = newAPRep(*m.ticket, *m.sessionKey, encPart) + if err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + + m.ourSequenceNumber = uint64(seqNum) + return aprep, err +} + +func (m *Krb5Mech) krbClientInit(service string) (err error) { + options := clientOptionsFromEnvironment() + if m.clientOptions != nil { + options = *m.clientOptions + } + cfgFile := normalizedKerberosPath(options.ConfigPath) + if cfgFile == "" { + cfgFile = defaultKrbConfFile() + } + cfg, err := config.Load(cfgFile) + if err != nil { + return fmt.Errorf("gssapi: loading krb5.conf: %w", err) + } + m.krbClient, err = newInitiatorClient(cfg, options) + if err != nil { + return err + } + + ticketService := configureServiceRealm(cfg, service) + tkt, key, err := m.krbClient.GetServiceTicket(ticketService) + if err != nil { + return fmt.Errorf("gssapi: getting service ticket for '%s': %s", service, err) + } + m.ticket, m.sessionKey, m.service = &tkt, &key, service + m.peerName = fmt.Sprintf("%s@%s", tkt.SName.PrincipalNameString(), tkt.Realm) + + return nil +} + +func configureServiceRealm(cfg *config.Config, service string) string { + principal, realm := types.ParseSPNString(strings.TrimSpace(service)) + if realm != "" && len(principal.NameString) > 0 { + host := strings.ToLower(strings.TrimSpace(principal.NameString[len(principal.NameString)-1])) + if host != "" { + cfg.DomainRealm[host] = realm + } + } + return principal.PrincipalNameString() +} + +func newInitiatorClient(cfg *config.Config, options ClientOptions) (*client.Client, error) { + settings := []func(*client.Settings){ + client.DisablePAFXFAST(options.DisablePAFXFAST), + } + useCCache := options.UseCCache || strings.TrimSpace(options.CCachePath) != "" + useKeytab := options.UseKeytab || strings.TrimSpace(options.KeytabPath) != "" + principal := strings.TrimSpace(options.Principal) + password := options.Password + + if useCCache { + ccachePath := normalizedKerberosPath(options.CCachePath) + if ccachePath == "" { + ccachePath = defaultKrbCCFile() + } + if ccachePath == "" { + return nil, errors.New("gssapi: credential cache path is empty") + } + ccache, err := credentials.LoadCCache(ccachePath) + if err != nil { + return nil, fmt.Errorf("gssapi: loading credentials cache: %w", err) + } + value, err := client.NewFromCCache(ccache, cfg, settings...) + if err != nil { + return nil, fmt.Errorf("gssapi: creating krb5 client from credentials cache: %w", err) + } + if err := value.AffirmLogin(); err != nil { + return nil, fmt.Errorf("gssapi: checking TGT: %w", err) + } + return value, nil + } + if useKeytab { + keytabPath := normalizedKerberosPath(options.KeytabPath) + if keytabPath == "" { + keytabPath = defaultKrbKTFile() + } + if keytabPath == "" { + return nil, errors.New("gssapi: keytab path is empty") + } + loadedKeytab, err := keytab.Load(keytabPath) + if err != nil { + return nil, fmt.Errorf("gssapi: loading keytab: %w", err) + } + if principal == "" { + principal, err = principalFromKeytab(loadedKeytab) + if err != nil { + return nil, err + } + } + username, realm, err := splitPrincipal(principal, cfg.LibDefaults.DefaultRealm) + if err != nil { + return nil, err + } + value := client.NewWithKeytab(username, realm, loadedKeytab, cfg, settings...) + if err := value.Login(); err != nil { + value.Destroy() + return nil, fmt.Errorf("gssapi: logging in with keytab: %w", err) + } + return value, nil + } + if principal != "" && password != "" { + username, realm, err := splitPrincipal(principal, cfg.LibDefaults.DefaultRealm) + if err != nil { + return nil, err + } + value := client.NewWithPassword(username, realm, password, cfg, settings...) + if err := value.Login(); err != nil { + value.Destroy() + return nil, fmt.Errorf("gssapi: logging in with password: %w", err) + } + return value, nil + } + return nil, errors.New("gssapi: no credential cache, keytab, or principal and password configured") +} + +func splitPrincipal(principal, defaultRealm string) (string, string, error) { + principal = strings.TrimSpace(principal) + separator := strings.LastIndexByte(principal, '@') + if separator < 1 || separator == len(principal)-1 { + if principal == "" || strings.TrimSpace(defaultRealm) == "" { + return "", "", fmt.Errorf("gssapi: invalid Kerberos principal %q", principal) + } + return principal, strings.TrimSpace(defaultRealm), nil + } + return principal[:separator], principal[separator+1:], nil +} + +func principalFromKeytab(value *keytab.Keytab) (string, error) { + principals := map[string]bool{} + for _, entry := range value.Entries { + principals[entry.Principal.String()] = true + } + if len(principals) != 1 { + return "", fmt.Errorf("gssapi: keytab contains %d principals; configure a client principal", len(principals)) + } + for principal := range principals { + return principal, nil + } + return "", errors.New("gssapi: keytab contains no principals") +} + +func environmentBool(key string) bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv(key))) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + +func clientOptionsFromEnvironment() ClientOptions { + return ClientOptions{ + ConfigPath: os.Getenv("KRB5_CONFIG"), + CCachePath: os.Getenv("KRB5CCNAME"), + KeytabPath: firstEnvironmentValue("KRB5_CLIENT_KTNAME", "KRB5_KTNAME"), + Principal: os.Getenv("DBX_KRB5_PRINCIPAL"), + Password: os.Getenv("DBX_KRB5_PASSWORD"), + UseCCache: environmentBool("DBX_KRB5_USE_CCACHE"), + UseKeytab: environmentBool("DBX_KRB5_USE_KEYTAB"), + DisablePAFXFAST: environmentBool("DBX_KRB5_DISABLE_PAFXFAST"), + } +} + +func firstEnvironmentValue(keys ...string) string { + for _, key := range keys { + if value := strings.TrimSpace(os.Getenv(key)); value != "" { + return value + } + } + return "" +} + +func normalizedKerberosPath(value string) string { + return strings.TrimPrefix(strings.TrimSpace(value), "FILE:") +} + +func krbConfFile() string { + cfgFile, ok := os.LookupEnv("KRB5_CONFIG") + if !ok { + cfgFile = defaultKrbConfFile() + } + + return cfgFile +} + +func krbCCFile() string { + ccFile, ok := os.LookupEnv("KRB5CCNAME") + if !ok { + ccFile = defaultKrbCCFile() + } + + return strings.TrimPrefix(ccFile, "FILE:") +} + +func krbKtFile() string { + ktFile, ok := os.LookupEnv("KRB5_CLIENT_KTNAME") + if !ok { + ktFile, ok = os.LookupEnv("KRB5_KTNAME") + } + if !ok { + ktFile = defaultKrbKTFile() + } + + return strings.TrimPrefix(ktFile, "FILE:") +} + +func (m *Krb5Mech) newWrapToken(payload []byte, sealed bool) (token wrapToken, err error) { + var flags gSSMessageTokenFlag + + if !m.isInitiator { + flags |= gSSMessageTokenFlagSentByAcceptor // send by acceptor + } + if sealed { + flags |= gSSMessageTokenFlagSealed // sealed + } + + // use the acceptor subkey if it was negotiated during auth + key := m.sessionKey + switch { + case m.acceptorSubKey != nil: + key = m.acceptorSubKey + flags |= gSSMessageTokenFlagAcceptorSubkey + case m.initiatorSubKey != nil: + key = m.initiatorSubKey + } + + token = wrapToken{ + Flags: flags, + SequenceNumber: m.ourSequenceNumber, + Payload: payload, + } + + // encrypt or sign the payload, see RFC 4121 § 4.2.4 + if sealed { + err = token.Seal(*key) + } else { + err = token.Sign(*key) + } + + if err == nil { + m.ourSequenceNumber++ // only bump the sequence number if everything is good + } + + return token, err +} + +// must return useful Kerberos error codes here so we can respond appropriately to the client if necessary +// err is also returned for local use (eg. logging an I/O error etc) +// +// This validation routine does *NOT* currently check addresses; the gokrb5 version in messages/APReq doesn't +// do this properly and in any case this behaviour should depend on the local kerberos configuration +func verifyAPReq(ktFile string, apreq *messages.APReq, skew time.Duration) (err error, krbError error) { + kt, err := keytab.Load(ktFile) + if err != nil { + krbError = messages.NewKRBError(apreq.Ticket.SName, apreq.Ticket.Realm, ianaerrcode.KRB_AP_ERR_NOKEY, "no key for service") + return + } + + err = apreq.Ticket.DecryptEncPart(kt, &apreq.Ticket.SName) + if _, ok := err.(messages.KRBError); ok { + krbError = err + return + } else if err != nil { + krbError = messages.NewKRBError(apreq.Ticket.SName, apreq.Ticket.Realm, ianaerrcode.KRB_AP_ERR_BAD_INTEGRITY, "could not decrypt ticket") + return + } + + // Check time validity of ticket + ok, err := apreq.Ticket.Valid(skew) + if err != nil || !ok { + krbError = err + return + } + + // Decrypt authenticator with session key from ticket's encrypted part + err = apreq.DecryptAuthenticator(apreq.Ticket.DecryptedEncPart.Key) + if err != nil { + krbError = messages.NewKRBError(apreq.Ticket.SName, apreq.Ticket.Realm, ianaerrcode.KRB_AP_ERR_BAD_INTEGRITY, "could not decrypt authenticator") + return + } + + // Check the authenticator checksum type + if apreq.Authenticator.Cksum.CksumType != chksumtype.GSSAPI { + krbError = messages.NewKRBError(apreq.Ticket.SName, apreq.Ticket.Realm, ianaerrcode.KRB_AP_ERR_BADMATCH, "wrong authenticator checksum type") + return + } + if len(apreq.Authenticator.Cksum.Checksum) < 24 { + krbError = messages.NewKRBError(apreq.Ticket.SName, apreq.Ticket.Realm, ianaerrcode.KRB_AP_ERR_BADMATCH, "authenticator checksum too short") + return + } + + // Check CName in authenticator is the same as that in the ticket + if !apreq.Authenticator.CName.Equal(apreq.Ticket.DecryptedEncPart.CName) { + krbError = messages.NewKRBError(apreq.Ticket.SName, apreq.Ticket.Realm, ianaerrcode.KRB_AP_ERR_BADMATCH, "CName in Authenticator does not match that in service ticket") + return + } + + // Check the clock skew between the client and the service server + ct := apreq.Authenticator.CTime.Add(time.Duration(apreq.Authenticator.Cusec) * time.Microsecond) + t := time.Now().UTC() + if t.Sub(ct) > skew || ct.Sub(t) > skew { + krbError = messages.NewKRBError(apreq.Ticket.SName, apreq.Ticket.Realm, ianaerrcode.KRB_AP_ERR_SKEW, fmt.Sprintf("clock skew with client too large. greater than %v seconds", skew)) + return + } + + return nil, nil +} + +func mkGssErrKrbCode(code int32, message string) (token []byte, err error) { + ke := messages.NewKRBError(types.PrincipalName{}, "", code, message) + return mkGssErrFromKrbErr(ke) +} + +func mkGssErrFromKrbErr(ke messages.KRBError) (token []byte, err error) { + tb, _ := hex.DecodeString(tokenIDKrbError) + gssToken := kRB5Token{ + oID: oID(), + tokID: tb, + kRBError: &ke, + } + + token, err = gssToken.marshal() + if err == nil { + // marshaled ok, return the kerberos error and token to the peer + err = ke + } else { + // marshal failed, return that error and no token to send to the peer + token = nil + } + + return +} + +// Generate a base key -- usually the same as GenerateEncryptionKey, except +// that the gokrb5 library doesn't handle the hash/integrity and the encryption +// keys being different lengths in aes256-cts-hmac-sha384-192 +// TODO: fix GenerateEncryptionKey at some point to cope with different +// uses like this case. +func GenerateBaseKey(etype etype.EType) (types.EncryptionKey, error) { + k := types.EncryptionKey{ + KeyType: etype.GetETypeID(), + } + + // special-case one encryption type + kl := etype.GetKeyByteSize() + if etype.GetETypeID() == etypeID.AES256_CTS_HMAC_SHA384_192 { + kl = 32 + } + + b := make([]byte, kl) + _, err := rand.Read(b) + if err != nil { + return k, err + } + k.KeyValue = b + return k, nil +} diff --git a/agents/go-common/go-gssapi/krb5/krb5_test.go b/agents/go-common/go-gssapi/krb5/krb5_test.go new file mode 100644 index 0000000000..211e49ecbe --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/krb5_test.go @@ -0,0 +1,8 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +package krb5 + +// TODO: testing this stuff really needs a KDC or a mock +// diff --git a/agents/go-common/go-gssapi/krb5/message_token.go b/agents/go-common/go-gssapi/krb5/message_token.go new file mode 100644 index 0000000000..9aa226547a --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/message_token.go @@ -0,0 +1,506 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +package krb5 + +import ( + "bytes" + "crypto/hmac" + "encoding/binary" + "errors" + "fmt" + + "github.com/jcmturner/gokrb5/v8/crypto" + "github.com/jcmturner/gokrb5/v8/iana/keyusage" + + "github.com/jcmturner/gokrb5/v8/types" +) + +/* + * Derived from github.com/jcmturner/gokrb5/gssapi/wrapToken.go + * + * The modified version adds functionality for sealing GSSAPI messages + * + */ + +// RFC 4121 § 4.2.6 +const ( + msgTokenHdrLen = 16 + msgTokenFillerByte byte = 0xFF +) + +// RFC 4121 § 4.2.2 +type gSSMessageTokenFlag uint8 + +const ( + gSSMessageTokenFlagSentByAcceptor gSSMessageTokenFlag = 1 << iota + gSSMessageTokenFlagSealed + gSSMessageTokenFlagAcceptorSubkey +) + +// RFC 4121 § 4.2.6.1 +type mICToken struct { + // 2 byte token ID (0x04, 0x04) + Flags gSSMessageTokenFlag + // 5 byte filler (0xFF) + SequenceNumber uint64 // 64-bit sequence number + Checksum []byte + signed bool +} + +// RFC 4121 § 4.2.6.2 +type wrapToken struct { + // 2 byte token ID (0x05, 0x04) + Flags gSSMessageTokenFlag + // 1 byte filler (0xFF) + EC uint16 // "Extra count" - the checksum or padding length + RRC uint16 // right rotation count for SSPI (we don't support this yet) + SequenceNumber uint64 // 64-bit sequence number + Payload []byte // signed or encrypted payload + signedOrSealed bool +} + +// Return the 2 bytes identifying a GSS API Wrap token +func getGssWrapTokenID() [2]byte { + return [2]byte{0x05, 0x04} +} + +// Return the 2 bytes identifying a GSS API MIC token +func getGssMICTokenID() [2]byte { + return [2]byte{0x04, 0x04} +} + +// RFC 4121 § 4.2.4 +// Checksum is calculated over the plaintext (supplied token payload), and +// the token header with EC and RRC set to zero +// The function modifies the Payload and EC/RRC fields of the WrapToken +func (wt *wrapToken) Sign(key types.EncryptionKey) error { + if wt.Payload == nil { + return errors.New("gssapi: attempt to sign token with no payload") + } + if wt.signedOrSealed { + return errors.New("gssapi: attempt to sign a signed/sealed token") + } + + sig, err := wt.computeChecksum(key) + if err != nil { + return fmt.Errorf("gssapi: %s", err) + } + + encType, err := crypto.GetEtype(key.KeyType) + if err != nil { + return fmt.Errorf("gssapi: %s", err) + } + + wt.Payload = append(wt.Payload, sig...) + wt.EC = uint16(encType.GetHMACBitLength() / 8) + wt.RRC = 0 + wt.signedOrSealed = true + + return nil +} + +// RFC 4121 § 4.2.4 +// Encrypts the Payload and sets EC/RRC on the WrapToken +func (wt *wrapToken) Seal(key types.EncryptionKey) (err error) { + if wt.Payload == nil { + return errors.New("gssapi: attempt to encrypt token with no payload") + } + if wt.signedOrSealed { + return errors.New("gssapi: attempt to seal a signed/sealed token") + } + + toEncrypt := make([]byte, 0, len(wt.Payload)+msgTokenHdrLen) + toEncrypt = append(toEncrypt, wt.Payload...) + toEncrypt = append(toEncrypt, wt.header()...) + + usage := keyusage.GSSAPI_INITIATOR_SEAL + if wt.Flags&gSSMessageTokenFlagSentByAcceptor != 0 { + usage = keyusage.GSSAPI_ACCEPTOR_SEAL + } + + encType, err := crypto.GetEtype(key.KeyType) + if err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + var encData []byte + _, encData, err = encType.EncryptMessage(key.KeyValue, toEncrypt, uint32(usage)) + if err != nil { + err = fmt.Errorf("gssapi: %s", err) + } + + wt.Payload = encData + wt.EC = 0 + wt.RRC = 0 + wt.signedOrSealed = true + + return err +} + +func (wt *wrapToken) header() (hdr []byte) { + hdr = make([]byte, msgTokenHdrLen) + + tokID := getGssWrapTokenID() + hdr1 := []byte{ + tokID[0], tokID[1], // token ID + byte(wt.Flags), // flags + 0xFF, // filler + 0x00, 0x00, // EC + 0x00, 0x00, // RRC + } + + copy(hdr, hdr1) + binary.BigEndian.PutUint64(hdr[8:], wt.SequenceNumber) + + return +} + +func (wt *wrapToken) computeChecksum(key types.EncryptionKey) (cksum []byte, err error) { + // wrap tokens always use the Seal key usage (RFC 4121 § 2) + usage := keyusage.GSSAPI_INITIATOR_SEAL + if wt.Flags&gSSMessageTokenFlagSentByAcceptor != 0 { + usage = keyusage.GSSAPI_ACCEPTOR_SEAL + } + + plLen := len(wt.Payload) + + // Build a slice containing { payload | header } + cksumData := make([]byte, 0, msgTokenHdrLen+plLen) + cksumData = append(cksumData, wt.Payload...) + cksumData = append(cksumData, wt.header()...) + + encType, err := crypto.GetEtype(key.KeyType) + if err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + + cksum, err = encType.GetChecksumHash(key.KeyValue, cksumData, uint32(usage)) + if err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + + return +} + +// Marshal a token that has already been signed or sealed +func (wt *wrapToken) Marshal() (token []byte, err error) { + if !wt.signedOrSealed { + err = errors.New("gssapi: wrap token is not signed or sealed") + return + } + + tokenID := getGssWrapTokenID() + token = make([]byte, msgTokenHdrLen+len(wt.Payload)) + + copy(token[0:], tokenID[:]) + token[2] = byte(wt.Flags) + token[3] = msgTokenFillerByte + binary.BigEndian.PutUint16(token[4:6], wt.EC) + binary.BigEndian.PutUint16(token[6:8], wt.RRC) + binary.BigEndian.PutUint64(token[8:16], wt.SequenceNumber) + copy(token[16:], wt.Payload) + + return +} + +// Unmarshal a signed or sealed token +func (wt *wrapToken) Unmarshal(token []byte) (err error) { + // zero everything in the token + *wt = wrapToken{} + + // token must be at least 16 bytes + if len(token) < msgTokenHdrLen { + return errors.New("gssapi: wrap token is too short") + } + + // Check for 0x60 as the first byte; As per RFC 4121 § 4.4, these Token IDs + // are reserved - and indicate 'Generic GSS-API token framing' that was used by + // GSS-API v1, and are not supported in GSS-API v2.. catch that specific case so + // we can emmit a useful message + if token[0] == 0x60 { + return errors.New("gssapi: GSS-API v1 message tokens are not supported") + } + + // check token ID + tokenID := getGssWrapTokenID() + if !bytes.Equal(tokenID[:], token[0:2]) { + return errors.New("gssapi: bad wrap token ID") + } + + wt.Flags = gSSMessageTokenFlag(token[2]) + + if token[3] != msgTokenFillerByte { + return errors.New("gssapi: invalid wrap token (bad filler)") + } + + wt.EC = binary.BigEndian.Uint16(token[4:6]) + wt.RRC = binary.BigEndian.Uint16(token[6:8]) + wt.SequenceNumber = binary.BigEndian.Uint64(token[8:16]) + + if len(token) > msgTokenHdrLen { + wt.Payload = token[16:] + } + + wt.signedOrSealed = true + return nil +} + +func (wt *wrapToken) VerifyAndDecode(key types.EncryptionKey, expectFromAcceptor bool) (isSealed bool, err error) { + if !wt.signedOrSealed { + return false, errors.New("gssapi: wrap token is not signed or sealed") + } + if wt.Payload == nil || len(wt.Payload) == 0 { + return false, errors.New("gssapi: cannot verify an empty wrap token payload") + } + + isFromAcceptor := wt.Flags&gSSMessageTokenFlagSentByAcceptor != 0 + if isFromAcceptor != expectFromAcceptor { + return false, fmt.Errorf("gssapi: wrap token from acceptor: %t, expect from acceptor: %t", isFromAcceptor, expectFromAcceptor) + } + + if wt.Flags&gSSMessageTokenFlagSealed != 0 { + return true, wt.decrypt(key) + } else { + return false, wt.checkSig(key) + } +} + +func (wt *wrapToken) decrypt(key types.EncryptionKey) (err error) { + usage := keyusage.GSSAPI_INITIATOR_SEAL + if wt.Flags&gSSMessageTokenFlagSentByAcceptor != 0 { + usage = keyusage.GSSAPI_ACCEPTOR_SEAL + } + + encType, err := crypto.GetEtype(key.KeyType) + if err != nil { + return fmt.Errorf("gssapi: wrap token: %s", err) + } + + var decrypted []byte + decrypted, err = encType.DecryptMessage(key.KeyValue, wt.Payload, uint32(usage)) + if err != nil { + return fmt.Errorf("gssapi: wrap token: %s", err) + } + + // check that the decrypted payload is big enough + if len(decrypted) < int(wt.EC+msgTokenHdrLen) { + return errors.New("gssapi: decrypted wrap token payload is too short") + } + + // save the decrypted header part from the end of the plaintext + decryptedHeader := decrypted[len(decrypted)-msgTokenHdrLen:] + + // check that plain text header wasn't modified + wt2 := wrapToken{} + if err = wt2.Unmarshal(decryptedHeader); err != nil { + return + } + if !(wt.Flags == wt2.Flags && + wt.EC == wt2.EC && + wt.SequenceNumber == wt2.SequenceNumber) { + return errors.New("gssapi: wrap token header was modified") + } + + // remove the header and extra-count bytes from the decrypted payload + wt.Payload = decrypted[0 : len(decrypted)-msgTokenHdrLen-int(wt.EC)] + wt.signedOrSealed = false + + return err +} + +func (wt *wrapToken) checkSig(key types.EncryptionKey) (err error) { + encType, err := crypto.GetEtype(key.KeyType) + if err != nil { + return fmt.Errorf("gssapi: wrap token: %s", err) + } + + // extra-count should be the crypto checksum length + if wt.EC != uint16(encType.GetHMACBitLength()/8) { + return errors.New("gssapi: bad wrap token checksum length") + } + + // check that the payload is big enough + if len(wt.Payload) < int(wt.EC) { + return errors.New("gssapi: signed wrap token payload is too short") + } + + tokCksum := wt.Payload[len(wt.Payload)-int(wt.EC):] + + wt2 := *wt + wt2.Payload = wt.Payload[0 : len(wt.Payload)-int(wt.EC)] + computedCksum, err := wt2.computeChecksum(key) + if err != nil { + return fmt.Errorf("gssapi: %s", err) + } + + if !hmac.Equal(tokCksum, computedCksum) { + return errors.New("gssapi: invalid wrap token checksum") + } + + // remove the signature from the payload + wt.Payload = wt.Payload[0 : len(wt.Payload)-int(wt.EC)] + wt.signedOrSealed = false + + return err +} + +// Ported from MIT source code (gss_krb5int_rotate_left) +func rotateLeft(buf []byte, rc uint) (out []byte) { + defer func() { + out = buf + }() + + if len(buf) == 0 || rc == 0 { + return + } + + rc %= uint(len(buf)) + if rc == 0 { + return + } + + tmpBuf := make([]byte, rc) + copy(tmpBuf, buf[0:rc]) + copy(buf, buf[rc:]) + copy(buf[uint(len(buf))-rc:], tmpBuf) + + return +} + +// RFC 4121 § 4.2.4 +// Checksum is calculated over the plaintext (supplied token payload), and +// the token header +func (mt *mICToken) Sign(payload []byte, key types.EncryptionKey) (err error) { + // mic tokens always use the Sign key usage + usage := keyusage.GSSAPI_INITIATOR_SIGN + if mt.Flags&gSSMessageTokenFlagSentByAcceptor != 0 { + usage = keyusage.GSSAPI_ACCEPTOR_SIGN + } + + cksumData := make([]byte, 0, msgTokenHdrLen+len(payload)) + cksumData = append(cksumData, payload...) + cksumData = append(cksumData, mt.header()...) + + encType, err := crypto.GetEtype(key.KeyType) + if err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + + mt.Checksum, err = encType.GetChecksumHash(key.KeyValue, cksumData, uint32(usage)) + if err != nil { + err = fmt.Errorf("gssapi: %s", err) + return + } + + mt.signed = true + + return +} + +func (mt *mICToken) header() (hdr []byte) { + hdr = make([]byte, msgTokenHdrLen) + + tokID := getGssMICTokenID() + hdr1 := []byte{ + tokID[0], tokID[1], // token ID + byte(mt.Flags), // flags + 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // filler + 0x00, 0x00, // EC + 0x00, 0x00, // RRC + } + + copy(hdr, hdr1) + binary.BigEndian.PutUint64(hdr[8:], mt.SequenceNumber) + + return +} + +func (mt *mICToken) Marshal() (token []byte, err error) { + if !mt.signed { + err = errors.New("gssapi: MIC token is not signed") + } + + tokenID := getGssMICTokenID() + token = make([]byte, msgTokenHdrLen+len(mt.Checksum)) + + copy(token[0:], tokenID[:]) + token[2] = byte(mt.Flags) + copy(token[3:8], []byte{0xFF, 0xFF, 0xFF, 0xFF, 0xFF}) + binary.BigEndian.PutUint64(token[8:16], mt.SequenceNumber) + copy(token[16:], mt.Checksum) + + return +} + +func (mt *mICToken) Unmarshal(token []byte) (err error) { + // zero out the MIC token + *mt = mICToken{} + + // token must be at least 16 bytes + if len(token) < msgTokenHdrLen { + return errors.New("gssapi: wrap token is too short") + } + + // Check for 0x60 as the first byte; As per RFC 4121 § 4.4, these Token IDs + // are reserved - and indicate 'Generic GSS-API token framing' that was used by + // GSS-API v1, and are not supported in GSS-API v2.. catch that specific case so + // we can emmit a useful message + if token[0] == 0x60 { + return errors.New("gssapi: GSS-API v1 message tokens are not supported") + } + + // check token ID + tokenID := getGssMICTokenID() + if !bytes.Equal(tokenID[:], token[0:2]) { + return errors.New("gssapi: bad MIC token ID") + } + + mt.Flags = gSSMessageTokenFlag(token[2]) + + if !bytes.Equal(token[3:8], []byte{0xFF, 0xFF, 0xFF, 0xFF, 0xFF}) { + return errors.New("gssapi: invalid MIC token (bad filler)") + } + + mt.SequenceNumber = binary.BigEndian.Uint64(token[8:16]) + + if len(token) > msgTokenHdrLen { + mt.Checksum = token[16:] + } + + mt.signed = true + + return err +} + +func (mt *mICToken) Verify(payload []byte, key types.EncryptionKey, expectFromAcceptor bool) (err error) { + if !mt.signed { + return errors.New("gssapi: MIC token is not signed") + } + + if len(payload) == 0 { + return errors.New("gssapi: cannot verify an empty MIC token payload") + } + + isFromAcceptor := mt.Flags&gSSMessageTokenFlagSentByAcceptor != 0 + if isFromAcceptor != expectFromAcceptor { + return fmt.Errorf("gssapi: MIC token from acceptor: %t, expect from acceptor: %t", isFromAcceptor, expectFromAcceptor) + } + + // copy the token and use it to sign the supplied payload + wt2 := *mt + if err = wt2.Sign(payload, key); err != nil { + return err + } + + // check the token's checksums + if !bytes.Equal(mt.Checksum, wt2.Checksum) { + return errors.New("gssapi: invalid MIC token checksum") + } + + return +} diff --git a/agents/go-common/go-gssapi/krb5/message_token_test.go b/agents/go-common/go-gssapi/krb5/message_token_test.go new file mode 100644 index 0000000000..f1b4323ac1 --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/message_token_test.go @@ -0,0 +1,208 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +package krb5 + +import ( + "encoding/hex" + "fmt" + "testing" + + "github.com/jcmturner/gokrb5/v8/iana/etypeID" + "github.com/stretchr/testify/assert" + + "github.com/jcmturner/gokrb5/v8/types" +) + +const ( + TestWrapPayload = "testing 123" + + // from kadmin: + // ank -kvno 123 -pw password -e test test + // ktadd -k test.kt -norandkey test + TestAES256KVNO = 123 + TestAES256Key = "93860ea9a3961f58f1e1370286c720ab8da6574cacb26396f7de6ebfbbfd00a0" + AESCksumLen = 12 + EncPayloadLen = 55 + + SampleWrapTokenSignature = "71914A5D08018A97375AB52A" + WrapTokenSignedHeader = "050400ff000c0000000000000000007B" + SampleSignedWrapTOken = "050404ff000c000000000000209bb2cb74657374696e6720313233efed11aa6caa6cf5a7e595a5" + SampleSignedWrapTokenWindows = "050400ff000c000c0000000000000000a79b6be6ce749f2f6102c78774657374" + SampleMICTokenSignature = "b479cc6b1a27beb60a815b26" + MICTokenHeader = "040404ffffffffff000000000000007B" + SampleMICToken = "040404ffffffffff000000000000007Bb479cc6b1a27beb60a815b26" +) + +func mkSampleWrapToken() (wt wrapToken) { + return wrapToken{ + Flags: 0, + SequenceNumber: 123, + Payload: []byte(TestWrapPayload), + } +} + +func mkSampleMICToken() (mt mICToken) { + return mICToken{ + Flags: 4, + SequenceNumber: 123, + } +} + +func mkSampleAESKey() (key types.EncryptionKey) { + b, _ := hex.DecodeString(TestAES256Key) + return types.EncryptionKey{ + KeyType: etypeID.AES256_CTS_HMAC_SHA1_96, + KeyValue: b, + } +} + +func TestWrapTokenSign(t *testing.T) { + key := mkSampleAESKey() + tok := mkSampleWrapToken() + + err := tok.Sign(key) + + assert.NoError(t, err, "signing operation failed") + assert.True(t, tok.signedOrSealed, "token was not signed") + assert.Equal(t, uint16(AESCksumLen), tok.EC, "wrong checksum length") + assert.Equal(t, len(TestWrapPayload)+AESCksumLen, len(tok.Payload), "wrong signed payload length") + + wantSig, _ := hex.DecodeString(SampleWrapTokenSignature) + assert.Equal(t, wantSig, tok.Payload[len(TestWrapPayload):], "signature not as expected") + assert.Equal(t, []byte(TestWrapPayload), tok.Payload[0:len(TestWrapPayload)], "corrupt payload") +} + +func TestWrapTokenSeal(t *testing.T) { + key := mkSampleAESKey() + tok := mkSampleWrapToken() + + err := tok.Seal(key) + + assert.NoError(t, err, "sealing operation failed") + assert.True(t, tok.signedOrSealed, "token was not sealed") + assert.Equal(t, uint16(0), tok.EC, "wrong extra-count") + assert.Equal(t, EncPayloadLen, len(tok.Payload), "sealed token length is wrong") +} + +func TestWrapTokenMarshal(t *testing.T) { + key := mkSampleAESKey() + tok := mkSampleWrapToken() + + _, err := tok.Marshal() + assert.Error(t, err, "Marshal of unsigned/sealed token should be an error") + + err = tok.Sign(key) + assert.NoError(t, err, "signing operation failed") + + tokBytes, err := tok.Marshal() + assert.NoError(t, err, "Marshal of signed token should succeed") + assert.Equal(t, 16+len(TestWrapPayload)+AESCksumLen, len(tokBytes), "bad token length") + + wantHeader, _ := hex.DecodeString(WrapTokenSignedHeader) + assert.Equal(t, wantHeader, tokBytes[0:16], "bad wrap token header") + + wantSig, _ := hex.DecodeString(SampleWrapTokenSignature) + assert.Equal(t, []byte(TestWrapPayload), tokBytes[16:16+len(TestWrapPayload)], "corrupt payload") + assert.Equal(t, wantSig, tokBytes[16+len(TestWrapPayload):], "signature not as expected") +} + +func TestWrapTokenUnmarshal(t *testing.T) { + tokBytes, _ := hex.DecodeString(SampleSignedWrapTOken) + + tok := wrapToken{} + err := tok.Unmarshal(tokBytes) + assert.NoError(t, err, "Unmarshal of signed token failed") + + assert.Equal(t, 0x04, int(tok.Flags), "bad token flags") + assert.Equal(t, uint16(AESCksumLen), tok.EC, "bad EC (signature length)") + assert.Equal(t, uint16(0), tok.RRC, "bad RRC") + assert.Equal(t, uint64(0x209bb2cb), tok.SequenceNumber, "bad sequence number") + assert.Equal(t, true, tok.signedOrSealed, "token is not signed/sealed") +} + +func TestWindowsWrapTokenUnmarshal(t *testing.T) { + tokBytes, _ := hex.DecodeString(SampleSignedWrapTokenWindows) + + tok := wrapToken{} + err := tok.Unmarshal(tokBytes) + assert.NoError(t, err, "Unmarshal of signed token failed") + + assert.Equal(t, 0x00, int(tok.Flags), "bad token flags") + assert.Equal(t, uint16(AESCksumLen), tok.EC, "bad EC (signature length)") + assert.Equal(t, uint16(12), tok.RRC, "bad RRC") + assert.Equal(t, uint64(0), tok.SequenceNumber, "bad sequence number") + assert.Equal(t, true, tok.signedOrSealed, "token is not signed/sealed") +} + +func TestRotateLeft(t *testing.T) { + var testData = "abcdefghijklmnop" + + var tests = []struct { + rc uint + expected string + }{ + {0, "abcdefghijklmnop"}, + {1, "bcdefghijklmnopa"}, + {15, "pabcdefghijklmno"}, + {16, "abcdefghijklmnop"}, + {17, "bcdefghijklmnopa"}, + {32, "abcdefghijklmnop"}, + {33, "bcdefghijklmnopa"}, + } + + for _, tt := range tests { + t.Run(fmt.Sprintf("rc=%d", tt.rc), func(t *testing.T) { + in := testData + out := rotateLeft([]byte(in), tt.rc) + assert.Equal(t, tt.expected, string(out)) + }) + } +} + +func TestMICTokenSign(t *testing.T) { + key := mkSampleAESKey() + tok := mkSampleMICToken() + + err := tok.Sign([]byte(TestWrapPayload), key) + + assert.NoError(t, err, "signing operation failed") + assert.True(t, tok.signed, "token was not signed") + + wantSig, _ := hex.DecodeString(SampleMICTokenSignature) + assert.Equal(t, wantSig, tok.Checksum, "signature not as expected") +} + +func TestMICTokenMarshal(t *testing.T) { + key := mkSampleAESKey() + tok := mkSampleMICToken() + + _, err := tok.Marshal() + assert.Error(t, err, "Marshal of unsigned MIC token should be an error") + + err = tok.Sign([]byte(TestWrapPayload), key) + assert.NoError(t, err, "signing operation failed") + + tokBytes, err := tok.Marshal() + assert.NoError(t, err, "Marshal of signed token should succeed") + assert.Equal(t, 16+AESCksumLen, len(tokBytes), "bad token length") + + wantHeader, _ := hex.DecodeString(MICTokenHeader) + assert.Equal(t, wantHeader, tokBytes[0:16], "bad MIC token header") + + wantSig, _ := hex.DecodeString(SampleMICTokenSignature) + assert.Equal(t, wantSig, tokBytes[16:], "signature not as expected") +} + +func TestMICTokenUnmarshal(t *testing.T) { + tokBytes, _ := hex.DecodeString(SampleMICToken) + + tok := mICToken{} + err := tok.Unmarshal(tokBytes) + assert.NoError(t, err, "Unmarshal of MIC token failed") + + assert.Equal(t, 0x04, int(tok.Flags), "bad token flags") + assert.Equal(t, uint64(123), tok.SequenceNumber, "bad sequence number") + assert.Equal(t, true, tok.signed, "token is not signed/sealed") +} diff --git a/agents/go-common/go-gssapi/krb5/sample_test.go b/agents/go-common/go-gssapi/krb5/sample_test.go new file mode 100644 index 0000000000..bffb4810e9 --- /dev/null +++ b/agents/go-common/go-gssapi/krb5/sample_test.go @@ -0,0 +1,110 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +package krb5 + +import ( + "encoding/binary" + "time" + + "github.com/jcmturner/gokrb5/v8/iana/msgtype" + "github.com/jcmturner/gokrb5/v8/test/testdata" + + "github.com/jcmturner/gokrb5/v8/messages" + "github.com/jcmturner/gokrb5/v8/types" +) + +// Sample data from MIT Kerberos v1.19.1 + +// from src/tests/asn.1/ktest.h +const ( + SampleUSec = 123456 + SampleSeqNumber = 17 + SampleNonce = 42 + SampleFlags = 0xFEDCBA98 + SampleError = 0x3C + SamplePrincipalName = "hftsai/extra@ATHENA.MIT.EDU" + SampleData = "krb5data" +) + +func ktestMakeSampleAPReqEncPart() encAPRepPart { + tm, _ := time.Parse(testdata.TEST_TIME_FORMAT, testdata.TEST_TIME) + return encAPRepPart{ + CTime: tm, + Cusec: SampleUSec, + Subkey: ktestMakeSampleKeyblock(), + SequenceNumber: SampleSeqNumber, + } +} + +func ktestMakeSampleKeyblock() types.EncryptionKey { + kv := []byte("12345678") + return types.EncryptionKey{ + KeyType: 1, + KeyValue: kv, + } +} + +func ktestMakeSampleEncData() types.EncryptedData { + return types.EncryptedData{ + EType: 0, + KVNO: 5, + Cipher: []byte(testdata.TEST_CIPHERTEXT), + } +} + +func ktestMakeSampleTicket() messages.Ticket { + pn, realm := types.ParseSPNString(SamplePrincipalName) + return messages.Ticket{ + TktVNO: 5, + Realm: realm, + SName: pn, + EncPart: ktestMakeSampleEncData(), + } +} + +func ktestMakeSampleApReq() (aprep messages.APReq) { + aprep = messages.APReq{ + PVNO: 5, + MsgType: msgtype.KRB_AP_REQ, + APOptions: types.NewKrbFlags(), + Ticket: ktestMakeSampleTicket(), + EncryptedAuthenticator: ktestMakeSampleEncData(), + } + + binary.BigEndian.PutUint32(aprep.APOptions.Bytes[0:], SampleFlags) + return +} + +func ktestMakeSampleApRep() (aprep aPRep) { + aprep = aPRep{ + PVNO: 5, + MsgType: msgtype.KRB_AP_REP, + EncPart: ktestMakeSampleEncData(), + } + + return +} + +func ktestMakeSampleError() (krberr messages.KRBError) { + pn, realm := types.ParseSPNString(SamplePrincipalName) + tm, _ := time.Parse(testdata.TEST_TIME_FORMAT, testdata.TEST_TIME) + krberr = messages.KRBError{ + PVNO: 5, + MsgType: msgtype.KRB_ERROR, + CTime: tm, + Cusec: SampleUSec, + STime: tm, + Susec: SampleUSec, + ErrorCode: SampleError, + CRealm: realm, + CName: pn, + Realm: realm, + SName: pn, + EText: SampleData, + EData: []byte(SampleData), + } + + return +} diff --git a/agents/go-common/go-gssapi/registry.go b/agents/go-common/go-gssapi/registry.go new file mode 100644 index 0000000000..19d02c11da --- /dev/null +++ b/agents/go-common/go-gssapi/registry.go @@ -0,0 +1,61 @@ +// Copyright 2021 Jake Scott. All rights reserved. +// Use of this source code is governed by the Apache License +// version 2.0 that can be found in the LICENSE file. + +package gssapi + +import "strings" + +type MechFactory func() Mech + +var mechs map[string]MechFactory + +func init() { + mechs = make(map[string]MechFactory) +} + +// Register should be called by Mech implementations to enable +// a mechanism to be used by clients +func Register(name string, f MechFactory) { + name = strings.ToLower(name) + _, ok := mechs[name] + + // can't register two mechs with the same name + if ok { + panic("Cannot have two mechs named " + name) + } + + mechs[name] = f +} + +// IsRegistered can be used to find out whether a named +// mechanism is registered or not +func IsRegistered(name string) bool { + name = strings.ToLower(name) + _, ok := mechs[name] + + return ok +} + +// NewMech returns a mechanism context by name +func NewMech(name string) Mech { + name = strings.ToLower(name) + f, ok := mechs[name] + + if ok { + return f() + } + + return nil +} + +// Mechs returns the list of registered mechanism names +func Mechs() (l []string) { + l = make([]string, 0, len(mechs)) + + for name := range mechs { + l = append(l, name) + } + + return +} diff --git a/agents/go-common/gohive/LICENSE b/agents/go-common/gohive/LICENSE new file mode 100644 index 0000000000..f4c25a66ec --- /dev/null +++ b/agents/go-common/gohive/LICENSE @@ -0,0 +1,20 @@ + +The MIT License (MIT) + +Copyright (c) 2018-2024 Jaume Marhuenda Beltran + +Permission is hereby granted, free of charge, to any person obtaining a copy of this +software and associated documentation files (the "Software"), to deal in the Software +without restriction, including without limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons +to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or +substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, +INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR +PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE +FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. diff --git a/agents/go-common/gohive/browser_auth.go b/agents/go-common/gohive/browser_auth.go new file mode 100644 index 0000000000..91633bab3a --- /dev/null +++ b/agents/go-common/gohive/browser_auth.go @@ -0,0 +1,242 @@ +package gohive + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + "net/url" + "os/exec" + "runtime" + "strconv" + "strings" + "sync" + "time" +) + +const defaultBrowserResponseTimeout = 120 * time.Second + +type browserSSOResponse struct { + Successful bool + Message string + Token string +} + +type browserSSOClient struct { + configuredPort int + timeout time.Duration + openURL func(string) error + responses chan browserSSOResponse + + mu sync.RWMutex + server *http.Server + listener net.Listener + port int + redirectURL *url.URL + clientIdentifier string + token string +} + +func newBrowserSSOClient(port int, timeout time.Duration) *browserSSOClient { + if timeout <= 0 { + timeout = defaultBrowserResponseTimeout + } + return &browserSSOClient{ + configuredPort: port, + timeout: timeout, + openURL: openBrowserURL, + responses: make(chan browserSSOResponse, 1), + } +} + +func (client *browserSSOClient) Start() error { + listener, err := net.Listen("tcp4", net.JoinHostPort("127.0.0.1", strconv.Itoa(client.configuredPort))) + if err != nil { + return fmt.Errorf("start Hive browser SSO callback listener: %w", err) + } + server := &http.Server{ + Handler: client, + ReadHeaderTimeout: 5 * time.Second, + IdleTimeout: client.timeout, + } + client.mu.Lock() + client.listener = listener + client.server = server + client.port = listener.Addr().(*net.TCPAddr).Port + client.mu.Unlock() + go func() { + _ = server.Serve(listener) + }() + return nil +} + +func (client *browserSSOClient) ServeHTTP(responseWriter http.ResponseWriter, request *http.Request) { + if request.Method != http.MethodPost { + http.Error(responseWriter, "method not allowed", http.StatusMethodNotAllowed) + return + } + request.Body = http.MaxBytesReader(responseWriter, request.Body, 1<<20) + if err := request.ParseForm(); err != nil { + http.Error(responseWriter, "invalid SSO response", http.StatusBadRequest) + return + } + response := browserSSOResponse{ + Successful: strings.EqualFold(request.FormValue("status"), "true"), + Message: request.FormValue("message"), + Token: request.FormValue("token"), + } + responseWriter.Header().Set("Content-Type", "text/html; charset=utf-8") + responseWriter.WriteHeader(http.StatusOK) + if response.Successful { + _, _ = responseWriter.Write([]byte("Successfully authenticated. You may close this window.")) + } else { + _, _ = responseWriter.Write([]byte("Authentication failed. You may close this window.")) + } + select { + case client.responses <- response: + default: + } +} + +func (client *browserSSOClient) Port() int { + client.mu.RLock() + defer client.mu.RUnlock() + return client.port +} + +func (client *browserSSOClient) SetRedirect(rawURL, clientIdentifier string) error { + parsed, err := url.Parse(rawURL) + if err != nil || !parsed.IsAbs() || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return fmt.Errorf("invalid Hive browser SSO redirect URL %q", rawURL) + } + if strings.TrimSpace(clientIdentifier) == "" { + return errors.New("Hive browser SSO response omitted X-Hive-Client-Identifier") + } + client.mu.Lock() + client.redirectURL = parsed + client.clientIdentifier = clientIdentifier + client.mu.Unlock() + return nil +} + +func (client *browserSSOClient) HasRedirect() bool { + client.mu.RLock() + defer client.mu.RUnlock() + return client.redirectURL != nil +} + +func (client *browserSSOClient) Authenticate(ctx context.Context) error { + client.mu.RLock() + redirectURL := client.redirectURL + openURL := client.openURL + timeout := client.timeout + client.mu.RUnlock() + if redirectURL == nil { + return errors.New("Hive browser SSO did not receive a redirect URL") + } + if err := openURL(redirectURL.String()); err != nil { + return fmt.Errorf("open Hive browser SSO URL: %w", err) + } + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return errors.New("timed out while waiting for Hive browser SSO response") + case response := <-client.responses: + if !response.Successful { + if strings.TrimSpace(response.Message) == "" { + return errors.New("Hive browser SSO authentication failed") + } + return fmt.Errorf("Hive browser SSO authentication failed: %s", response.Message) + } + if strings.TrimSpace(response.Token) == "" { + return errors.New("Hive browser SSO returned an empty token") + } + client.mu.Lock() + client.token = response.Token + client.mu.Unlock() + return nil + } +} + +func (client *browserSSOClient) Credentials() (string, string) { + client.mu.RLock() + defer client.mu.RUnlock() + return client.token, client.clientIdentifier +} + +func (client *browserSSOClient) Close() error { + client.mu.Lock() + server := client.server + listener := client.listener + client.server = nil + client.listener = nil + client.mu.Unlock() + if server != nil { + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + return server.Shutdown(ctx) + } + if listener != nil { + return listener.Close() + } + return nil +} + +type browserAuthRoundTripper struct { + Base http.RoundTripper + Client *browserSSOClient +} + +func (transport *browserAuthRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + base := transport.Base + if base == nil { + base = http.DefaultTransport + } + attempt := request.Clone(request.Context()) + attempt.Header = request.Header.Clone() + token, clientIdentifier := transport.Client.Credentials() + if token == "" { + port := transport.Client.Port() + if port <= 0 { + return nil, errors.New("Hive browser SSO callback listener is not running") + } + attempt.Header.Set("X-Hive-Token-Response-Port", strconv.Itoa(port)) + } else { + attempt.Header.Del("X-Hive-Token-Response-Port") + attempt.Header.Set("Authorization", "Bearer "+token) + attempt.Header.Set("X-Hive-Client-Identifier", clientIdentifier) + } + response, err := base.RoundTrip(attempt) + if err != nil || response == nil || (response.StatusCode != http.StatusFound && response.StatusCode != http.StatusSeeOther) { + return response, err + } + if err := transport.Client.SetRedirect(response.Header.Get("Location"), response.Header.Get("X-Hive-Client-Identifier")); err != nil { + if response.Body != nil { + _ = response.Body.Close() + } + return nil, err + } + return response, nil +} + +func openBrowserURL(rawURL string) error { + var command *exec.Cmd + switch runtime.GOOS { + case "darwin": + command = exec.Command("open", rawURL) + case "windows": + command = exec.Command("rundll32", "url.dll,FileProtocolHandler", rawURL) + case "linux": + command = exec.Command("xdg-open", rawURL) + default: + return fmt.Errorf("unsupported operating system %s", runtime.GOOS) + } + if err := command.Start(); err != nil { + return err + } + return command.Process.Release() +} diff --git a/agents/go-common/gohive/browser_auth_test.go b/agents/go-common/gohive/browser_auth_test.go new file mode 100644 index 0000000000..61fbb8976c --- /dev/null +++ b/agents/go-common/gohive/browser_auth_test.go @@ -0,0 +1,116 @@ +package gohive + +import ( + "context" + "net/http" + "net/url" + "strconv" + "strings" + "testing" + "time" +) + +func TestBrowserSSOClientAuthenticatesThroughLoopbackCallback(t *testing.T) { + client := newBrowserSSOClient(0, time.Second) + if err := client.Start(); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = client.Close() }) + if err := client.SetRedirect("https://identity.example.com/sso?RelayState=test", "browser-client"); err != nil { + t.Fatal(err) + } + client.openURL = func(rawURL string) error { + if rawURL != "https://identity.example.com/sso?RelayState=test" { + t.Fatalf("unexpected browser URL: %q", rawURL) + } + response, err := http.PostForm( + "http://127.0.0.1:"+strconv.Itoa(client.Port())+"/", + url.Values{"status": {"true"}, "token": {"signed-browser-token"}}, + ) + if err != nil { + return err + } + return response.Body.Close() + } + if err := client.Authenticate(context.Background()); err != nil { + t.Fatal(err) + } + token, clientIdentifier := client.Credentials() + if token != "signed-browser-token" || clientIdentifier != "browser-client" { + t.Fatalf("unexpected browser credentials: token=%q client=%q", token, clientIdentifier) + } +} + +func TestBrowserAuthRoundTripperCompletesRedirectAndBearerFlow(t *testing.T) { + client := newBrowserSSOClient(0, time.Second) + if err := client.Start(); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = client.Close() }) + calls := 0 + transport := &browserAuthRoundTripper{ + Client: client, + Base: roundTripFunc(func(request *http.Request) (*http.Response, error) { + calls++ + if calls == 1 { + if request.Header.Get("X-Hive-Token-Response-Port") != strconv.Itoa(client.Port()) || request.Header.Get("Authorization") != "" { + t.Fatalf("unexpected initial browser headers: %#v", request.Header) + } + response := httpResponse(request, http.StatusFound) + response.Header.Set("Location", "https://identity.example.com/sso") + response.Header.Set("X-Hive-Client-Identifier", "browser-client") + return response, nil + } + if request.Header.Get("Authorization") != "Bearer signed-browser-token" || request.Header.Get("X-Hive-Client-Identifier") != "browser-client" || request.Header.Get("X-Hive-Token-Response-Port") != "" { + t.Fatalf("unexpected authenticated browser headers: %#v", request.Header) + } + return httpResponse(request, http.StatusOK), nil + }), + } + request, err := http.NewRequest(http.MethodPost, "https://hs2.example.com/cliservice", strings.NewReader("payload")) + if err != nil { + t.Fatal(err) + } + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if !client.HasRedirect() { + t.Fatal("browser redirect was not captured") + } + client.openURL = func(string) error { + client.responses <- browserSSOResponse{Successful: true, Token: "signed-browser-token"} + return nil + } + if err := client.Authenticate(context.Background()); err != nil { + t.Fatal(err) + } + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if calls != 2 || request.Header.Get("Authorization") != "" || request.Header.Get("X-Hive-Token-Response-Port") != "" { + t.Fatalf("browser transport mutated request or skipped a call: calls=%d headers=%#v", calls, request.Header) + } +} + +func TestBrowserSSORedirectValidation(t *testing.T) { + client := newBrowserSSOClient(0, time.Second) + for _, rawURL := range []string{"", "/relative", "file:///tmp/token"} { + if err := client.SetRedirect(rawURL, "browser-client"); err == nil { + t.Fatalf("expected redirect %q to be rejected", rawURL) + } + } + if err := client.SetRedirect("https://identity.example.com/sso", ""); err == nil { + t.Fatal("expected missing client identifier to be rejected") + } +} + +func TestBrowserAuthenticationRequiresHTTPAndTLS(t *testing.T) { + _, err := innerConnect(context.Background(), "127.0.0.1", 1, "BROWSER", &connectConfiguration{}) + if err == nil || !strings.Contains(err.Error(), "HTTP transport") { + t.Fatalf("unexpected binary browser error: %v", err) + } + _, err = innerConnect(context.Background(), "127.0.0.1", 1, "BROWSER", &connectConfiguration{TransportMode: "http"}) + if err == nil || !strings.Contains(err.Error(), "ssl=true") { + t.Fatalf("unexpected insecure browser error: %v", err) + } +} diff --git a/agents/go-common/gohive/connector.go b/agents/go-common/gohive/connector.go new file mode 100644 index 0000000000..f11e2f8a2d --- /dev/null +++ b/agents/go-common/gohive/connector.go @@ -0,0 +1,153 @@ +package gohive + +import ( + "context" + "crypto/tls" + "crypto/x509" + "database/sql" + "database/sql/driver" + "fmt" + "os" + "time" + + "github.com/beltran/gosasl" +) + +// Config holds everything needed to connect to Hive/Impala via gohive v2. +type Config struct { + Host string + Port int + Auth string // "NONE", "KERBEROS", "NOSASL", etc. + Username string + Password string + Database string + TransportMode string // "binary" or "http" + HTTPPath string + Service string // Kerberos service name + HTTPKerberosChannelBinding bool + GSSAPIOptions gosasl.GSSAPIOptions + TLSConfig *tls.Config + SSLCertFile string + SSLKeyFile string + SSLCAFile string + SSLInsecureSkip bool + HiveConfiguration map[string]string + ConnectTimeout time.Duration + SocketTimeout time.Duration + HTTPTimeout time.Duration + FetchSize int64 + MaxMessageSize int32 + HTTPHeaders map[string]string + HTTPCookies map[string]string + RequestTracking bool + DisableCookieAuth bool + CookieName string + JWT string + DelegationToken string + BrowserToken string + BrowserClientID string + BrowserResponsePort int + BrowserResponseTimeout time.Duration + BrowserDisableSSLCheck bool +} + +var _ driver.Connector = (*HiveConnector)(nil) + +// HiveConnector implements driver.Connector using gohive v2 under the hood. +type HiveConnector struct { + cfg Config +} + +// NewConnector creates a new HiveConnector with the given config. +func NewConnector(cfg Config) *HiveConnector { + return &HiveConnector{cfg: cfg} +} + +// OpenDB is a convenience that returns a *sql.DB ready for use with GORM. +func OpenDB(cfg Config) *sql.DB { + return sql.OpenDB(NewConnector(cfg)) +} + +// Connect establishes a new connection using gohive v2. +func (c *HiveConnector) Connect(ctx context.Context) (driver.Conn, error) { + connCfg := newConnectConfiguration() + connCfg.Username = c.cfg.Username + connCfg.Password = c.cfg.Password + connCfg.Database = c.cfg.Database + if c.cfg.TransportMode != "" { + connCfg.TransportMode = c.cfg.TransportMode + } + if c.cfg.HTTPPath != "" { + connCfg.HTTPPath = c.cfg.HTTPPath + } + connCfg.Service = c.cfg.Service + if connCfg.Service == "" { + connCfg.Service = "hive" + } + connCfg.TLSConfig = c.cfg.TLSConfig + connCfg.HTTPKerberosChannelBinding = c.cfg.HTTPKerberosChannelBinding + connCfg.GSSAPIOptions = c.cfg.GSSAPIOptions + connCfg.HiveConfiguration = c.cfg.HiveConfiguration + connCfg.ConnectTimeout = c.cfg.ConnectTimeout + connCfg.SocketTimeout = c.cfg.SocketTimeout + connCfg.HttpTimeout = c.cfg.HTTPTimeout + if c.cfg.FetchSize > 0 { + connCfg.FetchSize = c.cfg.FetchSize + } + if c.cfg.MaxMessageSize > 0 { + connCfg.MaxMessageSize = c.cfg.MaxMessageSize + connCfg.MaxSize = uint32(c.cfg.MaxMessageSize) + } + connCfg.HTTPHeaders = c.cfg.HTTPHeaders + connCfg.HTTPCookies = c.cfg.HTTPCookies + connCfg.RequestTracking = c.cfg.RequestTracking + connCfg.DisableCookieAuth = c.cfg.DisableCookieAuth + connCfg.CookieName = c.cfg.CookieName + connCfg.JWT = c.cfg.JWT + connCfg.DelegationToken = c.cfg.DelegationToken + connCfg.BrowserToken = c.cfg.BrowserToken + connCfg.BrowserClientID = c.cfg.BrowserClientID + connCfg.BrowserResponsePort = c.cfg.BrowserResponsePort + connCfg.BrowserResponseTimeout = c.cfg.BrowserResponseTimeout + connCfg.BrowserDisableSSLCheck = c.cfg.BrowserDisableSSLCheck + + // Fallback: build TLS config from cert/key files if TLSConfig not provided directly + if connCfg.TLSConfig == nil { + + if c.cfg.SSLCertFile != "" && c.cfg.SSLKeyFile != "" { + tlsConfig, err := getTlsConfiguration(c.cfg.SSLCertFile, c.cfg.SSLKeyFile) + if err != nil { + return nil, fmt.Errorf("failed to configure SSL: %v", err) + } + tlsConfig.InsecureSkipVerify = c.cfg.SSLInsecureSkip + connCfg.TLSConfig = tlsConfig + } else if c.cfg.SSLCAFile != "" { + + caPEM, err := os.ReadFile(c.cfg.SSLCAFile) + if err != nil { + return nil, fmt.Errorf("invalid ca path: %s (%w)", c.cfg.SSLCAFile, err) + } + pool := x509.NewCertPool() + if !pool.AppendCertsFromPEM(caPEM) { + return nil, fmt.Errorf("invalid certification %q", c.cfg.SSLCAFile) + } + connCfg.TLSConfig = &tls.Config{ + RootCAs: pool, + InsecureSkipVerify: c.cfg.SSLInsecureSkip, + } + } else if c.cfg.SSLInsecureSkip { + connCfg.TLSConfig = &tls.Config{InsecureSkipVerify: true} + } + } + + conn, err := connect(ctx, c.cfg.Host, c.cfg.Port, c.cfg.Auth, connCfg) + if err != nil { + return nil, fmt.Errorf("gohive connect: %w", err) + } + return &sqlConnection{conn: conn}, nil +} + +// Driver returns a placeholder driver (unused by sql.OpenDB). +func (c *HiveConnector) Driver() driver.Driver { + return &Driver{} +} diff --git a/agents/go-common/gohive/driver.go b/agents/go-common/gohive/driver.go new file mode 100644 index 0000000000..2e47ca3b83 --- /dev/null +++ b/agents/go-common/gohive/driver.go @@ -0,0 +1,346 @@ +package gohive + +import ( + "context" + "database/sql" + "database/sql/driver" + "fmt" + "io" + "reflect" + "strings" + "time" +) + +const operationCloseTimeout = 5 * time.Second + +type fetchSizeContextKey struct{} + +// NonQueryResult reports that HiveServer2 executed a statement without a +// result set. QueryContext returns this value so callers can mirror JDBC +// Statement.execute() without guessing from the SQL text. +type NonQueryResult struct { + AffectedRows int64 +} + +func (result *NonQueryResult) Error() string { + return "Hive statement completed without a result set" +} + +// WithFetchSize applies the HiveServer2 fetch batch size to a single +// statement without mutating the shared connection configuration. +func WithFetchSize(ctx context.Context, fetchSize int) context.Context { + if fetchSize <= 0 { + return ctx + } + return context.WithValue(ctx, fetchSizeContextKey{}, int64(fetchSize)) +} + +// Driver is the interface that must be implemented by a database driver. +type Driver struct{} + +// Open returns a new connection to the database. +// The name is a string in a driver-specific format. +func (d *Driver) Open(name string) (driver.Conn, error) { + connector, err := d.OpenConnector(name) + if err != nil { + return nil, err + } + return connector.Connect(context.Background()) +} + +// OpenConnector implements driver.DriverContext +func (d *Driver) OpenConnector(name string) (driver.Connector, error) { + // Parse the DSN + dsn, err := ParseDSN(name) + if err != nil { + return nil, err + } + + // Create configuration + config := newConnectConfiguration() + config.Username = dsn.Username + config.Password = dsn.Password + config.Database = dsn.Database + config.TransportMode = dsn.TransportMode + config.Service = dsn.Service + config.HiveConfiguration = dsn.HiveConfiguration + + // Configure SSL if paths are provided + if dsn.SSLCertFile != "" && dsn.SSLKeyFile != "" { + tlsConfig, err := getTlsConfiguration(dsn.SSLCertFile, dsn.SSLKeyFile) + if err != nil { + return nil, fmt.Errorf("failed to configure SSL: %v", err) + } + tlsConfig.InsecureSkipVerify = dsn.SSLInsecureSkip + config.TLSConfig = tlsConfig + } + + return &connector{host: dsn.Host, port: dsn.Port, auth: dsn.Auth, config: config}, nil +} + +// connector implements driver.Connector +type connector struct { + host string + port int + auth string + config *connectConfiguration +} + +// Connect returns a connection to the database. +func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { + // Connect to Hive + conn, err := connect(ctx, c.host, c.port, c.auth, c.config) + if err != nil { + return nil, err + } + return &sqlConnection{conn: conn}, nil +} + +// Driver returns the underlying Driver of the Connector. +func (c *connector) Driver() driver.Driver { + return &Driver{} +} + +// sqlConnection implements driver.Conn +type sqlConnection struct { + conn *connection +} + +// Close invalidates and potentially stops any current +// prepared statements and transactions, marking this +// connection as no longer in use. +func (c *sqlConnection) Close() error { + return c.conn.close() +} + +// Begin starts and returns a new transaction. +func (c *sqlConnection) Begin() (driver.Tx, error) { + return c.BeginTx(context.Background(), driver.TxOptions{}) +} + +// BeginTx starts and returns a new transaction. +func (c *sqlConnection) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) { + return nil, driver.ErrSkip +} + +// Prepare returns a prepared statement, bound to this connection. +func (c *sqlConnection) Prepare(query string) (driver.Stmt, error) { + return nil, fmt.Errorf("prepared statements are not supported by Hive") +} + +// PrepareContext returns a prepared statement, bound to this connection. +func (c *sqlConnection) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) { + return nil, fmt.Errorf("prepared statements are not supported by Hive") +} + +// Exec executes a query that doesn't return rows. +// Implements driver.Execer +func (c *sqlConnection) Exec(query string, args []driver.Value) (driver.Result, error) { + return c.ExecContext(context.Background(), query, args) +} + +// ExecContext executes a query that doesn't return rows. +func (c *sqlConnection) ExecContext(ctx context.Context, query string, args []driver.Value) (driver.Result, error) { + if len(args) > 0 { + return nil, fmt.Errorf("query parameters are not supported by Hive") + } + + cursor := c.conn.cursor() + cursor.fetchSize = fetchSizeFromContext(ctx) + cursor.exec(ctx, query) + if cursor.error() != nil { + return nil, cursor.error() + } + affectedRows := cursorAffectedRows(cursor) + if err := closeCursor(cursor); err != nil { + return nil, err + } + return &result{affectedRows: affectedRows}, nil +} + +// Query executes a query that may return rows. +// Implements driver.Queryer +func (c *sqlConnection) Query(query string, args []driver.Value) (driver.Rows, error) { + return c.QueryContext(context.Background(), query, args) +} + +// QueryContext executes a query that may return rows. +func (c *sqlConnection) QueryContext(ctx context.Context, query string, args []driver.Value) (driver.Rows, error) { + if len(args) > 0 { + return nil, fmt.Errorf("query parameters are not supported by Hive") + } + + cursor := c.conn.cursor() + cursor.fetchSize = fetchSizeFromContext(ctx) + cursor.exec(ctx, query) + if cursor.error() != nil { + return nil, cursor.error() + } + if cursor.operationHandle == nil { + return nil, fmt.Errorf("HiveServer2 returned no operation handle") + } + if !cursor.operationHandle.HasResultSet { + result := &NonQueryResult{AffectedRows: cursorAffectedRows(cursor)} + if err := closeCursor(cursor); err != nil { + return nil, err + } + return nil, result + } + return &rows{cursor: cursor, ctx: ctx, descriptionValid: false}, nil +} + +// result implements driver.Result +type result struct { + affectedRows int64 +} + +// LastInsertId returns the database's auto-generated ID +// after, for example, an INSERT into a table with primary key. +func (r *result) LastInsertId() (int64, error) { + return 0, driver.ErrSkip +} + +// RowsAffected returns the number of rows affected by the query. +func (r *result) RowsAffected() (int64, error) { + return r.affectedRows, nil +} + +// rows implements driver.Rows +type rows struct { + cursor *cursor + ctx context.Context + description [][]string + descriptionValid bool +} + +// Columns returns the names of the columns. +func (r *rows) Columns() []string { + if !r.descriptionValid { + r.description = r.cursor.description(r.ctx) + r.descriptionValid = true + } + columns := make([]string, len(r.description)) + for i, col := range r.description { + columns[i] = col[0] + } + return columns +} + +// Close closes the rows iterator. +func (r *rows) Close() error { + if r.cursor != nil { + return closeCursor(r.cursor) + } + return nil +} + +// Next is called to populate the next row of data into +// the provided slice. +func (r *rows) Next(dest []driver.Value) error { + if r.cursor == nil { + return io.EOF + } + + if !r.cursor.hasMore(r.ctx) { + return io.EOF + } + + // Fetch the row directly into the destination slice + r.cursor.fetchOneDriver(r.ctx, dest) + if r.cursor.Err != nil { + return r.cursor.Err + } + + if !r.descriptionValid { + r.description = r.cursor.description(r.ctx) + r.descriptionValid = true + } + + return nil +} + +// ColumnTypeScanType returns the scan type for the given column index +func (r *rows) ColumnTypeScanType(index int) reflect.Type { + if r.cursor == nil { + return nil + } + if !r.descriptionValid { + r.description = r.cursor.description(r.ctx) + r.descriptionValid = true + } + if r.cursor.Err != nil || index >= len(r.description) { + return nil + } + colType := strings.TrimSuffix(strings.ToUpper(r.description[index][1]), "_TYPE") + switch colType { + case "BOOLEAN": + return reflect.TypeOf(false) + case "TINYINT": + return reflect.TypeOf(int8(0)) + case "SMALLINT": + return reflect.TypeOf(int16(0)) + case "INT": + return reflect.TypeOf(int32(0)) + case "BIGINT": + return reflect.TypeOf(int64(0)) + case "FLOAT": + return reflect.TypeOf(float32(0)) // Return float32 for FLOAT type + case "DOUBLE": + return reflect.TypeOf(float64(0)) + case "STRING", "VARCHAR", "CHAR": + return reflect.TypeOf("") + case "TIMESTAMP": + return reflect.TypeOf("") + case "DATE": + return reflect.TypeOf("") + case "BINARY": + return reflect.TypeOf([]byte{}) + default: + return reflect.TypeOf("") + } +} + +// ColumnTypeDatabaseTypeName returns the database system type name. +func (r *rows) ColumnTypeDatabaseTypeName(index int) string { + if !r.descriptionValid { + r.description = r.cursor.description(r.ctx) + r.descriptionValid = true + } + if index >= len(r.description) { + return "" + } + return strings.TrimSuffix(r.description[index][1], "_TYPE") +} + +func cursorAffectedRows(cursor *cursor) int64 { + if cursor == nil || cursor.operationHandle == nil || cursor.operationHandle.ModifiedRowCount == nil { + return 0 + } + count := int64(*cursor.operationHandle.ModifiedRowCount) + if count < 0 { + return 0 + } + return count +} + +func closeCursor(cursor *cursor) error { + if cursor == nil { + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), operationCloseTimeout) + defer cancel() + cursor.close(ctx) + return cursor.error() +} + +func fetchSizeFromContext(ctx context.Context) int64 { + if ctx == nil { + return 0 + } + value, _ := ctx.Value(fetchSizeContextKey{}).(int64) + return value +} + +func init() { + sql.Register("hive", &Driver{}) +} diff --git a/agents/go-common/gohive/driver_test.go b/agents/go-common/gohive/driver_test.go new file mode 100644 index 0000000000..6b43efd2de --- /dev/null +++ b/agents/go-common/gohive/driver_test.go @@ -0,0 +1,53 @@ +package gohive + +import ( + "context" + "testing" + + "github.com/beltran/gohive/v2/hiveserver" +) + +func TestWithFetchSize(t *testing.T) { + ctx := WithFetchSize(context.Background(), 512) + if value := fetchSizeFromContext(ctx); value != 512 { + t.Fatalf("unexpected fetch size: %d", value) + } + if value := fetchSizeFromContext(WithFetchSize(ctx, 0)); value != 512 { + t.Fatalf("non-positive fetch size should preserve context value, got %d", value) + } +} + +func TestCursorEffectiveFetchSize(t *testing.T) { + connection := &connection{configuration: &connectConfiguration{FetchSize: 1000}} + if value := (&cursor{conn: connection}).effectiveFetchSize(); value != 1000 { + t.Fatalf("unexpected default fetch size: %d", value) + } + if value := (&cursor{conn: connection, fetchSize: 128}).effectiveFetchSize(); value != 128 { + t.Fatalf("unexpected statement fetch size: %d", value) + } +} + +func TestCursorAffectedRows(t *testing.T) { + modified := float64(7) + value := cursorAffectedRows(&cursor{operationHandle: &hiveserver.TOperationHandle{ + HasResultSet: false, + ModifiedRowCount: &modified, + }}) + if value != 7 { + t.Fatalf("unexpected affected rows: %d", value) + } +} + +func TestCursorAffectedRowsNormalizesMissingAndNegativeValues(t *testing.T) { + negative := float64(-1) + for _, cursor := range []*cursor{ + nil, + {}, + {operationHandle: &hiveserver.TOperationHandle{HasResultSet: false}}, + {operationHandle: &hiveserver.TOperationHandle{HasResultSet: false, ModifiedRowCount: &negative}}, + } { + if value := cursorAffectedRows(cursor); value != 0 { + t.Fatalf("expected zero affected rows, got %d", value) + } + } +} diff --git a/agents/go-common/gohive/dsn.go b/agents/go-common/gohive/dsn.go new file mode 100644 index 0000000000..7ca5301502 --- /dev/null +++ b/agents/go-common/gohive/dsn.go @@ -0,0 +1,107 @@ +package gohive + +import ( + "fmt" + "net/url" + "strconv" + "strings" +) + +// DSN represents a parsed Data Source Name +type DSN struct { + Username string + Password string + Host string + Port int + Database string + Auth string + TransportMode string + Service string + SSLCertFile string + SSLKeyFile string + SSLInsecureSkip bool + HiveConfiguration map[string]string +} + +// ParseDSN parses a DSN string into a DSN struct +func ParseDSN(dsn string) (*DSN, error) { + // Check if the DSN starts with "hive://" + if !strings.HasPrefix(dsn, "hive://") { + return nil, fmt.Errorf("DSN must start with 'hive://'") + } + + // Do NOT remove the "hive://" prefix; let url.Parse handle it + u, err := url.Parse(dsn) + if err != nil { + return nil, fmt.Errorf("invalid DSN format: %v", err) + } + + // Create a new DSN struct + parsedDSN := &DSN{ + Database: strings.TrimPrefix(u.Path, "/"), + Auth: "NONE", + TransportMode: "binary", // Default transport mode + Service: "hive", // Default service + HiveConfiguration: make(map[string]string), + } + + // Parse username and password if present + if u.User != nil { + parsedDSN.Username = u.User.Username() + if password, ok := u.User.Password(); ok { + parsedDSN.Password = password + } + } + + // Parse host and port + parsedDSN.Host = u.Hostname() + if portStr := u.Port(); portStr != "" { + port, err := strconv.Atoi(portStr) + if err != nil { + return nil, fmt.Errorf("invalid port number: %v", err) + } + parsedDSN.Port = port + } else { + parsedDSN.Port = 10000 // Default port + } + + // Parse query parameters + query := u.Query() + if auth := query.Get("auth"); auth != "" { + parsedDSN.Auth = auth + } + if transport := query.Get("transport"); transport != "" { + parsedDSN.TransportMode = transport + } + if service := query.Get("service"); service != "" { + parsedDSN.Service = service + } + if sslCert := query.Get("sslcert"); sslCert != "" { + parsedDSN.SSLCertFile = sslCert + } + if sslKey := query.Get("sslkey"); sslKey != "" { + parsedDSN.SSLKeyFile = sslKey + } + if sslInsecureSkip := query.Get("sslinsecureskipverify"); sslInsecureSkip != "" { + parsedDSN.SSLInsecureSkip = sslInsecureSkip == "true" + } + + // Collect all other query parameters as Hive configuration properties + for key, values := range query { + // Skip standard parameters + if key == "auth" || key == "transport" || key == "service" || + key == "sslcert" || key == "sslkey" || key == "sslinsecureskipverify" { + continue + } + if len(values) > 0 { + parsedDSN.HiveConfiguration[key] = values[0] + } + } + + // Validate required fields + if parsedDSN.Database == "" { + return nil, fmt.Errorf("database name is required") + } + + return parsedDSN, nil +} diff --git a/agents/go-common/gohive/dsn_test.go b/agents/go-common/gohive/dsn_test.go new file mode 100644 index 0000000000..14e491660e --- /dev/null +++ b/agents/go-common/gohive/dsn_test.go @@ -0,0 +1,23 @@ +package gohive + +import "testing" + +func TestParseDSNDefaultsToPlainAuthentication(t *testing.T) { + value, err := ParseDSN("hive://hs2.example.com/default") + if err != nil { + t.Fatal(err) + } + if value.Auth != "NONE" || value.TransportMode != "binary" || value.Port != 10000 { + t.Fatalf("unexpected DSN defaults: %#v", value) + } +} + +func TestParseDSNPreservesUserInfo(t *testing.T) { + value, err := ParseDSN("hive://user%40example.com:p%40ss%3Aword@hs2.example.com:10001/default?auth=LDAP&transport=http") + if err != nil { + t.Fatal(err) + } + if value.Username != "user@example.com" || value.Password != "p@ss:word" || value.Auth != "LDAP" || value.TransportMode != "http" { + t.Fatalf("unexpected DSN: %#v", value) + } +} diff --git a/agents/go-common/gohive/go.mod b/agents/go-common/gohive/go.mod new file mode 100644 index 0000000000..e573a28163 --- /dev/null +++ b/agents/go-common/gohive/go.mod @@ -0,0 +1,28 @@ +module github.com/t8y2/dbx/agents/go-common/gohive + +go 1.23.0 + +require ( + github.com/apache/thrift v0.22.0 + github.com/beltran/gohive/v2 v2.1.0 + github.com/beltran/gosasl v1.0.0 + github.com/pkg/errors v0.9.1 + golang.org/x/net v0.41.0 +) + +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.39.0 // indirect +) + +replace github.com/beltran/gosasl => ../gosasl + +replace github.com/golang-auth/go-gssapi/v2 => ../go-gssapi diff --git a/agents/go-common/gohive/go.sum b/agents/go-common/gohive/go.sum new file mode 100644 index 0000000000..7312aa360f --- /dev/null +++ b/agents/go-common/gohive/go.sum @@ -0,0 +1,56 @@ +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/apache/thrift v0.22.0 h1:r7mTJdj51TMDe6RtcmNdQxgn9XcyfGDOzegMDRg47uc= +github.com/apache/thrift v0.22.0/go.mod h1:1e7J/O1Ae6ZQMTYdy9xa3w9k+XHWPfRvdPyJeynQ+/g= +github.com/beltran/gohive/v2 v2.1.0 h1:+XPfODaLXsZkD8mEnKCYqinC4OmVMx4/oR4ys8IGKeI= +github.com/beltran/gohive/v2 v2.1.0/go.mod h1:MrxyI0sdriG52I+Bo2lJpGR75bxjcI/Hx3+MBnYqR10= +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/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/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +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/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/crypto v0.39.0 h1:SHs+kF4LP+f+p14esP5jAoDpHU8Gu/v9lFRK6IT5imM= +golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U= +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/go.mod h1:uSPa2vr4CLtc/ILN5odXGNXS6mhrKVzTaCXzk9m6W3k= +golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw= +golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA= +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/go-common/gohive/hive.go b/agents/go-common/gohive/hive.go new file mode 100644 index 0000000000..76e9ca5bc7 --- /dev/null +++ b/agents/go-common/gohive/hive.go @@ -0,0 +1,1650 @@ +package gohive + +import ( + "context" + "crypto/tls" + "fmt" + "log" + "net" + "net/http" + "net/http/cookiejar" + "net/url" + "os/user" + "strconv" + "strings" + "sync" + "time" + + "database/sql/driver" + + "github.com/apache/thrift/lib/go/thrift" + "github.com/beltran/gohive/v2/hiveserver" + "github.com/beltran/gosasl" + "github.com/pkg/errors" + "golang.org/x/net/publicsuffix" +) + +const defaultFetchSize int64 = 1000 +const zookeeperDefaultNamespace = "hiveserver2" +const defaultMaxLength = 16384000 + +// Cursor states +const ( + _NONE = iota + _RUNNING + _FINISHED + _ERROR + _CONTEXT_DONE + _ASYNC_ENDED +) + +type dialContextFunc func(ctx context.Context, network, addr string) (net.Conn, error) + +type basicAuthRoundTripper struct { + Base http.RoundTripper + Username string + Password string +} + +type bearerAuthRoundTripper struct { + Base http.RoundTripper + Token string + ClientIdentifier string +} + +func (transport *bearerAuthRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + base := transport.Base + if base == nil { + base = http.DefaultTransport + } + attempt := request.Clone(request.Context()) + attempt.Header = request.Header.Clone() + attempt.Header.Set("Authorization", "Bearer "+transport.Token) + if transport.ClientIdentifier != "" { + attempt.Header.Set("X-Hive-Client-Identifier", transport.ClientIdentifier) + } + return base.RoundTrip(attempt) +} + +type headerAuthRoundTripper struct { + Base http.RoundTripper + Name string + Value string +} + +func (transport *headerAuthRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + base := transport.Base + if base == nil { + base = http.DefaultTransport + } + attempt := request.Clone(request.Context()) + attempt.Header = request.Header.Clone() + attempt.Header.Set(transport.Name, transport.Value) + return base.RoundTrip(attempt) +} + +type customHTTPRoundTripper struct { + Base http.RoundTripper + Headers map[string]string + Cookies map[string]string + RequestTracker *httpRequestTracker +} + +type cookieAuthRoundTripper struct { + Base http.RoundTripper + Auth http.RoundTripper + CookieName string + HasStaticCookie bool +} + +func (transport *cookieAuthRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + if !transport.HasStaticCookie && !hasNamedCookie(request, transport.CookieName) { + return transport.Auth.RoundTrip(request) + } + if request.Body != nil && request.GetBody == nil { + return transport.Auth.RoundTrip(request) + } + response, err := transport.Base.RoundTrip(request) + if err != nil || response.StatusCode != http.StatusUnauthorized { + return response, err + } + if response.Body != nil { + _ = response.Body.Close() + } + retry := request.Clone(request.Context()) + retry.Header = request.Header.Clone() + if request.GetBody != nil { + retry.Body, err = request.GetBody() + if err != nil { + return nil, err + } + } + return transport.Auth.RoundTrip(retry) +} + +func hasNamedCookie(request *http.Request, name string) bool { + if strings.TrimSpace(name) == "" { + return false + } + _, err := request.Cookie(name) + return err == nil +} + +func withCookieAuthentication(configuration *connectConfiguration, base, auth http.RoundTripper) http.RoundTripper { + if configuration.DisableCookieAuth { + return auth + } + return &cookieAuthRoundTripper{ + Base: base, + Auth: auth, + CookieName: configuration.CookieName, + HasStaticCookie: hasConfiguredCookie(configuration.HTTPCookies, configuration.CookieName), + } +} + +func hasConfiguredCookie(cookies map[string]string, name string) bool { + if strings.TrimSpace(name) == "" { + return false + } + _, exists := cookies[name] + return exists +} + +type httpRequestTracker struct { + mu sync.Mutex + sessionID string + counter uint64 +} + +func newHTTPRequestTracker() *httpRequestTracker { + return &httpRequestTracker{sessionID: "NO_SESSION"} +} + +func (tracker *httpRequestTracker) next() string { + tracker.mu.Lock() + defer tracker.mu.Unlock() + tracker.counter++ + return fmt.Sprintf("HIVE_%s_%020d", tracker.sessionID, tracker.counter) +} + +func (tracker *httpRequestTracker) setSessionHandle(handle *hiveserver.TSessionHandle) { + if handle == nil || handle.SessionId == nil || len(handle.SessionId.GUID) == 0 { + return + } + tracker.mu.Lock() + tracker.sessionID = fmt.Sprintf("%x", handle.SessionId.GUID) + tracker.mu.Unlock() +} + +func (transport *customHTTPRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + base := transport.Base + if base == nil { + base = http.DefaultTransport + } + attempt := request.Clone(request.Context()) + attempt.Header = request.Header.Clone() + for name, value := range transport.Headers { + attempt.Header.Add(name, value) + } + for name, value := range transport.Cookies { + attempt.AddCookie(&http.Cookie{Name: name, Value: value}) + } + attempt.Header.Add("X-XSRF-HEADER", "true") + attempt.Header.Add("X-CSRF-TOKEN", "true") + if transport.RequestTracker != nil { + attempt.Header.Set("X-Request-ID", transport.RequestTracker.next()) + } + return base.RoundTrip(attempt) +} + +func (transport *basicAuthRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + base := transport.Base + if base == nil { + base = http.DefaultTransport + } + attempt := request.Clone(request.Context()) + attempt.Header = request.Header.Clone() + attempt.SetBasicAuth(transport.Username, transport.Password) + return base.RoundTrip(attempt) +} + +func hiveHTTPURL(protocol, host string, port int, path string) string { + return (&url.URL{ + Scheme: protocol, + Host: net.JoinHostPort(host, strconv.Itoa(port)), + Path: "/" + strings.Trim(path, "/"), + }).String() +} + +// connection holds the information for getting a cursor to hive. +type connection struct { + host string + port int + username string + database string + auth string + kerberosServiceName string + password string + sessionHandle *hiveserver.TSessionHandle + client *hiveserver.TCLIServiceClient + configuration *connectConfiguration + transport thrift.TTransport + authCloser interface{ Close() error } + mu sync.Mutex // Mutex to protect connection operations + clientMu sync.Mutex // Mutex to protect client operations +} + +// connectConfiguration is the configuration for the connection +// The fields have to be filled manually but not all of them are required +// Depends on the auth and kind of connection. +type connectConfiguration struct { + Username string + Principal string + Password string + Service string + HiveConfiguration map[string]string + PollIntervalInMillis int + FetchSize int64 + TransportMode string + HTTPPath string + TLSConfig *tls.Config + ZookeeperNamespace string + Database string + ConnectTimeout time.Duration + SocketTimeout time.Duration + HttpTimeout time.Duration + DialContext dialContextFunc + DisableKeepAlives bool + HTTPKerberosChannelBinding bool + GSSAPIOptions gosasl.GSSAPIOptions + HTTPHeaders map[string]string + HTTPCookies map[string]string + RequestTracking bool + RequestTracker *httpRequestTracker + DisableCookieAuth bool + CookieName string + JWT string + DelegationToken string + BrowserToken string + BrowserClientID string + BrowserResponsePort int + BrowserResponseTimeout time.Duration + BrowserDisableSSLCheck bool + // Maximum length of the data in bytes. Used for SASL. + MaxSize uint32 + MaxMessageSize int32 +} + +// newConnectConfiguration returns a connect configuration, all with empty fields +func newConnectConfiguration() *connectConfiguration { + return &connectConfiguration{ + Username: "", + Password: "", + Service: "", + HiveConfiguration: nil, + PollIntervalInMillis: 200, + FetchSize: defaultFetchSize, + TransportMode: "binary", + HTTPPath: "cliservice", + TLSConfig: nil, + ZookeeperNamespace: zookeeperDefaultNamespace, + CookieName: "hive.server2.auth", + MaxSize: defaultMaxLength, + } +} + +// Error represents an error surfaced from HiveServer2. +type Error struct { + Err error + + // Simple error message, without the full stack trace. Surfaced from Thrift. + Message string + // See https://github.com/apache/hive/blob/master/common/src/java/org/apache/hadoop/hive/ql/ErrorMsg.java for info about error codes. + ErrorCode int + SQLState string +} + +func (err *Error) Error() string { + if err == nil || err.Err == nil { + return "HiveServer2 error" + } + return err.Err.Error() +} + +func (err *Error) Unwrap() error { + if err == nil { + return nil + } + return err.Err +} + +// connect to hive server +func connect(ctx context.Context, host string, port int, auth string, + configuration *connectConfiguration) (conn *connection, err error) { + return innerConnect(ctx, host, port, auth, configuration) +} + +func parseHiveServer2Info(hsInfos []string) []map[string]string { + results := make([]map[string]string, len(hsInfos)) + actualCount := 0 + + for _, hsInfo := range hsInfos { + validFormat := false + node := make(map[string]string) + + for _, param := range strings.Split(hsInfo, ";") { + kvPair := strings.Split(param, "=") + if len(kvPair) < 2 { + break + } + if kvPair[0] == "serverUri" { + hostAndPort := strings.Split(kvPair[1], ":") + if len(hostAndPort) == 2 { + node["host"] = hostAndPort[0] + node["port"] = hostAndPort[1] + validFormat = len(node["host"]) != 0 && len(node["port"]) != 0 + } else { + break + } + } else { + node[kvPair[0]] = kvPair[1] + } + } + if validFormat { + results[actualCount] = node + actualCount++ + } + } + return results[0:actualCount] +} + +func dial(ctx context.Context, addr string, dialFn dialContextFunc, timeout time.Duration) (net.Conn, error) { + dctx := ctx + if timeout > 0 { + var cancel context.CancelFunc + dctx, cancel = context.WithTimeout(ctx, timeout) + defer cancel() + } + return dialFn(dctx, "tcp", addr) +} + +func innerConnect(ctx context.Context, host string, port int, auth string, + configuration *connectConfiguration) (conn *connection, err error) { + if configuration == nil { + configuration = newConnectConfiguration() + } + if auth == "BROWSER" { + if configuration.TransportMode != "http" { + return nil, errors.New("Hive browser authentication requires HTTP transport mode") + } + if configuration.TLSConfig == nil && !configuration.BrowserDisableSSLCheck { + return nil, errors.New("Hive browser authentication requires ssl=true unless browserDisableSslCheck=true") + } + } + var authCloser interface{ Close() error } + var browserClient *browserSSOClient + var transport thrift.TTransport + defer func() { + if err != nil { + if transport != nil { + _ = transport.Close() + } + if authCloser != nil { + _ = authCloser.Close() + } + } + }() + + var socket thrift.TTransport + addr := fmt.Sprintf("%s:%d", host, port) + if configuration.DialContext != nil { + var netConn net.Conn + netConn, err = dial(ctx, addr, configuration.DialContext, configuration.ConnectTimeout) + if err != nil { + return + } + if configuration.TLSConfig != nil { + socket = thrift.NewTSSLSocketFromConnConf(netConn, &thrift.TConfiguration{ + ConnectTimeout: configuration.ConnectTimeout, + SocketTimeout: configuration.SocketTimeout, + MaxMessageSize: configuration.MaxMessageSize, + TLSConfig: configuration.TLSConfig, + }) + } else { + socket = thrift.NewTSocketFromConnConf(netConn, &thrift.TConfiguration{ + ConnectTimeout: configuration.ConnectTimeout, + SocketTimeout: configuration.SocketTimeout, + MaxMessageSize: configuration.MaxMessageSize, + }) + } + } else { + if configuration.TLSConfig != nil { + socket = thrift.NewTSSLSocketConf(addr, &thrift.TConfiguration{ + ConnectTimeout: configuration.ConnectTimeout, + SocketTimeout: configuration.SocketTimeout, + MaxMessageSize: configuration.MaxMessageSize, + TLSConfig: configuration.TLSConfig, + }) + } else { + socket = thrift.NewTSocketConf(addr, &thrift.TConfiguration{ + ConnectTimeout: configuration.ConnectTimeout, + SocketTimeout: configuration.SocketTimeout, + MaxMessageSize: configuration.MaxMessageSize, + }) + } + if err = socket.Open(); err != nil { + return + } + } + + if configuration.Username == "" { + _user, err := user.Current() + if err != nil { + return nil, errors.New("Can't determine the username") + } + configuration.Username = strings.Replace(_user.Name, " ", "", -1) + } + // password may not matter but can't be empty + if configuration.Password == "" { + configuration.Password = "x" + } + + if configuration.TransportMode == "http" { + if auth == "NONE" || auth == "LDAP" || auth == "CUSTOM" { + httpClient, protocol, err := prepareHTTPClient(configuration) + if err != nil { + return nil, err + } + + baseTransport := httpClient.Transport + authTransport := &basicAuthRoundTripper{ + Base: baseTransport, + Username: configuration.Username, + Password: configuration.Password, + } + httpClient.Transport = withCookieAuthentication(configuration, baseTransport, authTransport) + httpOptions := thrift.THttpClientOptions{Client: httpClient} + transport, err = thrift.NewTHttpClientTransportFactoryWithOptions(hiveHTTPURL(protocol, host, port, configuration.HTTPPath), httpOptions).GetTransport(socket) + if err != nil { + return nil, err + } + } else if auth == "NOSASL" { + httpClient, protocol, err := prepareHTTPClient(configuration) + if err != nil { + return nil, err + } + httpOptions := thrift.THttpClientOptions{Client: httpClient} + transport, err = thrift.NewTHttpClientTransportFactoryWithOptions(hiveHTTPURL(protocol, host, port, configuration.HTTPPath), httpOptions).GetTransport(socket) + if err != nil { + return nil, err + } + } else if auth == "KERBEROS" { + httpClient, protocol, err := prepareHTTPClient(configuration) + if err != nil { + return nil, err + } + baseTransport := httpClient.Transport + spnegoTransport := gosasl.NewSPNEGORoundTripperWithOptions( + baseTransport, + configuration.Service, + host, + configuration.HTTPKerberosChannelBinding, + configuration.GSSAPIOptions, + ) + httpClient.Transport = withCookieAuthentication(configuration, baseTransport, spnegoTransport) + authCloser = spnegoTransport + httpOptions := thrift.THttpClientOptions{ + Client: httpClient, + } + transport, err = thrift.NewTHttpClientTransportFactoryWithOptions(hiveHTTPURL(protocol, host, port, configuration.HTTPPath), httpOptions).GetTransport(socket) + if err != nil { + return nil, err + } + } else if auth == "JWT" { + if strings.TrimSpace(configuration.JWT) == "" { + return nil, errors.New("Hive JWT authentication requires a token") + } + httpClient, protocol, err := prepareHTTPClient(configuration) + if err != nil { + return nil, err + } + baseTransport := httpClient.Transport + authTransport := &bearerAuthRoundTripper{Base: baseTransport, Token: configuration.JWT} + httpClient.Transport = withCookieAuthentication(configuration, baseTransport, authTransport) + httpOptions := thrift.THttpClientOptions{Client: httpClient} + transport, err = thrift.NewTHttpClientTransportFactoryWithOptions(hiveHTTPURL(protocol, host, port, configuration.HTTPPath), httpOptions).GetTransport(socket) + if err != nil { + return nil, err + } + } else if auth == "BROWSER" { + httpClient, protocol, err := prepareHTTPClient(configuration) + if err != nil { + return nil, err + } + baseTransport := httpClient.Transport + var authTransport http.RoundTripper + if strings.TrimSpace(configuration.BrowserToken) != "" { + authTransport = &bearerAuthRoundTripper{ + Base: baseTransport, + Token: configuration.BrowserToken, + ClientIdentifier: configuration.BrowserClientID, + } + } else { + browserClient = newBrowserSSOClient(configuration.BrowserResponsePort, configuration.BrowserResponseTimeout) + if err := browserClient.Start(); err != nil { + return nil, err + } + defer browserClient.Close() + httpClient.CheckRedirect = func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + } + authTransport = &browserAuthRoundTripper{Base: baseTransport, Client: browserClient} + } + httpClient.Transport = withCookieAuthentication(configuration, baseTransport, authTransport) + httpOptions := thrift.THttpClientOptions{Client: httpClient} + transport, err = thrift.NewTHttpClientTransportFactoryWithOptions(hiveHTTPURL(protocol, host, port, configuration.HTTPPath), httpOptions).GetTransport(socket) + if err != nil { + return nil, err + } + } else if auth == "DIGEST-MD5" { + if strings.TrimSpace(configuration.DelegationToken) == "" { + return nil, errors.New("Hive HTTP delegation token authentication requires a token") + } + httpClient, protocol, err := prepareHTTPClient(configuration) + if err != nil { + return nil, err + } + baseTransport := httpClient.Transport + authTransport := &headerAuthRoundTripper{ + Base: baseTransport, + Name: "X-Hive-Delegation-Token", + Value: configuration.DelegationToken, + } + httpClient.Transport = withCookieAuthentication(configuration, baseTransport, authTransport) + httpOptions := thrift.THttpClientOptions{Client: httpClient} + transport, err = thrift.NewTHttpClientTransportFactoryWithOptions(hiveHTTPURL(protocol, host, port, configuration.HTTPPath), httpOptions).GetTransport(socket) + if err != nil { + return nil, err + } + } else { + return nil, fmt.Errorf("unsupported Hive HTTP authentication %q", auth) + } + } else if configuration.TransportMode == "binary" { + if auth == "NOSASL" { + transport = thrift.NewTBufferedTransport(socket, 4096) + if transport == nil { + return nil, errors.New("BufferedTransport was nil") + } + } else if auth == "NONE" || auth == "LDAP" || auth == "CUSTOM" { + saslConfiguration := map[string]string{"username": configuration.Username, "password": configuration.Password} + transport, err = NewTSaslTransport(socket, host, "PLAIN", saslConfiguration, configuration.MaxSize) + if err != nil { + return + } + } else if auth == "KERBEROS" { + saslConfiguration := map[string]string{"service": configuration.Service} + transport, err = NewTSaslTransportWithGSSAPIOptions( + socket, + host, + "GSSAPI", + saslConfiguration, + configuration.MaxSize, + configuration.GSSAPIOptions, + ) + if err != nil { + return + } + } else if auth == "DIGEST-MD5" { + saslConfiguration := map[string]string{"username": configuration.Username, "password": configuration.Password, "service": configuration.Service} + transport, err = NewTSaslTransport(socket, host, "DIGEST-MD5", saslConfiguration, configuration.MaxSize) + if err != nil { + return + } + } else { + return nil, fmt.Errorf("unsupported Hive binary authentication %q", auth) + } + if !transport.IsOpen() { + if err = transport.Open(); err != nil { + return + } + } + } else { + return nil, fmt.Errorf("unsupported Hive transport mode %q", configuration.TransportMode) + } + + protocolFactory := thrift.NewTBinaryProtocolFactoryConf(&thrift.TConfiguration{MaxMessageSize: configuration.MaxMessageSize}) + client := hiveserver.NewTCLIServiceClientFactory(transport, protocolFactory) + + openSession := hiveserver.NewTOpenSessionReq() + openSession.ClientProtocol = hiveserver.TProtocolVersion_HIVE_CLI_SERVICE_PROTOCOL_V10 + openSession.Configuration = configuration.HiveConfiguration + openSession.Username = &configuration.Username + openSession.Password = &configuration.Password + // Context is ignored + response, err := client.OpenSession(ctx, openSession) + if auth == "BROWSER" && browserClient != nil && err != nil && browserClient.HasRedirect() { + if browserErr := browserClient.Authenticate(ctx); browserErr != nil { + return nil, browserErr + } + response, err = client.OpenSession(ctx, openSession) + } + if err != nil { + return + } + if response == nil { + return nil, errors.New("HiveServer2 OpenSession returned no response") + } + if !success(safeStatus(response.GetStatus())) { + return nil, hiveStatusError("opening session", response.GetStatus()) + } + if response.SessionHandle == nil { + return nil, errors.New("HiveServer2 OpenSession returned no session handle") + } + if configuration.RequestTracker != nil { + configuration.RequestTracker.setSessionHandle(response.SessionHandle) + } + + database := configuration.Database + if database == "" { + database = "default" + } + conn = &connection{ + host: host, + port: port, + database: database, + auth: auth, + kerberosServiceName: "", + sessionHandle: response.SessionHandle, + client: client, + configuration: configuration, + transport: transport, + authCloser: authCloser, + } + + if configuration.Database != "" { + cursor := conn.cursor() + defer cursor.close(ctx) + cursor.exec(ctx, "USE "+quoteHiveIdentifier(configuration.Database)) + if cursor.Err != nil { + return nil, cursor.Err + } + } + + return conn, nil +} + +type cookieDedupTransport struct { + http.RoundTripper +} + +// RoundTrip removes duplicate cookies (cookies with the same name) from the request +// This is a mitigation for the issue where Hive/Impala cookies get duplicated in the response +func (d *cookieDedupTransport) RoundTrip(req *http.Request) (*http.Response, error) { + attempt := req.Clone(req.Context()) + attempt.Header = req.Header.Clone() + cookieMap := map[string]string{} + for _, cookie := range attempt.Cookies() { + cookieMap[cookie.Name] = cookie.Value + } + + attempt.Header.Del("Cookie") + + for key, value := range cookieMap { + attempt.AddCookie(&http.Cookie{Name: key, Value: value}) + } + + base := d.RoundTripper + if base == nil { + base = http.DefaultTransport + } + resp, err := base.RoundTrip(attempt) + + return resp, err +} + +func getHTTPClient(configuration *connectConfiguration) (httpClient *http.Client, protocol string, err error) { + dialContext := configuration.DialContext + if dialContext == nil { + dialContext = (&net.Dialer{Timeout: configuration.ConnectTimeout}).DialContext + } + if configuration.TLSConfig != nil { + httpClient = &http.Client{ + Timeout: configuration.HttpTimeout, + Transport: &http.Transport{ + TLSClientConfig: configuration.TLSConfig, + DialContext: dialContext, + DisableKeepAlives: configuration.DisableKeepAlives, + }, + } + protocol = "https" + } else { + httpClient = &http.Client{ + Timeout: configuration.HttpTimeout, + Transport: &http.Transport{ + DialContext: dialContext, + DisableKeepAlives: configuration.DisableKeepAlives, + }, + } + protocol = "http" + } + + httpClient.Transport = &cookieDedupTransport{httpClient.Transport} + + return +} + +func prepareHTTPClient(configuration *connectConfiguration) (*http.Client, string, error) { + httpClient, protocol, err := getHTTPClient(configuration) + if err != nil { + return nil, "", err + } + if configuration.RequestTracking && configuration.RequestTracker == nil { + configuration.RequestTracker = newHTTPRequestTracker() + } + httpClient.Transport = &customHTTPRoundTripper{ + Base: httpClient.Transport, + Headers: configuration.HTTPHeaders, + Cookies: configuration.HTTPCookies, + RequestTracker: configuration.RequestTracker, + } + if !configuration.DisableCookieAuth { + httpClient.Jar, err = cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List}) + if err != nil { + return nil, "", err + } + } + return httpClient, protocol, nil +} + +// cursor is used for fetching the rows after a query +type cursor struct { + conn *connection + operationHandle *hiveserver.TOperationHandle + fetchSize int64 + queue []*hiveserver.TColumn + response *hiveserver.TFetchResultsResp + columnIndex int + totalRows int + state int + newData bool + Err error + descriptionData [][]string + + // Caller is responsible for managing this channel + Logs chan<- []string +} + +// exec issues a synchronous query. +func (c *cursor) exec(ctx context.Context, query string) { + c.execute(ctx, query) +} + +// execute sends a query to hive for execution with a context +func (c *cursor) execute(ctx context.Context, query string) { + c.executeSync(ctx, query) + // We cannot trust in setting executeReq.RunAsync = true + // because if the context ends the operation can't be cancelled cleanly + if c.Err != nil { + if c.state == _CONTEXT_DONE { + c.handleDoneContext(ctx) + } + return + } + + if c.Err != nil { + if c.state == _CONTEXT_DONE { + c.handleDoneContext(ctx) + } else if c.state == _ERROR { + c.Err = errors.New("Probably the context was over when passed to execute. This probably resulted in the message being sent but we didn't get an operation handle so it's most likely a bug in thrift") + } + return + } + + // Flush logs after execution is finished + if c.Logs != nil { + logs := c.fetchLogs() + if c.error() != nil { + c.state = _ASYNC_ENDED + return + } + c.Logs <- logs + } + + c.state = _ASYNC_ENDED +} + +func (c *cursor) handleDoneContext(ctx context.Context) { + originalError := c.Err + if c.operationHandle != nil { + c.cancel() + if c.Err != nil { + return + } + } + c.resetState(ctx) + c.Err = originalError + c.state = _FINISHED +} + +// executeSync sends a query to hive for execution with a context +func (c *cursor) executeSync(ctx context.Context, query string) { + c.resetState(ctx) + + c.state = _RUNNING + executeReq := hiveserver.NewTExecuteStatementReq() + c.conn.clientMu.Lock() + executeReq.SessionHandle = c.conn.sessionHandle + executeReq.Statement = query + executeReq.RunAsync = false + var responseExecute *hiveserver.TExecuteStatementResp = nil + + responseExecute, c.Err = c.conn.client.ExecuteStatement(ctx, executeReq) + c.conn.clientMu.Unlock() + + if c.Err != nil { + if strings.Contains(c.Err.Error(), "context deadline exceeded") { + c.state = _CONTEXT_DONE + if responseExecute == nil { + c.state = _ERROR + } else if responseExecute != nil { + // We may need this to cancel the operation + c.operationHandle = responseExecute.OperationHandle + } + } + return + } + if responseExecute == nil { + c.Err = errors.New("HiveServer2 ExecuteStatement returned no response") + c.state = _ERROR + return + } + if !success(safeStatus(responseExecute.GetStatus())) { + status := safeStatus(responseExecute.GetStatus()) + c.Err = &Error{ + Err: errors.New("Error while executing query: " + status.String()), + Message: status.GetErrorMessage(), + ErrorCode: int(status.GetErrorCode()), + SQLState: status.GetSqlState(), + } + return + } + + c.operationHandle = responseExecute.OperationHandle + if c.operationHandle == nil { + c.Err = errors.New("HiveServer2 ExecuteStatement returned no operation handle") + c.state = _ERROR + return + } + if !c.operationHandle.HasResultSet { + c.state = _FINISHED + } +} + +// fetchLogs returns all the Hive execution logs for the latest query up to the current point +func (c *cursor) fetchLogs() []string { + logRequest := hiveserver.NewTFetchResultsReq() + logRequest.OperationHandle = c.operationHandle + logRequest.Orientation = hiveserver.TFetchOrientation_FETCH_NEXT + c.conn.clientMu.Lock() + logRequest.MaxRows = c.conn.configuration.FetchSize + // FetchType 1 is "logs" + logRequest.FetchType = 1 + + resp, err := c.conn.client.FetchResults(context.Background(), logRequest) + c.conn.clientMu.Unlock() + if err != nil || resp == nil || resp.Results == nil { + c.Err = err + return nil + } + + // resp contains 1 row, with a column for each line in the log + cols := resp.Results.GetColumns() + var logs []string + + for _, col := range cols { + logs = append(logs, col.StringVal.Values...) + } + + return logs +} + +func success(status *hiveserver.TStatus) bool { + statusCode := status.GetStatusCode() + return statusCode == hiveserver.TStatusCode_SUCCESS_STATUS || statusCode == hiveserver.TStatusCode_SUCCESS_WITH_INFO_STATUS +} + +func (c *cursor) fetchIfEmpty(ctx context.Context) { + c.Err = nil + if c.totalRows == c.columnIndex { + c.queue = nil + if !c.hasMore(ctx) { + // print stack trace + c.Err = errors.New("No more rows are left") + return + } + if c.Err != nil { + return + } + } +} + +// rowMap returns one row as a map. Advances the cursor one +func (c *cursor) rowMap(ctx context.Context) map[string]interface{} { + c.Err = nil + c.fetchIfEmpty(ctx) + if c.Err != nil { + return nil + } + + d := c.description(ctx) + if c.Err != nil || len(d) != len(c.queue) { + return nil + } + m := make(map[string]interface{}, len(c.queue)) + for i := 0; i < len(c.queue); i++ { + columnName := d[i][0] + columnType := d[i][1] + if columnType == "BOOLEAN_TYPE" { + if isNull(c.queue[i].BoolVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].BoolVal.Values[c.columnIndex] + } + } else if columnType == "TINYINT_TYPE" { + if isNull(c.queue[i].ByteVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].ByteVal.Values[c.columnIndex] + } + } else if columnType == "SMALLINT_TYPE" { + if isNull(c.queue[i].I16Val.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].I16Val.Values[c.columnIndex] + } + } else if columnType == "INT_TYPE" { + if isNull(c.queue[i].I32Val.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].I32Val.Values[c.columnIndex] + } + } else if columnType == "BIGINT_TYPE" { + if isNull(c.queue[i].I64Val.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].I64Val.Values[c.columnIndex] + } + } else if columnType == "FLOAT_TYPE" { + if isNull(c.queue[i].DoubleVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].DoubleVal.Values[c.columnIndex] + } + } else if columnType == "DOUBLE_TYPE" { + if isNull(c.queue[i].DoubleVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].DoubleVal.Values[c.columnIndex] + } + } else if columnType == "STRING_TYPE" || columnType == "VARCHAR_TYPE" || columnType == "CHAR_TYPE" { + if isNull(c.queue[i].StringVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].StringVal.Values[c.columnIndex] + } + } else if columnType == "TIMESTAMP_TYPE" { + if isNull(c.queue[i].StringVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].StringVal.Values[c.columnIndex] + } + } else if columnType == "DATE_TYPE" { + if isNull(c.queue[i].StringVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].StringVal.Values[c.columnIndex] + } + } else if columnType == "BINARY_TYPE" { + if isNull(c.queue[i].BinaryVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].BinaryVal.Values[c.columnIndex] + } + } else if columnType == "ARRAY_TYPE" { + if isNull(c.queue[i].StringVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].StringVal.Values[c.columnIndex] + } + } else if columnType == "MAP_TYPE" { + if isNull(c.queue[i].StringVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].StringVal.Values[c.columnIndex] + } + } else if columnType == "STRUCT_TYPE" { + if isNull(c.queue[i].StringVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].StringVal.Values[c.columnIndex] + } + } else if columnType == "UNION_TYPE" { + if isNull(c.queue[i].StringVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].StringVal.Values[c.columnIndex] + } + } else if columnType == "DECIMAL_TYPE" { + if isNull(c.queue[i].StringVal.Nulls, c.columnIndex) { + m[columnName] = nil + } else { + m[columnName] = c.queue[i].StringVal.Values[c.columnIndex] + } + } + } + if len(m) != len(d) { + log.Printf("Some columns have the same name as per the description: %v, this makes it impossible to get the values using the RowMap API, please use the FetchOne API", d) + } + c.columnIndex++ + return m +} + +func (c *cursor) fetchOneDriver(ctx context.Context, dests []driver.Value) { + c.Err = nil + c.fetchIfEmpty(ctx) + if c.Err != nil { + return + } + + if len(c.queue) != len(dests) { + c.Err = errors.Errorf("%d arguments where passed for filling but the number of columns is %d", len(dests), len(c.queue)) + return + } + + for i := 0; i < len(c.queue); i++ { + if c.queue[i].IsSetBinaryVal() { + if isNull(c.queue[i].BinaryVal.Nulls, c.columnIndex) { + dests[i] = nil + } else { + dests[i] = c.queue[i].BinaryVal.Values[c.columnIndex] + } + } else if c.queue[i].IsSetByteVal() { + if isNull(c.queue[i].ByteVal.Nulls, c.columnIndex) { + dests[i] = nil + } else { + dests[i] = int64(c.queue[i].ByteVal.Values[c.columnIndex]) + } + } else if c.queue[i].IsSetI16Val() { + if isNull(c.queue[i].I16Val.Nulls, c.columnIndex) { + dests[i] = nil + } else { + dests[i] = int64(c.queue[i].I16Val.Values[c.columnIndex]) + } + } else if c.queue[i].IsSetI32Val() { + if isNull(c.queue[i].I32Val.Nulls, c.columnIndex) { + dests[i] = nil + } else { + dests[i] = int64(c.queue[i].I32Val.Values[c.columnIndex]) + } + } else if c.queue[i].IsSetI64Val() { + if isNull(c.queue[i].I64Val.Nulls, c.columnIndex) { + dests[i] = nil + } else { + dests[i] = c.queue[i].I64Val.Values[c.columnIndex] + } + } else if c.queue[i].IsSetStringVal() { + if isNull(c.queue[i].StringVal.Nulls, c.columnIndex) { + dests[i] = nil + } else { + dests[i] = c.queue[i].StringVal.Values[c.columnIndex] + } + } else if c.queue[i].IsSetDoubleVal() { + if isNull(c.queue[i].DoubleVal.Nulls, c.columnIndex) { + dests[i] = nil + } else { + dests[i] = c.queue[i].DoubleVal.Values[c.columnIndex] + } + } else if c.queue[i].IsSetBoolVal() { + if isNull(c.queue[i].BoolVal.Nulls, c.columnIndex) { + dests[i] = nil + } else { + dests[i] = c.queue[i].BoolVal.Values[c.columnIndex] + } + } else { + c.Err = errors.Errorf("Empty column %v", c.queue[i]) + return + } + } + c.columnIndex++ +} + +// fetchOne returns one row and advances the cursor one +func (c *cursor) fetchOne(ctx context.Context, dests ...interface{}) { + c.Err = nil + c.fetchIfEmpty(ctx) + if c.Err != nil { + return + } + + if len(c.queue) != len(dests) { + c.Err = errors.Errorf("%d arguments where passed for filling but the number of columns is %d", len(dests), len(c.queue)) + return + } + for i := 0; i < len(c.queue); i++ { + if c.queue[i].IsSetBinaryVal() { + if dests[i] == nil { + dests[i] = c.queue[i].BinaryVal.Values[c.columnIndex] + continue + } + d, ok := dests[i].(*[]byte) + if !ok { + c.Err = errors.Errorf("Unexpected data type %T for value %v (should be %T) index is %v", dests[i], c.queue[i].BinaryVal.Values[c.columnIndex], c.queue[i].BinaryVal.Values[c.columnIndex], i) + return + } + if isNull(c.queue[i].BinaryVal.Nulls, c.columnIndex) { + *d = nil + } else { + *d = c.queue[i].BinaryVal.Values[c.columnIndex] + } + } else if c.queue[i].IsSetByteVal() { + if dests[i] == nil { + dests[i] = c.queue[i].ByteVal.Values[c.columnIndex] + continue + } + d, ok := dests[i].(*int8) + if !ok { + d, ok := dests[i].(**int8) + if !ok { + c.Err = errors.Errorf("Unexpected data type %T for value %v (should be %T) index is %v", dests[i], c.queue[i].ByteVal.Values[c.columnIndex], c.queue[i].ByteVal.Values[c.columnIndex], i) + return + } + + if isNull(c.queue[i].ByteVal.Nulls, c.columnIndex) { + *d = nil + } else { + if *d == nil { + *d = new(int8) + } + **d = c.queue[i].ByteVal.Values[c.columnIndex] + } + } else { + *d = c.queue[i].ByteVal.Values[c.columnIndex] + } + + } else if c.queue[i].IsSetI16Val() { + if dests[i] == nil { + dests[i] = c.queue[i].I16Val.Values[c.columnIndex] + continue + } + d, ok := dests[i].(*int16) + if !ok { + d, ok := dests[i].(**int16) + if !ok { + c.Err = errors.Errorf("Unexpected data type %T for value %v (should be %T) index is %v", dests[i], c.queue[i].I16Val.Values[c.columnIndex], c.queue[i].I16Val.Values[c.columnIndex], i) + return + } + + if isNull(c.queue[i].I16Val.Nulls, c.columnIndex) { + *d = nil + } else { + if *d == nil { + *d = new(int16) + } + **d = c.queue[i].I16Val.Values[c.columnIndex] + } + } else { + *d = c.queue[i].I16Val.Values[c.columnIndex] + } + } else if c.queue[i].IsSetI32Val() { + if dests[i] == nil { + dests[i] = c.queue[i].I32Val.Values[c.columnIndex] + continue + } + d, ok := dests[i].(*int32) + if !ok { + d, ok := dests[i].(**int32) + if !ok { + c.Err = errors.Errorf("Unexpected data type %T for value %v (should be %T) index is %v", dests[i], c.queue[i].I32Val.Values[c.columnIndex], c.queue[i].I32Val.Values[c.columnIndex], i) + return + } + + if isNull(c.queue[i].I32Val.Nulls, c.columnIndex) { + *d = nil + } else { + if *d == nil { + *d = new(int32) + } + **d = c.queue[i].I32Val.Values[c.columnIndex] + } + } else { + *d = c.queue[i].I32Val.Values[c.columnIndex] + } + } else if c.queue[i].IsSetI64Val() { + if dests[i] == nil { + dests[i] = c.queue[i].I64Val.Values[c.columnIndex] + continue + } + d, ok := dests[i].(*int64) + if !ok { + d, ok := dests[i].(**int64) + if !ok { + c.Err = errors.Errorf("Unexpected data type %T for value %v (should be %T) index is %v", dests[i], c.queue[i].I64Val.Values[c.columnIndex], c.queue[i].I64Val.Values[c.columnIndex], i) + return + } + + if isNull(c.queue[i].I64Val.Nulls, c.columnIndex) { + *d = nil + } else { + if *d == nil { + *d = new(int64) + } + **d = c.queue[i].I64Val.Values[c.columnIndex] + } + } else { + *d = c.queue[i].I64Val.Values[c.columnIndex] + } + } else if c.queue[i].IsSetStringVal() { + if dests[i] == nil { + dests[i] = c.queue[i].StringVal.Values[c.columnIndex] + continue + } + d, ok := dests[i].(*string) + if !ok { + d, ok := dests[i].(**string) + if !ok { + c.Err = errors.Errorf("Unexpected data type %T for value %v (should be %T) index is %v", dests[i], c.queue[i].StringVal.Values[c.columnIndex], c.queue[i].StringVal.Values[c.columnIndex], i) + return + } + + if isNull(c.queue[i].StringVal.Nulls, c.columnIndex) { + *d = nil + } else { + if *d == nil { + *d = new(string) + } + **d = c.queue[i].StringVal.Values[c.columnIndex] + } + } else { + *d = c.queue[i].StringVal.Values[c.columnIndex] + } + } else if c.queue[i].IsSetDoubleVal() { + if dests[i] == nil { + dests[i] = c.queue[i].DoubleVal.Values[c.columnIndex] + continue + } + d, ok := dests[i].(*float64) + if !ok { + d, ok := dests[i].(**float64) + if !ok { + c.Err = errors.Errorf("Unexpected data type %T for value %v (should be %T) index is %v", dests[i], c.queue[i].DoubleVal.Values[c.columnIndex], c.queue[i].DoubleVal.Values[c.columnIndex], i) + return + } + + if isNull(c.queue[i].DoubleVal.Nulls, c.columnIndex) { + *d = nil + } else { + if *d == nil { + *d = new(float64) + } + **d = c.queue[i].DoubleVal.Values[c.columnIndex] + } + } else { + *d = c.queue[i].DoubleVal.Values[c.columnIndex] + } + } else if c.queue[i].IsSetBoolVal() { + if dests[i] == nil { + dests[i] = c.queue[i].BoolVal.Values[c.columnIndex] + continue + } + d, ok := dests[i].(*bool) + if !ok { + d, ok := dests[i].(**bool) + if !ok { + c.Err = errors.Errorf("Unexpected data type %T for value %v (should be %T) index is %v", dests[i], c.queue[i].BoolVal.Values[c.columnIndex], c.queue[i].BoolVal.Values[c.columnIndex], i) + return + } + + if isNull(c.queue[i].BoolVal.Nulls, c.columnIndex) { + *d = nil + } else { + if *d == nil { + *d = new(bool) + } + **d = c.queue[i].BoolVal.Values[c.columnIndex] + } + } else { + *d = c.queue[i].BoolVal.Values[c.columnIndex] + } + } else { + c.Err = errors.Errorf("Empty column %v", c.queue[i]) + return + } + } + c.columnIndex++ +} + +func isNull(nulls []byte, position int) bool { + index := position / 8 + if len(nulls) > index { + b := nulls[index] + return (b & (1 << (uint)(position%8))) != 0 + } + return false +} + +// description return a map with the names of the columns and their types +// must be called after a FetchResult request +// a context should be added here but seems to be ignored by thrift +func (c *cursor) description(ctx context.Context) [][]string { + if c.descriptionData != nil { + return c.descriptionData + } + if c.operationHandle == nil { + c.Err = errors.Errorf("Description can only be called after after a Poll or after an async request") + return nil + } + + metaRequest := hiveserver.NewTGetResultSetMetadataReq() + metaRequest.OperationHandle = c.operationHandle + c.conn.clientMu.Lock() + metaResponse, err := c.conn.client.GetResultSetMetadata(ctx, metaRequest) + c.conn.clientMu.Unlock() + if err != nil { + c.Err = err + return nil + } + if metaResponse == nil { + c.Err = errors.New("HiveServer2 GetResultSetMetadata returned no response") + return nil + } + if !success(safeStatus(metaResponse.GetStatus())) { + c.Err = hiveStatusError("retrieving result-set metadata", metaResponse.GetStatus()) + return nil + } + if metaResponse.Schema == nil { + c.Err = errors.New("HiveServer2 GetResultSetMetadata returned no schema") + return nil + } + m := make([][]string, len(metaResponse.Schema.Columns)) + for i, column := range metaResponse.Schema.Columns { + if column == nil || column.TypeDesc == nil { + c.Err = fmt.Errorf("HiveServer2 metadata column %d has no type description", i) + return nil + } + for _, typeDesc := range column.TypeDesc.Types { + if typeDesc == nil || typeDesc.PrimitiveEntry == nil { + continue + } + m[i] = []string{column.ColumnName, typeDesc.PrimitiveEntry.Type.String()} + } + if len(m[i]) == 0 { + c.Err = fmt.Errorf("HiveServer2 metadata column %d has no primitive type", i) + return nil + } + } + c.descriptionData = m + return m +} + +// hasMore returns whether more rows can be fetched from the server +func (c *cursor) hasMore(ctx context.Context) bool { + c.Err = nil + if c.response == nil && c.state != _FINISHED { + c.Err = c.pollUntilData(ctx, 1) + return c.state != _FINISHED || c.totalRows != c.columnIndex + } + // *c.response.HasMoreRows is always false + // so it can be checked and another roundtrip has to be done if extra data has been added + if c.totalRows == c.columnIndex && c.state != _FINISHED { + c.Err = c.pollUntilData(ctx, 1) + } + + return c.state != _FINISHED || c.totalRows != c.columnIndex +} + +func (c *cursor) error() error { + return c.Err +} + +func (c *cursor) effectiveFetchSize() int64 { + if c.fetchSize > 0 { + return c.fetchSize + } + return c.conn.configuration.FetchSize +} + +func (c *cursor) pollUntilData(ctx context.Context, n int) (err error) { + rowsAvailable := make(chan error) + var stopLock sync.Mutex + var done = false + go func() { + defer close(rowsAvailable) + for true { + stopLock.Lock() + if done { + stopLock.Unlock() + rowsAvailable <- nil + return + } + stopLock.Unlock() + + fetchRequest := hiveserver.NewTFetchResultsReq() + fetchRequest.OperationHandle = c.operationHandle + fetchRequest.Orientation = hiveserver.TFetchOrientation_FETCH_NEXT + c.conn.clientMu.Lock() + fetchRequest.MaxRows = c.effectiveFetchSize() + responseFetch, err := c.conn.client.FetchResults(ctx, fetchRequest) + c.conn.clientMu.Unlock() + if err != nil { + rowsAvailable <- err + return + } + if responseFetch == nil { + rowsAvailable <- errors.New("HiveServer2 FetchResults returned no response") + return + } + c.response = responseFetch + + if !success(safeStatus(responseFetch.GetStatus())) { + rowsAvailable <- hiveStatusError("fetching results", responseFetch.GetStatus()) + return + } + err = c.parseResults(responseFetch) + if err != nil { + rowsAvailable <- err + return + } + + if len(c.queue) > 0 { + rowsAvailable <- nil + return + } + time.Sleep(time.Duration(c.conn.configuration.PollIntervalInMillis) * time.Millisecond) + } + }() + + select { + case err = <-rowsAvailable: + case <-ctx.Done(): + stopLock.Lock() + done = true + stopLock.Unlock() + select { + // Wait for goroutine to finish + case <-rowsAvailable: + } + err = errors.New("Context is done") + } + + if err != nil { + return err + } + + if len(c.queue) < n { + return errors.Errorf("Only %d rows where received", len(c.queue)) + } + return nil +} + +// cancel cancels the current operation +func (c *cursor) cancel() { + c.Err = nil + cancelRequest := hiveserver.NewTCancelOperationReq() + cancelRequest.OperationHandle = c.operationHandle + var responseCancel *hiveserver.TCancelOperationResp + // This context is simply ignored + c.conn.clientMu.Lock() + responseCancel, c.Err = c.conn.client.CancelOperation(context.Background(), cancelRequest) + c.conn.clientMu.Unlock() + if c.Err != nil { + return + } + if responseCancel == nil { + c.Err = errors.New("HiveServer2 CancelOperation returned no response") + return + } + if !success(safeStatus(responseCancel.GetStatus())) { + c.Err = hiveStatusError("cancelling operation", responseCancel.GetStatus()) + } +} + +// close closes the cursor +func (c *cursor) close(ctx context.Context) { + c.Err = c.resetState(ctx) +} + +func (c *cursor) resetState(ctx context.Context) error { + c.response = nil + c.Err = nil + c.queue = nil + c.columnIndex = 0 + c.totalRows = 0 + c.state = _NONE + c.descriptionData = nil + c.newData = false + if c.operationHandle != nil { + closeRequest := hiveserver.NewTCloseOperationReq() + closeRequest.OperationHandle = c.operationHandle + + c.conn.clientMu.Lock() + responseClose, err := c.conn.client.CloseOperation(ctx, closeRequest) + c.conn.clientMu.Unlock() + c.operationHandle = nil + if err != nil { + return err + } + if responseClose == nil { + return errors.New("HiveServer2 CloseOperation returned no response") + } + if !success(safeStatus(responseClose.GetStatus())) { + return hiveStatusError("closing operation", responseClose.GetStatus()) + } + return nil + } + return nil +} + +func (c *cursor) parseResults(response *hiveserver.TFetchResultsResp) (err error) { + c.queue = response.Results.GetColumns() + c.columnIndex = 0 + c.totalRows, err = getTotalRows(c.queue) + c.newData = c.totalRows > 0 + if !c.newData { + c.state = _FINISHED + } + return +} + +func getTotalRows(columns []*hiveserver.TColumn) (int, error) { + for _, el := range columns { + if el.IsSetBinaryVal() { + return len(el.BinaryVal.Values), nil + } else if el.IsSetByteVal() { + return len(el.ByteVal.Values), nil + } else if el.IsSetI16Val() { + return len(el.I16Val.Values), nil + } else if el.IsSetI32Val() { + return len(el.I32Val.Values), nil + } else if el.IsSetI64Val() { + return len(el.I64Val.Values), nil + } else if el.IsSetBoolVal() { + return len(el.BoolVal.Values), nil + } else if el.IsSetDoubleVal() { + return len(el.DoubleVal.Values), nil + } else if el.IsSetStringVal() { + return len(el.StringVal.Values), nil + } else { + return -1, errors.Errorf("Unrecognized column type %T", el) + } + } + return 0, errors.New("All columns seem empty") +} + +func safeStatus(status *hiveserver.TStatus) *hiveserver.TStatus { + if status == nil { + return &DEFAULT_STATUS + } + return status +} + +func hiveStatusError(action string, status *hiveserver.TStatus) error { + status = safeStatus(status) + return &Error{ + Err: fmt.Errorf("HiveServer2 error while %s: %s", action, status.String()), + Message: status.GetErrorMessage(), + ErrorCode: int(status.GetErrorCode()), + SQLState: status.GetSqlState(), + } +} + +func quoteHiveIdentifier(value string) string { + return "`" + strings.ReplaceAll(value, "`", "``") + "`" +} + +var DEFAULT_SQL_STATE = "" +var DEFAULT_ERROR_CODE = int32(-1) +var DEFAULT_ERROR_MESSAGE = "unknown error" +var DEFAULT_STATUS = hiveserver.TStatus{ + StatusCode: hiveserver.TStatusCode_ERROR_STATUS, + InfoMessages: nil, + SqlState: &DEFAULT_SQL_STATE, + ErrorCode: &DEFAULT_ERROR_CODE, + ErrorMessage: &DEFAULT_ERROR_MESSAGE, +} + +// cursor creates a cursor from a connection +func (c *connection) cursor() *cursor { + return &cursor{ + conn: c, + } +} + +// close closes a session +func (c *connection) close() error { + c.mu.Lock() + defer c.mu.Unlock() + + var closeErr error + if c.transport != nil { + closeErr = c.transport.Close() + } + if c.authCloser != nil { + if err := c.authCloser.Close(); closeErr == nil { + closeErr = err + } + } + return closeErr +} + +// getTlsConfiguration returns a tls.Config with the provided certificate and key +func getTlsConfiguration(sslPemPath, sslKeyPath string) (tlsConfig *tls.Config, err error) { + cert, err := tls.LoadX509KeyPair(sslPemPath, sslKeyPath) + if err != nil { + return nil, err + } + tlsConfig = &tls.Config{ + Certificates: []tls.Certificate{cert}, + InsecureSkipVerify: true, + } + return tlsConfig, nil +} diff --git a/agents/go-common/gohive/http_auth_test.go b/agents/go-common/gohive/http_auth_test.go new file mode 100644 index 0000000000..d205062cbb --- /dev/null +++ b/agents/go-common/gohive/http_auth_test.go @@ -0,0 +1,375 @@ +package gohive + +import ( + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/beltran/gohive/v2/hiveserver" +) + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return function(request) +} + +func TestBasicAuthRoundTripperAddsCredentialsWithoutMutatingRequest(t *testing.T) { + var authorization string + transport := &basicAuthRoundTripper{ + Base: roundTripFunc(func(request *http.Request) (*http.Response, error) { + authorization = request.Header.Get("Authorization") + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader("ok")), + Request: request, + }, nil + }), + Username: "user@example.com", + Password: "p@ss:word", + } + request, err := http.NewRequest(http.MethodPost, "http://hs2.example.com/cliservice", strings.NewReader("payload")) + if err != nil { + t.Fatal(err) + } + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if authorization != "Basic dXNlckBleGFtcGxlLmNvbTpwQHNzOndvcmQ=" { + t.Fatalf("unexpected Authorization header: %q", authorization) + } + if request.Header.Get("Authorization") != "" { + t.Fatal("original request was mutated") + } +} + +func TestBearerAndDelegationAuthRoundTrippers(t *testing.T) { + for name, transport := range map[string]http.RoundTripper{ + "jwt": &bearerAuthRoundTripper{ + Base: captureRoundTripper(t, "Authorization", "Bearer signed-token"), + Token: "signed-token", + ClientIdentifier: "browser-client", + }, + "delegation": &headerAuthRoundTripper{ + Base: captureRoundTripper(t, "X-Hive-Delegation-Token", "delegation-token"), + Name: "X-Hive-Delegation-Token", + Value: "delegation-token", + }, + } { + t.Run(name, func(t *testing.T) { + request, err := http.NewRequest(http.MethodPost, "http://hs2.example.com/cliservice", strings.NewReader("payload")) + if err != nil { + t.Fatal(err) + } + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if request.Header.Get("Authorization") != "" || request.Header.Get("X-Hive-Delegation-Token") != "" { + t.Fatal("original request was mutated") + } + }) + } +} + +func captureRoundTripper(t *testing.T, header, expected string) http.RoundTripper { + t.Helper() + return roundTripFunc(func(request *http.Request) (*http.Response, error) { + if actual := request.Header.Get(header); actual != expected { + t.Fatalf("%s = %q, expected %q", header, actual, expected) + } + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody, Request: request}, nil + }) +} + +func TestCustomHTTPRoundTripperAddsHeadersAndCookies(t *testing.T) { + transport := &customHTTPRoundTripper{ + Base: roundTripFunc(func(request *http.Request) (*http.Response, error) { + if request.Header.Get("X-Trace-ID") != "trace-value" { + t.Fatalf("custom header missing: %#v", request.Header) + } + if request.Header.Get("X-XSRF-HEADER") != "true" || request.Header.Get("X-CSRF-TOKEN") != "true" { + t.Fatalf("Hive HTTP compatibility headers missing: %#v", request.Header) + } + cookie, err := request.Cookie("SessionID") + if err != nil || cookie.Value != "cookie-value" { + t.Fatalf("custom cookie missing: cookie=%#v err=%v", cookie, err) + } + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody, Request: request}, nil + }), + Headers: map[string]string{"X-Trace-ID": "trace-value"}, + Cookies: map[string]string{"SessionID": "cookie-value"}, + } + request, err := http.NewRequest(http.MethodGet, "http://hs2.example.com", nil) + if err != nil { + t.Fatal(err) + } + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if request.Header.Get("X-Trace-ID") != "" || request.Header.Get("Cookie") != "" { + t.Fatal("original request was mutated") + } +} + +func TestCustomHTTPRoundTripperTracksRequestsAcrossOpenSession(t *testing.T) { + tracker := newHTTPRequestTracker() + var requestIDs []string + transport := &customHTTPRoundTripper{ + Base: roundTripFunc(func(request *http.Request) (*http.Response, error) { + requestIDs = append(requestIDs, request.Header.Get("X-Request-ID")) + return httpResponse(request, http.StatusOK), nil + }), + RequestTracker: tracker, + } + request, err := http.NewRequest(http.MethodPost, "http://hs2.example.com", strings.NewReader("payload")) + if err != nil { + t.Fatal(err) + } + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + tracker.setSessionHandle(&hiveserver.TSessionHandle{SessionId: &hiveserver.THandleIdentifier{GUID: []byte{0x01, 0xab}}}) + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if len(requestIDs) != 2 || requestIDs[0] != "HIVE_NO_SESSION_00000000000000000001" || requestIDs[1] != "HIVE_01ab_00000000000000000002" { + t.Fatalf("unexpected request IDs: %v", requestIDs) + } + if request.Header.Get("X-Request-ID") != "" { + t.Fatal("original request was mutated") + } +} + +func TestCookieAuthRoundTripperUsesAuthWithoutCookie(t *testing.T) { + baseCalls := 0 + authCalls := 0 + transport := &cookieAuthRoundTripper{ + Base: roundTripFunc(func(request *http.Request) (*http.Response, error) { + baseCalls++ + return nil, nil + }), + Auth: roundTripFunc(func(request *http.Request) (*http.Response, error) { + authCalls++ + return httpResponse(request, http.StatusOK), nil + }), + CookieName: "hive.server2.auth", + } + request, err := http.NewRequest(http.MethodGet, "http://hs2.example.com", nil) + if err != nil { + t.Fatal(err) + } + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if baseCalls != 0 || authCalls != 1 { + t.Fatalf("unexpected transport calls: base=%d auth=%d", baseCalls, authCalls) + } +} + +func TestCookieAuthRoundTripperUsesCookieWithoutAuth(t *testing.T) { + baseCalls := 0 + authCalls := 0 + transport := &cookieAuthRoundTripper{ + Base: roundTripFunc(func(request *http.Request) (*http.Response, error) { + baseCalls++ + return httpResponse(request, http.StatusOK), nil + }), + Auth: roundTripFunc(func(request *http.Request) (*http.Response, error) { + authCalls++ + return httpResponse(request, http.StatusOK), nil + }), + CookieName: "hive.server2.auth", + } + request, err := http.NewRequest(http.MethodGet, "http://hs2.example.com", nil) + if err != nil { + t.Fatal(err) + } + request.AddCookie(&http.Cookie{Name: "hive.server2.auth", Value: "session-token"}) + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if baseCalls != 1 || authCalls != 0 { + t.Fatalf("unexpected transport calls: base=%d auth=%d", baseCalls, authCalls) + } +} + +func TestCookieAuthRoundTripperRetriesWithAuthAfterUnauthorized(t *testing.T) { + baseCalls := 0 + authCalls := 0 + transport := &cookieAuthRoundTripper{ + Base: roundTripFunc(func(request *http.Request) (*http.Response, error) { + baseCalls++ + payload, err := io.ReadAll(request.Body) + if err != nil { + t.Fatal(err) + } + if string(payload) != "payload" { + t.Fatalf("unexpected initial payload: %q", payload) + } + return httpResponse(request, http.StatusUnauthorized), nil + }), + Auth: roundTripFunc(func(request *http.Request) (*http.Response, error) { + authCalls++ + payload, err := io.ReadAll(request.Body) + if err != nil { + t.Fatal(err) + } + if string(payload) != "payload" { + t.Fatalf("unexpected retry payload: %q", payload) + } + return httpResponse(request, http.StatusOK), nil + }), + CookieName: "hive.server2.auth", + } + request, err := http.NewRequest(http.MethodPost, "http://hs2.example.com", strings.NewReader("payload")) + if err != nil { + t.Fatal(err) + } + request.Header.Set("X-Original", "preserved") + request.AddCookie(&http.Cookie{Name: "hive.server2.auth", Value: "expired-token"}) + originalHeaders := request.Header.Clone() + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if baseCalls != 1 || authCalls != 1 { + t.Fatalf("unexpected transport calls: base=%d auth=%d", baseCalls, authCalls) + } + if request.Header.Get("Authorization") != "" || request.Header.Get("X-Original") != originalHeaders.Get("X-Original") || request.Header.Get("Cookie") != originalHeaders.Get("Cookie") { + t.Fatalf("original request was mutated: %#v", request.Header) + } +} + +func TestWithCookieAuthenticationDisabledAlwaysUsesAuth(t *testing.T) { + baseCalls := 0 + authCalls := 0 + base := roundTripFunc(func(request *http.Request) (*http.Response, error) { + baseCalls++ + return httpResponse(request, http.StatusOK), nil + }) + auth := roundTripFunc(func(request *http.Request) (*http.Response, error) { + authCalls++ + return httpResponse(request, http.StatusOK), nil + }) + transport := withCookieAuthentication(&connectConfiguration{DisableCookieAuth: true}, base, auth) + request, err := http.NewRequest(http.MethodGet, "http://hs2.example.com", nil) + if err != nil { + t.Fatal(err) + } + request.AddCookie(&http.Cookie{Name: "hive.server2.auth", Value: "session-token"}) + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if baseCalls != 0 || authCalls != 1 { + t.Fatalf("unexpected transport calls: base=%d auth=%d", baseCalls, authCalls) + } +} + +func TestCookieAuthenticationUsesConfiguredAuthenticationCookie(t *testing.T) { + baseCalls := 0 + authCalls := 0 + base := roundTripFunc(func(request *http.Request) (*http.Response, error) { + baseCalls++ + return httpResponse(request, http.StatusOK), nil + }) + auth := roundTripFunc(func(request *http.Request) (*http.Response, error) { + authCalls++ + return httpResponse(request, http.StatusOK), nil + }) + transport := withCookieAuthentication(&connectConfiguration{ + CookieName: "hive.server2.auth", + HTTPCookies: map[string]string{"hive.server2.auth": "session-token"}, + }, base, auth) + request, err := http.NewRequest(http.MethodGet, "http://hs2.example.com", nil) + if err != nil { + t.Fatal(err) + } + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if baseCalls != 1 || authCalls != 0 { + t.Fatalf("unexpected transport calls: base=%d auth=%d", baseCalls, authCalls) + } +} + +func httpResponse(request *http.Request, statusCode int) *http.Response { + return &http.Response{ + StatusCode: statusCode, + Header: make(http.Header), + Body: http.NoBody, + Request: request, + } +} + +func TestHiveHTTPURLSupportsIPv6AndEscapesPath(t *testing.T) { + value := hiveHTTPURL("https", "2001:db8::1", 10001, "/gateway/hive service/") + if value != "https://[2001:db8::1]:10001/gateway/hive%20service" { + t.Fatalf("unexpected Hive HTTP URL: %q", value) + } +} + +func TestGetHTTPClientUsesConfiguredDialTimeout(t *testing.T) { + client, protocol, err := getHTTPClient(&connectConfiguration{ConnectTimeout: 3 * time.Second}) + if err != nil { + t.Fatal(err) + } + if protocol != "http" { + t.Fatalf("unexpected protocol: %q", protocol) + } + deduplicating, ok := client.Transport.(*cookieDedupTransport) + if !ok { + t.Fatalf("unexpected transport: %T", client.Transport) + } + base, ok := deduplicating.RoundTripper.(*http.Transport) + if !ok || base.DialContext == nil { + t.Fatalf("HTTP transport has no dial context: %T", deduplicating.RoundTripper) + } +} + +func TestPrepareHTTPClientHonorsCookieAuth(t *testing.T) { + enabled, _, err := prepareHTTPClient(&connectConfiguration{}) + if err != nil { + t.Fatal(err) + } + if enabled.Jar == nil { + t.Fatal("cookie authentication should be enabled by default") + } + disabled, _, err := prepareHTTPClient(&connectConfiguration{DisableCookieAuth: true}) + if err != nil { + t.Fatal(err) + } + if disabled.Jar != nil { + t.Fatal("cookie authentication was not disabled") + } +} + +func TestCookieDedupTransportDoesNotMutateOriginalRequest(t *testing.T) { + var cookieHeader string + transport := &cookieDedupTransport{RoundTripper: roundTripFunc(func(request *http.Request) (*http.Response, error) { + cookieHeader = request.Header.Get("Cookie") + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody, Request: request}, nil + })} + request, err := http.NewRequest(http.MethodGet, "http://hs2.example.com", nil) + if err != nil { + t.Fatal(err) + } + request.Header.Add("Cookie", "session=first; session=second; route=node-a") + original := request.Header.Get("Cookie") + if _, err := transport.RoundTrip(request); err != nil { + t.Fatal(err) + } + if request.Header.Get("Cookie") != original { + t.Fatalf("original Cookie header was mutated: %q", request.Header.Get("Cookie")) + } + if strings.Count(cookieHeader, "session=") != 1 || !strings.Contains(cookieHeader, "session=second") || !strings.Contains(cookieHeader, "route=node-a") { + t.Fatalf("cookies were not deduplicated: %q", cookieHeader) + } +} + +func TestQuoteHiveIdentifierEscapesBackticks(t *testing.T) { + if value := quoteHiveIdentifier("analytics`prod"); value != "`analytics``prod`" { + t.Fatalf("unexpected quoted identifier: %q", value) + } +} diff --git a/agents/go-common/gohive/metadata.go b/agents/go-common/gohive/metadata.go new file mode 100644 index 0000000000..bc833b0452 --- /dev/null +++ b/agents/go-common/gohive/metadata.go @@ -0,0 +1,161 @@ +package gohive + +import ( + "context" + "database/sql/driver" + "fmt" + "io" + "strings" + + "github.com/beltran/gohive/v2/hiveserver" + "github.com/pkg/errors" +) + +// MetadataResult contains the tabular result returned by a HiveServer2 +// metadata operation. +type MetadataResult struct { + Columns []string + Rows [][]driver.Value +} + +// MetadataProvider exposes HiveServer2 metadata RPCs through database/sql's +// Conn.Raw API. +type MetadataProvider interface { + GetHiveSchemas(context.Context, string) (MetadataResult, error) + GetHiveTables(context.Context, string, string, []string) (MetadataResult, error) + GetHiveColumns(context.Context, string, string, string) (MetadataResult, error) + GetHiveTypeInfo(context.Context) (MetadataResult, error) +} + +var _ MetadataProvider = (*sqlConnection)(nil) + +func (connection *sqlConnection) GetHiveSchemas(ctx context.Context, schemaPattern string) (MetadataResult, error) { + request := hiveserver.NewTGetSchemasReq() + request.SessionHandle = connection.conn.sessionHandle + request.SchemaName = hivePattern(schemaPattern) + connection.conn.clientMu.Lock() + response, err := connection.conn.client.GetSchemas(ctx, request) + connection.conn.clientMu.Unlock() + if err != nil { + return MetadataResult{}, err + } + if response == nil { + return MetadataResult{}, errors.New("HiveServer2 GetSchemas returned no response") + } + return connection.readHiveMetadata(ctx, "GetSchemas", response.GetStatus(), response.GetOperationHandle()) +} + +func (connection *sqlConnection) GetHiveTables( + ctx context.Context, + schemaPattern string, + tablePattern string, + tableTypes []string, +) (MetadataResult, error) { + request := hiveserver.NewTGetTablesReq() + request.SessionHandle = connection.conn.sessionHandle + request.SchemaName = hivePattern(schemaPattern) + request.TableName = hivePattern(tablePattern) + request.TableTypes = append([]string(nil), tableTypes...) + connection.conn.clientMu.Lock() + response, err := connection.conn.client.GetTables(ctx, request) + connection.conn.clientMu.Unlock() + if err != nil { + return MetadataResult{}, err + } + if response == nil { + return MetadataResult{}, errors.New("HiveServer2 GetTables returned no response") + } + return connection.readHiveMetadata(ctx, "GetTables", response.GetStatus(), response.GetOperationHandle()) +} + +func (connection *sqlConnection) GetHiveColumns( + ctx context.Context, + schemaPattern string, + tablePattern string, + columnPattern string, +) (MetadataResult, error) { + request := hiveserver.NewTGetColumnsReq() + request.SessionHandle = connection.conn.sessionHandle + request.SchemaName = hivePattern(schemaPattern) + request.TableName = hivePattern(tablePattern) + request.ColumnName = hivePattern(columnPattern) + connection.conn.clientMu.Lock() + response, err := connection.conn.client.GetColumns(ctx, request) + connection.conn.clientMu.Unlock() + if err != nil { + return MetadataResult{}, err + } + if response == nil { + return MetadataResult{}, errors.New("HiveServer2 GetColumns returned no response") + } + return connection.readHiveMetadata(ctx, "GetColumns", response.GetStatus(), response.GetOperationHandle()) +} + +func (connection *sqlConnection) GetHiveTypeInfo(ctx context.Context) (MetadataResult, error) { + request := hiveserver.NewTGetTypeInfoReq() + request.SessionHandle = connection.conn.sessionHandle + connection.conn.clientMu.Lock() + response, err := connection.conn.client.GetTypeInfo(ctx, request) + connection.conn.clientMu.Unlock() + if err != nil { + return MetadataResult{}, err + } + if response == nil { + return MetadataResult{}, errors.New("HiveServer2 GetTypeInfo returned no response") + } + return connection.readHiveMetadata(ctx, "GetTypeInfo", response.GetStatus(), response.GetOperationHandle()) +} + +func (connection *sqlConnection) readHiveMetadata( + ctx context.Context, + action string, + status *hiveserver.TStatus, + operationHandle *hiveserver.TOperationHandle, +) (MetadataResult, error) { + status = safeStatus(status) + if !success(status) { + return MetadataResult{}, &Error{ + Err: fmt.Errorf("HiveServer2 %s failed: %s", action, status.String()), + Message: status.GetErrorMessage(), + ErrorCode: int(status.GetErrorCode()), + SQLState: status.GetSqlState(), + } + } + if operationHandle == nil { + return MetadataResult{}, fmt.Errorf("HiveServer2 %s returned no operation handle", action) + } + metadataRows := &rows{ + cursor: &cursor{ + conn: connection.conn, + operationHandle: operationHandle, + }, + ctx: ctx, + } + defer metadataRows.Close() + columns := metadataRows.Columns() + if err := metadataRows.cursor.error(); err != nil { + return MetadataResult{}, err + } + result := MetadataResult{Columns: columns, Rows: make([][]driver.Value, 0)} + for { + values := make([]driver.Value, len(columns)) + err := metadataRows.Next(values) + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return MetadataResult{}, err + } + result.Rows = append(result.Rows, values) + } + return result, nil +} + +func hivePattern(value string) *hiveserver.TPatternOrIdentifier { + value = strings.TrimSpace(value) + if value == "" { + value = "%" + } + pattern := hiveserver.TPatternOrIdentifier(value) + return &pattern +} diff --git a/agents/go-common/gohive/sasl_transport.go b/agents/go-common/gohive/sasl_transport.go new file mode 100644 index 0000000000..f89efd2de9 --- /dev/null +++ b/agents/go-common/gohive/sasl_transport.go @@ -0,0 +1,304 @@ +package gohive + +import ( + "bytes" + "context" + "encoding/binary" + "fmt" + "io" + + "github.com/apache/thrift/lib/go/thrift" + "github.com/beltran/gosasl" + "github.com/pkg/errors" +) + +const ( + START = 1 + OK = 2 + BAD = 3 + ERROR = 4 + COMPLETE = 5 +) + +// TSaslTransport is a tranport thrift struct that uses SASL +type TSaslTransport struct { + service string + saslClient *gosasl.Client + tp thrift.TTransport + tpFramed thrift.TFramedTransport + mechanism string + writeBuf bytes.Buffer + readBuf bytes.Buffer + buffer [4]byte + rawFrameSize uint32 //Current remaining size of the frame. if ==0 read next frame header + frameSize int //Current remaining size of the frame. if ==0 read next frame header + maxLength uint32 + principal string + OpeningContext context.Context +} + +// NewTSaslTransport return a TSaslTransport +func NewTSaslTransport(trans thrift.TTransport, host string, mechanismName string, configuration map[string]string, maxLength uint32) (*TSaslTransport, error) { + return newTSaslTransport(trans, host, mechanismName, configuration, maxLength, nil) +} + +// NewTSaslTransportWithGSSAPIOptions creates a SASL transport with +// connection-scoped Kerberos configuration. +func NewTSaslTransportWithGSSAPIOptions( + trans thrift.TTransport, + host string, + mechanismName string, + configuration map[string]string, + maxLength uint32, + options gosasl.GSSAPIOptions, +) (*TSaslTransport, error) { + return newTSaslTransport(trans, host, mechanismName, configuration, maxLength, &options) +} + +func newTSaslTransport( + trans thrift.TTransport, + host string, + mechanismName string, + configuration map[string]string, + maxLength uint32, + gssapiOptions *gosasl.GSSAPIOptions, +) (*TSaslTransport, error) { + var mechanism gosasl.Mechanism + switch mechanismName { + case "PLAIN": + mechanism = gosasl.NewPlainMechanism(configuration["username"], configuration["password"]) + case "GSSAPI": + var err error + if gssapiOptions == nil { + mechanism, err = gosasl.NewGSSAPIMechanism(configuration["service"]) + } else { + mechanism, err = gosasl.NewGSSAPIMechanismWithOptions(configuration["service"], *gssapiOptions) + } + if err != nil { + return nil, err + } + case "DIGEST-MD5": + mechanism = gosasl.NewDigestMD5Mechanism(configuration["service"], configuration["username"], configuration["password"]) + default: + return nil, fmt.Errorf("unsupported SASL mechanism %q", mechanismName) + } + if trans == nil { + return nil, errors.New("SASL transport requires an underlying transport") + } + if maxLength == 0 { + return nil, errors.New("SASL maximum frame length must be positive") + } + client := gosasl.NewSaslClient(host, mechanism) + return &TSaslTransport{ + saslClient: client, + tp: trans, + mechanism: mechanismName, + maxLength: maxLength, + principal: configuration["principal"], + OpeningContext: context.Background(), + }, nil +} + +// IsOpen opens a SASL connection +func (p *TSaslTransport) IsOpen() bool { + return p.tp.IsOpen() && p.saslClient.Complete() +} + +// Open check if a SASL transport connection is opened +func (p *TSaslTransport) Open() (err error) { + if !p.tp.IsOpen() { + err = p.tp.Open() + if err != nil { + return err + } + } + if err = p.sendSaslMsg(p.OpeningContext, START, []byte(p.mechanism)); err != nil { + return err + } + + proccessed, err := p.saslClient.Start() + if err != nil { + return err + } + + if err = p.sendSaslMsg(p.OpeningContext, OK, proccessed); err != nil { + return err + } + + for true { + status, challenge, recvErr := p.recvSaslMsg(p.OpeningContext) + if recvErr != nil { + return recvErr + } + if status == OK { + proccessed, err = p.saslClient.Step(challenge) + if err != nil { + return + } + if err = p.sendSaslMsg(p.OpeningContext, OK, proccessed); err != nil { + return err + } + } else if status == COMPLETE { + if !p.saslClient.Complete() { + return thrift.NewTTransportException(thrift.NOT_OPEN, "The server erroneously indicated that SASL negotiation was complete") + } + break + } else { + return thrift.NewTTransportExceptionFromError(errors.Errorf("Bad SASL negotiation status: %d (%s)", status, challenge)) + } + } + return nil +} + +// Close close a SASL transport connection +func (p *TSaslTransport) Close() (err error) { + p.saslClient.Dispose() + return p.tp.Close() +} + +func (p *TSaslTransport) sendSaslMsg(ctx context.Context, status uint8, body []byte) error { + header := make([]byte, 5) + header[0] = status + length := uint32(len(body)) + binary.BigEndian.PutUint32(header[1:], length) + + _, err := p.tp.Write(append(header[:], body[:]...)) + if err != nil { + return err + } + + err = p.tp.Flush(ctx) + if err != nil { + return err + } + return nil +} + +func (p *TSaslTransport) recvSaslMsg(ctx context.Context) (int8, []byte, error) { + header := make([]byte, 5) + _, err := io.ReadFull(p.tp, header) + if err != nil { + return ERROR, nil, thrift.NewTTransportExceptionFromError(fmt.Errorf("reading SASL negotiation header: %w", err)) + } + + status := int8(header[0]) + length := binary.BigEndian.Uint32(header[1:]) + if length > p.maxLength { + return ERROR, nil, thrift.NewTTransportException( + thrift.UNKNOWN_TRANSPORT_EXCEPTION, + fmt.Sprintf("SASL negotiation frame size %d exceeds maximum %d", length, p.maxLength), + ) + } + + if length > 0 { + payload := make([]byte, length) + _, err = io.ReadFull(p.tp, payload) + if err != nil { + return ERROR, nil, thrift.NewTTransportExceptionFromError(fmt.Errorf("reading SASL negotiation payload: %w", err)) + } + return status, payload, nil + } + return status, nil, nil +} + +func (p *TSaslTransport) Read(buf []byte) (l int, err error) { + if p.rawFrameSize == 0 && p.frameSize == 0 { + p.rawFrameSize, err = p.readFrameHeader() + if err != nil { + return + } + } + + var got int + if p.rawFrameSize > 0 { + rawBuf := make([]byte, p.rawFrameSize) + got, err = io.ReadFull(p.tp, rawBuf) + if err != nil { + return + } + p.rawFrameSize = p.rawFrameSize - uint32(got) + + var unwrappedBuf []byte + unwrappedBuf, err = p.saslClient.Decode(rawBuf) + if err != nil { + return + } + p.frameSize += len(unwrappedBuf) + p.readBuf.Write(unwrappedBuf) + } + + // totalBytes := p.readBuf.Len() + got, err = p.readBuf.Read(buf) + p.frameSize = p.frameSize - got + + /* + if p.readBuf.Len() > 0 { + err = thrift.NewTTransportExceptionFromError(fmt.Errorf("Not enough frame size %d to read %d bytes", p.frameSize, totalBytes)) + return + } + */ + if p.frameSize < 0 { + return 0, thrift.NewTTransportException(thrift.UNKNOWN_TRANSPORT_EXCEPTION, "Negative frame size") + } + return got, thrift.NewTTransportExceptionFromError(err) +} + +func (p *TSaslTransport) readFrameHeader() (uint32, error) { + buf := p.buffer[:4] + if _, err := io.ReadFull(p.tp, buf); err != nil { + return 0, err + } + size := binary.BigEndian.Uint32(buf) + if size > p.maxLength { + return 0, thrift.NewTTransportException(thrift.UNKNOWN_TRANSPORT_EXCEPTION, fmt.Sprintf("Frame size is bigger than allowed, set configuration.MaxLength (%d)", size)) + } + return size, nil +} + +func (p *TSaslTransport) Write(buf []byte) (int, error) { + n, err := p.writeBuf.Write(buf) + return n, thrift.NewTTransportExceptionFromError(err) +} + +// Flush the bytes in the buffer +func (p *TSaslTransport) Flush(ctx context.Context) (err error) { + wrappedBuf, err := p.saslClient.Encode(p.writeBuf.Bytes()) + if err != nil { + return thrift.NewTTransportExceptionFromError(err) + } + + p.writeBuf.Reset() + + size := len(wrappedBuf) + if uint64(size) > uint64(p.maxLength) { + return thrift.NewTTransportException( + thrift.UNKNOWN_TRANSPORT_EXCEPTION, + fmt.Sprintf("Frame size is bigger than allowed, set configuration.MaxLength (%d)", size), + ) + } + buf := p.buffer[:4] + binary.BigEndian.PutUint32(buf, uint32(size)) + _, err = p.tp.Write(buf) + + if err != nil { + return thrift.NewTTransportExceptionFromError(err) + } + + if size > 0 { + if _, err := p.tp.Write(wrappedBuf); err != nil { + return thrift.NewTTransportExceptionFromError(err) + } + } + err = p.tp.Flush(ctx) + return thrift.NewTTransportExceptionFromError(err) +} + +// RemainingBytes return the size of the unwrapped bytes +func (p *TSaslTransport) RemainingBytes() uint64 { + return uint64(p.frameSize) +} + +// SetMaxLength set the maxLength +func (p *TSaslTransport) SetMaxLength(maxLength uint32) { + p.maxLength = maxLength +} diff --git a/agents/go-common/gohive/sasl_transport_test.go b/agents/go-common/gohive/sasl_transport_test.go new file mode 100644 index 0000000000..9505960cdd --- /dev/null +++ b/agents/go-common/gohive/sasl_transport_test.go @@ -0,0 +1,44 @@ +package gohive + +import ( + "context" + "encoding/binary" + "strings" + "testing" + + "github.com/apache/thrift/lib/go/thrift" +) + +func TestNewTSaslTransportRejectsUnsupportedMechanism(t *testing.T) { + transport := thrift.NewTMemoryBufferLen(32) + if _, err := NewTSaslTransport(transport, "hs2.example.com", "UNKNOWN", nil, 1024); err == nil || !strings.Contains(err.Error(), "unsupported SASL mechanism") { + t.Fatalf("expected unsupported mechanism error, got %v", err) + } +} + +func TestNewTSaslTransportRejectsInvalidTransportAndLimit(t *testing.T) { + if _, err := NewTSaslTransport(nil, "hs2.example.com", "PLAIN", map[string]string{}, 1024); err == nil { + t.Fatal("expected missing transport error") + } + if _, err := NewTSaslTransport(thrift.NewTMemoryBufferLen(32), "hs2.example.com", "PLAIN", map[string]string{}, 0); err == nil { + t.Fatal("expected invalid maximum length error") + } +} + +func TestRecvSaslMsgRejectsOversizedNegotiationFrame(t *testing.T) { + underlying := thrift.NewTMemoryBufferLen(16) + header := make([]byte, 5) + header[0] = OK + binary.BigEndian.PutUint32(header[1:], 2048) + if _, err := underlying.Write(header); err != nil { + t.Fatal(err) + } + transport, err := NewTSaslTransport(underlying, "hs2.example.com", "PLAIN", map[string]string{}, 1024) + if err != nil { + t.Fatal(err) + } + _, _, err = transport.recvSaslMsg(context.Background()) + if err == nil || !strings.Contains(err.Error(), "exceeds maximum") { + t.Fatalf("expected oversized negotiation frame error, got %v", err) + } +} diff --git a/agents/go-common/gosasl/LICENSE b/agents/go-common/gosasl/LICENSE new file mode 100644 index 0000000000..b84038a6ae --- /dev/null +++ b/agents/go-common/gosasl/LICENSE @@ -0,0 +1,19 @@ +The MIT License (MIT) + +Copyright (c) 2018-2024 Jaume Marhuenda Beltran + +Permission is hereby granted, free of charge, to any person obtaining a copy of this +software and associated documentation files (the "Software"), to deal in the Software +without restriction, including without limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons +to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or +substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, +INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR +PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE +FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. diff --git a/agents/go-common/gosasl/README.md b/agents/go-common/gosasl/README.md new file mode 100644 index 0000000000..f3e783406e --- /dev/null +++ b/agents/go-common/gosasl/README.md @@ -0,0 +1,65 @@ +# Go SASL library + +[![Build Status](https://app.travis-ci.com/beltran/gosasl.svg?branch=master)](https://app.travis-ci.com/beltran/gosasl) + +gosasl is a library for different SASL mechanisms. Currently GSSAPI, DIGEST-MD5, PLAIN and ANONYMOUS are implemented. +Support for other mechanisms may be added in the future. Only GSSAPI supports a QOP higher than auth. + + +## Installation +Gosasl can be installed with: +``` +go get github.com/beltran/gosasl +``` + +To add kerberos support gosasl requires header files to build against the GSSAPI C library. They can be installed with: +- Ubuntu: `sudo apt-get install libkrb5-dev` +- MacOS: `brew install homebrew/dupes/heimdal --without-x11` +- Debian: `yum install -y krb5-devel` + +Then: +``` +go get -tags kerberos github.com/beltran/gosasl +``` + +## Example Usage + +```go + mechanism, err := NewGSSAPIMechanism("service") + if err != nil { + log.Fatal(err) + } + conn = getConnection("somehost") + client := NewSaslClientWithMechanism("somehost", mechanism) + response, err := client.Start() + if err != nil { + log.Fatal(err) + } + conn.sendResponse(response) + + for true { + status, challenge = conn.getChallenge() + if status == COMPLETE { + break + } else if status == OK { + response = client.Step(challenge) + conn.sendResponse(response) + } else { + log.Fatal("Failed to establish connection") + } + } + if !client.Complete() { + log.Fatal("SASL negotiation did not complete") + } + + // begin normal communication + encoded := conn.fetchData() + decoded := client.Decode(encoded) + response = processData(decoded) + conn.sendData(client.Encode(response)) + + client.Dispose() +``` + + +This library is inspired by [pure-sasl](https://github.com/thobbs/pure-sasl). diff --git a/agents/go-common/gosasl/go.mod b/agents/go-common/gosasl/go.mod new file mode 100644 index 0000000000..80e61279d3 --- /dev/null +++ b/agents/go-common/gosasl/go.mod @@ -0,0 +1,22 @@ +module github.com/beltran/gosasl + +go 1.23.0 + +require ( + github.com/alexbrainman/sspi v0.0.0-20250919150558-7d374ff0d59e + github.com/golang-auth/go-gssapi/v2 v2.0.0 +) + +require ( + 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/golang-auth/go-gssapi/v2 => ../go-gssapi diff --git a/agents/go-common/gosasl/go.sum b/agents/go-common/gosasl/go.sum new file mode 100644 index 0000000000..85fd4e89e1 --- /dev/null +++ b/agents/go-common/gosasl/go.sum @@ -0,0 +1,48 @@ +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/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/go-common/gosasl/gssapi.go b/agents/go-common/gosasl/gssapi.go new file mode 100644 index 0000000000..c1cd4983be --- /dev/null +++ b/agents/go-common/gosasl/gssapi.go @@ -0,0 +1,399 @@ +package gosasl + +import ( + "errors" + "fmt" + "net" + "os" + "strings" +) + +const maxSASLBufferLength = 1<<24 - 1 + +type gssapiBackend interface { + Initiate(serviceName string, channelBinding []byte) error + Continue(token []byte) ([]byte, error) + IsEstablished() bool + InitiatorName() string + SupportsIntegrity() bool + SupportsConfidentiality() bool + Wrap(payload []byte, confidentiality bool) ([]byte, error) + Unwrap(token []byte) ([]byte, error) + Dispose() error +} + +var gssapiBackendFactory = newPlatformGSSAPIBackend +var gssapiBackendFactoryWithOptions = newPlatformGSSAPIBackendWithOptions +var lookupCanonicalHostname = net.LookupCNAME + +// GSSAPIOptions contains connection-scoped Kerberos and SPN settings. +type GSSAPIOptions struct { + ConfigPath string + CCachePath string + KeytabPath string + Principal string + Password string + QOP string + AuthorizationID string + ServerName string + ServiceHost string + UseCCache bool + UseKeytab bool + UseSSPI bool + CanonicalizeHost bool + DisablePAFXFAST bool +} + +// GSSAPIMechanism corresponds to the GSSAPI SASL mechanism described by RFC 4752. +type GSSAPIMechanism struct { + config *MechanismConfig + host string + service string + context gssapiBackend + contextStarted bool + securityLayer bool + qop byte + supportedQop byte + serverMaxLength int + UserSelectQop byte + MaxLength int + options GSSAPIOptions +} + +// NewGSSAPIMechanism returns a GSSAPI mechanism backed by pure Go Kerberos, or +// Windows SSPI when DBX_KRB5_USE_SSPI is enabled. +func NewGSSAPIMechanism(service string) (*GSSAPIMechanism, error) { + context, err := gssapiBackendFactory() + return newGSSAPIMechanism(service, gssapiOptionsFromEnvironment(), context, err) +} + +// NewGSSAPIMechanismWithOptions creates a GSSAPI mechanism without reading +// connection credentials from process-global environment variables. +func NewGSSAPIMechanismWithOptions(service string, options GSSAPIOptions) (*GSSAPIMechanism, error) { + context, err := gssapiBackendFactoryWithOptions(options) + return newGSSAPIMechanism(service, options, context, err) +} + +func newGSSAPIMechanism(service string, options GSSAPIOptions, context gssapiBackend, err error) (*GSSAPIMechanism, error) { + if err != nil { + return nil, err + } + userSelectQop, err := configuredQOPMask(options.QOP) + if err != nil { + _ = context.Dispose() + return nil, err + } + return &GSSAPIMechanism{ + config: newDefaultConfig("GSSAPI"), + service: service, + context: context, + UserSelectQop: userSelectQop, + MaxLength: DEFAULT_MAX_LENGTH, + options: options, + }, nil +} + +func (mechanism *GSSAPIMechanism) start() ([]byte, error) { + if mechanism.contextStarted { + return nil, errors.New("GSSAPI negotiation already started") + } + serviceName := qualifiedServiceNameWithOptions(mechanism.service, mechanism.host, mechanism.options) + if err := mechanism.context.Initiate(serviceName, nil); err != nil { + return nil, fmt.Errorf("initiating GSSAPI context for %q: %w", serviceName, err) + } + mechanism.contextStarted = true + token, err := mechanism.context.Continue(nil) + if err != nil { + return nil, fmt.Errorf("creating initial GSSAPI token: %w", err) + } + return token, nil +} + +func (mechanism *GSSAPIMechanism) step(challenge []byte) ([]byte, error) { + if !mechanism.contextStarted { + return nil, errors.New("GSSAPI negotiation has not started") + } + if mechanism.config.complete { + return nil, errors.New("GSSAPI negotiation is already complete") + } + if !mechanism.context.IsEstablished() { + token, err := mechanism.context.Continue(challenge) + if err != nil { + return nil, fmt.Errorf("continuing GSSAPI context: %w", err) + } + return token, nil + } + if mechanism.securityLayer { + return nil, errors.New("GSSAPI security layer negotiation is already complete") + } + return mechanism.negotiateSecurityLayer(challenge) +} + +func (mechanism *GSSAPIMechanism) negotiateSecurityLayer(challenge []byte) ([]byte, error) { + if !mechanism.context.SupportsIntegrity() { + return nil, errors.New("GSSAPI context does not provide integrity required by SASL negotiation") + } + data, err := mechanism.context.Unwrap(challenge) + if err != nil { + return nil, fmt.Errorf("unwrapping GSSAPI security-layer challenge: %w", err) + } + if len(data) != 4 { + return nil, fmt.Errorf("invalid GSSAPI security-layer challenge length %d, expected 4", len(data)) + } + + mechanism.supportedQop = QOP_TO_FLAG[AUTH] + if mechanism.context.SupportsIntegrity() { + mechanism.supportedQop |= QOP_TO_FLAG[AUTH_INT] + } + if mechanism.context.SupportsConfidentiality() { + mechanism.supportedQop |= QOP_TO_FLAG[AUTH_CONF] + } + mechanism.serverMaxLength = int(data[1])<<16 | int(data[2])<<8 | int(data[3]) + mechanism.qop, err = mechanism.selectQop(data[0]) + if err != nil { + return nil, err + } + + maxLength := 0 + if mechanism.qop != QOP_TO_FLAG[AUTH] { + maxLength = mechanism.MaxLength + if maxLength < 0 { + maxLength = 0 + } + if maxLength > maxSASLBufferLength { + maxLength = maxSASLBufferLength + } + if mechanism.serverMaxLength > 0 && maxLength > mechanism.serverMaxLength { + maxLength = mechanism.serverMaxLength + } + } + header := []byte{ + mechanism.qop, + byte(maxLength >> 16), + byte(maxLength >> 8), + byte(maxLength), + } + authorizationID := mechanism.config.AuthorizationID + if authorizationID == "" { + authorizationID = strings.TrimSpace(mechanism.options.AuthorizationID) + } + if authorizationID == "" { + authorizationID = mechanism.context.InitiatorName() + } + payload := append(header, []byte(authorizationID)...) + response, err := mechanism.context.Wrap(payload, false) + if err != nil { + return nil, fmt.Errorf("wrapping GSSAPI security-layer response: %w", err) + } + mechanism.securityLayer = true + mechanism.config.complete = true + return response, nil +} + +func qualifiedServiceName(service, host string) string { + return qualifiedServiceNameWithOptions(service, host, gssapiOptionsFromEnvironment()) +} + +func qualifiedServiceNameWithOptions(service, host string, options GSSAPIOptions) string { + if configuredHost := strings.TrimSpace(options.ServiceHost); configuredHost != "" { + host = configuredHost + } else if options.CanonicalizeHost { + host = canonicalKerberosHost(host) + } + if configuredService := strings.TrimSpace(options.ServerName); configuredService != "" { + service = configuredService + } + service = replaceSPNHostWildcard(service, host) + if strings.Contains(service, "/") || strings.Contains(service, "@") { + return service + } + return service + "/" + host +} + +func canonicalKerberosHost(host string) string { + host = strings.Trim(strings.TrimSpace(host), "[]") + if host == "" || net.ParseIP(host) != nil { + return host + } + canonical, err := lookupCanonicalHostname(host) + if err != nil || strings.TrimSpace(canonical) == "" { + return host + } + return strings.TrimSuffix(strings.TrimSpace(canonical), ".") +} + +func configuredEnvironmentBool(key string) bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv(key))) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + +func gssapiOptionsFromEnvironment() GSSAPIOptions { + return GSSAPIOptions{ + ConfigPath: os.Getenv("KRB5_CONFIG"), + CCachePath: os.Getenv("KRB5CCNAME"), + KeytabPath: firstConfiguredEnvironment("KRB5_CLIENT_KTNAME", "KRB5_KTNAME"), + Principal: os.Getenv("DBX_KRB5_PRINCIPAL"), + Password: os.Getenv("DBX_KRB5_PASSWORD"), + QOP: os.Getenv("DBX_KRB5_QOP"), + AuthorizationID: os.Getenv("DBX_KRB5_AUTHORIZATION_ID"), + ServerName: os.Getenv("DBX_KRB5_SERVER_NAME"), + ServiceHost: os.Getenv("SERVICE_HOST_QUALIFIED"), + UseCCache: configuredEnvironmentBool("DBX_KRB5_USE_CCACHE"), + UseKeytab: configuredEnvironmentBool("DBX_KRB5_USE_KEYTAB"), + UseSSPI: configuredEnvironmentBool("DBX_KRB5_USE_SSPI"), + CanonicalizeHost: configuredEnvironmentBool("DBX_KRB5_CANONICALIZE_HOST"), + DisablePAFXFAST: configuredEnvironmentBool("DBX_KRB5_DISABLE_PAFXFAST"), + } +} + +func firstConfiguredEnvironment(keys ...string) string { + for _, key := range keys { + if value := strings.TrimSpace(os.Getenv(key)); value != "" { + return value + } + } + return "" +} + +// GSSAPIContextClient exposes the reusable authentication context used by +// non-SASL protocols such as HTTP Negotiate/SPNEGO. +type GSSAPIContextClient struct { + context gssapiBackend + service string + host string + options GSSAPIOptions + started bool +} + +// NewGSSAPIContextClient creates a raw GSSAPI initiator for service/host. +func NewGSSAPIContextClient(service, host string) (*GSSAPIContextClient, error) { + context, err := gssapiBackendFactory() + if err != nil { + return nil, err + } + return &GSSAPIContextClient{context: context, service: service, host: host, options: gssapiOptionsFromEnvironment()}, nil +} + +// NewGSSAPIContextClientWithOptions creates a raw GSSAPI initiator with +// connection-scoped credentials. +func NewGSSAPIContextClientWithOptions(service, host string, options GSSAPIOptions) (*GSSAPIContextClient, error) { + context, err := gssapiBackendFactoryWithOptions(options) + if err != nil { + return nil, err + } + return &GSSAPIContextClient{context: context, service: service, host: host, options: options}, nil +} + +// Start initiates the context and returns the first token. channelBinding is +// optional application data, such as a tls-server-end-point binding. +func (client *GSSAPIContextClient) Start(channelBinding []byte) ([]byte, error) { + if client.started { + return nil, errors.New("GSSAPI context already started") + } + serviceName := qualifiedServiceNameWithOptions(client.service, client.host, client.options) + if err := client.context.Initiate(serviceName, append([]byte(nil), channelBinding...)); err != nil { + return nil, fmt.Errorf("initiating GSSAPI context for %q: %w", serviceName, err) + } + client.started = true + return client.context.Continue(nil) +} + +// Continue consumes a peer token and returns the next initiator token. +func (client *GSSAPIContextClient) Continue(challenge []byte) ([]byte, error) { + if !client.started { + return nil, errors.New("GSSAPI context has not started") + } + return client.context.Continue(challenge) +} + +// Complete reports whether mutual authentication has completed. +func (client *GSSAPIContextClient) Complete() bool { + return client.context.IsEstablished() +} + +// Dispose releases credentials and context state. +func (client *GSSAPIContextClient) Dispose() error { + if client.context == nil { + return nil + } + return client.context.Dispose() +} + +func (mechanism *GSSAPIMechanism) selectQop(serverQop byte) (byte, error) { + availableQops := mechanism.UserSelectQop & mechanism.supportedQop & serverQop + for _, qop := range []byte{QOP_TO_FLAG[AUTH_CONF], QOP_TO_FLAG[AUTH_INT], QOP_TO_FLAG[AUTH]} { + if qop&availableQops != 0 { + return qop, nil + } + } + return 0, fmt.Errorf( + "server GSSAPI QOP mask %#02x does not satisfy requested mask %#02x and client-supported mask %#02x", + serverQop, + mechanism.UserSelectQop, + mechanism.supportedQop, + ) +} + +func configuredQOPMask(value string) (byte, error) { + switch strings.ToLower(strings.TrimSpace(value)) { + case "": + return QOP_TO_FLAG[AUTH] | QOP_TO_FLAG[AUTH_INT] | QOP_TO_FLAG[AUTH_CONF], nil + case AUTH: + return QOP_TO_FLAG[AUTH], nil + case AUTH_INT: + return QOP_TO_FLAG[AUTH_INT], nil + case AUTH_CONF: + return QOP_TO_FLAG[AUTH_CONF], nil + default: + return 0, fmt.Errorf("unsupported Kerberos SASL QOP %q", value) + } +} + +// replaceSPNHostWildcard substitutes _HOST in a service principal name. +func replaceSPNHostWildcard(spn, host string) string { + match := krbSPNHost.FindStringSubmatchIndex(spn) + if match == nil || match[2] == -1 { + return spn + } + return spn[:match[2]] + host + spn[match[3]:] +} + +func (mechanism *GSSAPIMechanism) encode(outgoing []byte) ([]byte, error) { + if !mechanism.config.complete { + return nil, errors.New("GSSAPI negotiation is not complete") + } + switch mechanism.qop { + case QOP_TO_FLAG[AUTH]: + return append([]byte(nil), outgoing...), nil + case QOP_TO_FLAG[AUTH_INT]: + return mechanism.context.Wrap(append([]byte(nil), outgoing...), false) + case QOP_TO_FLAG[AUTH_CONF]: + return mechanism.context.Wrap(append([]byte(nil), outgoing...), true) + default: + return nil, fmt.Errorf("unsupported negotiated GSSAPI QOP %#02x", mechanism.qop) + } +} + +func (mechanism *GSSAPIMechanism) decode(incoming []byte) ([]byte, error) { + if !mechanism.config.complete { + return nil, errors.New("GSSAPI negotiation is not complete") + } + if mechanism.qop == QOP_TO_FLAG[AUTH] { + return append([]byte(nil), incoming...), nil + } + return mechanism.context.Unwrap(append([]byte(nil), incoming...)) +} + +func (mechanism *GSSAPIMechanism) dispose() { + if mechanism.context != nil { + _ = mechanism.context.Dispose() + } +} + +func (mechanism *GSSAPIMechanism) getConfig() *MechanismConfig { + return mechanism.config +} diff --git a/agents/go-common/gosasl/gssapi_backend_nonwindows.go b/agents/go-common/gosasl/gssapi_backend_nonwindows.go new file mode 100644 index 0000000000..31b6292d67 --- /dev/null +++ b/agents/go-common/gosasl/gssapi_backend_nonwindows.go @@ -0,0 +1,11 @@ +//go:build !windows + +package gosasl + +func newPlatformGSSAPIBackend() (gssapiBackend, error) { + return newPureGoGSSAPIBackend(gssapiOptionsFromEnvironment()) +} + +func newPlatformGSSAPIBackendWithOptions(options GSSAPIOptions) (gssapiBackend, error) { + return newPureGoGSSAPIBackend(options) +} diff --git a/agents/go-common/gosasl/gssapi_backend_windows.go b/agents/go-common/gosasl/gssapi_backend_windows.go new file mode 100644 index 0000000000..9a61916b5e --- /dev/null +++ b/agents/go-common/gosasl/gssapi_backend_windows.go @@ -0,0 +1,132 @@ +//go:build windows + +package gosasl + +import ( + "os" + "strings" + + "github.com/alexbrainman/sspi" + "github.com/alexbrainman/sspi/kerberos" +) + +const sspiQOPWrapNoEncrypt uint32 = 0x80000001 + +type sspiGSSAPIBackend struct { + credentials *sspi.Credentials + context *kerberos.ClientContext + established bool + initial []byte + sendSeq uint32 + receiveSeq uint32 + principal string +} + +func newPlatformGSSAPIBackend() (gssapiBackend, error) { + return newPlatformGSSAPIBackendWithOptions(gssapiOptionsFromEnvironment()) +} + +func newPlatformGSSAPIBackendWithOptions(options GSSAPIOptions) (gssapiBackend, error) { + if !options.UseSSPI { + return newPureGoGSSAPIBackend(options) + } + credentials, err := kerberos.AcquireCurrentUserCredentials() + if err != nil { + return nil, err + } + return &sspiGSSAPIBackend{credentials: credentials, principal: strings.TrimSpace(options.Principal)}, nil +} + +func (backend *sspiGSSAPIBackend) Initiate(serviceName string, channelBinding []byte) error { + flags := uint32(sspi.ISC_REQ_CONNECTION | + sspi.ISC_REQ_MUTUAL_AUTH | + sspi.ISC_REQ_REPLAY_DETECT | + sspi.ISC_REQ_SEQUENCE_DETECT | + sspi.ISC_REQ_INTEGRITY | + sspi.ISC_REQ_CONFIDENTIALITY) + context, established, token, err := kerberos.NewClientContextWithChannelBindings( + backend.credentials, + serviceName, + flags, + append([]byte(nil), channelBinding...), + ) + if err != nil { + return err + } + backend.context = context + backend.established = established + backend.initial = token + return nil +} + +func (backend *sspiGSSAPIBackend) Continue(token []byte) ([]byte, error) { + if backend.context == nil { + return nil, os.ErrInvalid + } + if backend.initial != nil { + initial := backend.initial + backend.initial = nil + return initial, nil + } + if backend.established { + return nil, nil + } + established, output, err := backend.context.Update(token) + if err != nil { + return nil, err + } + backend.established = established + return output, nil +} + +func (backend *sspiGSSAPIBackend) IsEstablished() bool { + return backend.established +} + +func (backend *sspiGSSAPIBackend) InitiatorName() string { + return backend.principal +} + +func (backend *sspiGSSAPIBackend) SupportsIntegrity() bool { + return backend.context != nil && backend.context.VerifySelectiveFlags(sspi.ISC_REQ_INTEGRITY) == nil +} + +func (backend *sspiGSSAPIBackend) SupportsConfidentiality() bool { + return backend.context != nil && backend.context.VerifySelectiveFlags(sspi.ISC_REQ_CONFIDENTIALITY) == nil +} + +func (backend *sspiGSSAPIBackend) Wrap(payload []byte, confidentiality bool) ([]byte, error) { + qop := uint32(0) + if !confidentiality { + qop = sspiQOPWrapNoEncrypt + } + token, err := backend.context.EncryptMessage(append([]byte(nil), payload...), qop, backend.sendSeq) + if err == nil { + backend.sendSeq++ + } + return token, err +} + +func (backend *sspiGSSAPIBackend) Unwrap(token []byte) ([]byte, error) { + _, payload, err := backend.context.DecryptMessage(append([]byte(nil), token...), backend.receiveSeq) + if err == nil { + backend.receiveSeq++ + } + return payload, err +} + +func (backend *sspiGSSAPIBackend) Dispose() error { + var firstErr error + if backend.context != nil { + firstErr = backend.context.Release() + backend.context = nil + } + if backend.credentials != nil { + if err := backend.credentials.Release(); firstErr == nil { + firstErr = err + } + backend.credentials = nil + } + backend.initial = nil + return firstErr +} diff --git a/agents/go-common/gosasl/gssapi_purego.go b/agents/go-common/gosasl/gssapi_purego.go new file mode 100644 index 0000000000..2a3142c301 --- /dev/null +++ b/agents/go-common/gosasl/gssapi_purego.go @@ -0,0 +1,89 @@ +package gosasl + +import ( + "errors" + "fmt" + + gssapi "github.com/golang-auth/go-gssapi/v2" + "github.com/golang-auth/go-gssapi/v2/common" + "github.com/golang-auth/go-gssapi/v2/krb5" +) + +type pureGoGSSAPIBackend struct { + mechanism gssapi.Mech +} + +func newPureGoGSSAPIBackend(options GSSAPIOptions) (gssapiBackend, error) { + mechanism := krb5.NewKrb5MechWithOptions(krb5.ClientOptions{ + ConfigPath: options.ConfigPath, + CCachePath: options.CCachePath, + KeytabPath: options.KeytabPath, + Principal: options.Principal, + Password: options.Password, + UseCCache: options.UseCCache, + UseKeytab: options.UseKeytab, + DisablePAFXFAST: options.DisablePAFXFAST, + }) + if mechanism == nil { + return nil, errors.New("pure Go Kerberos GSSAPI mechanism is not registered") + } + return &pureGoGSSAPIBackend{mechanism: mechanism}, nil +} + +func (backend *pureGoGSSAPIBackend) Initiate(serviceName string, channelBinding []byte) error { + flags := gssapi.ContextFlagMutual | + gssapi.ContextFlagReplay | + gssapi.ContextFlagSequence | + gssapi.ContextFlagInteg | + gssapi.ContextFlagConf + var binding *common.ChannelBinding + if len(channelBinding) > 0 { + binding = &common.ChannelBinding{Data: append([]byte(nil), channelBinding...)} + } + return backend.mechanism.Initiate(serviceName, flags, binding) +} + +func (backend *pureGoGSSAPIBackend) Continue(token []byte) ([]byte, error) { + return backend.mechanism.Continue(token) +} + +func (backend *pureGoGSSAPIBackend) IsEstablished() bool { + return backend.mechanism.IsEstablished() +} + +func (backend *pureGoGSSAPIBackend) InitiatorName() string { + if named, ok := backend.mechanism.(interface{ InitiatorName() string }); ok { + return named.InitiatorName() + } + return "" +} + +func (backend *pureGoGSSAPIBackend) SupportsIntegrity() bool { + return backend.mechanism.ContextFlags()&gssapi.ContextFlagInteg != 0 +} + +func (backend *pureGoGSSAPIBackend) SupportsConfidentiality() bool { + return backend.mechanism.ContextFlags()&gssapi.ContextFlagConf != 0 +} + +func (backend *pureGoGSSAPIBackend) Wrap(payload []byte, confidentiality bool) ([]byte, error) { + return backend.mechanism.Wrap(payload, confidentiality) +} + +func (backend *pureGoGSSAPIBackend) Unwrap(token []byte) ([]byte, error) { + payload, _, err := backend.mechanism.Unwrap(token) + return payload, err +} + +func (backend *pureGoGSSAPIBackend) Dispose() error { + disposable, ok := backend.mechanism.(interface{ Dispose() }) + if !ok { + return nil + } + disposable.Dispose() + return nil +} + +func (backend *pureGoGSSAPIBackend) String() string { + return fmt.Sprintf("pure-go-gssapi(%T)", backend.mechanism) +} diff --git a/agents/go-common/gosasl/http_spnego.go b/agents/go-common/gosasl/http_spnego.go new file mode 100644 index 0000000000..4868dc91ea --- /dev/null +++ b/agents/go-common/gosasl/http_spnego.go @@ -0,0 +1,272 @@ +package gosasl + +import ( + "crypto/sha256" + "crypto/sha512" + "crypto/tls" + "crypto/x509" + "encoding/base64" + "errors" + "fmt" + "io" + "net/http" + "strings" + "sync" +) + +const defaultSPNEGOMaxRounds = 5 + +// SPNEGORoundTripper performs HTTP Negotiate authentication using the same +// pure-Go Kerberos and Windows SSPI backends as the SASL mechanism. +type SPNEGORoundTripper struct { + Base http.RoundTripper + Service string + Host string + UseTLSChannelBinding bool + MaxRounds int + Options GSSAPIOptions + useEnvironment bool + + mutex sync.Mutex + client *GSSAPIContextClient +} + +// NewSPNEGORoundTripper creates an HTTP Negotiate transport. +func NewSPNEGORoundTripper(base http.RoundTripper, service, host string, useTLSChannelBinding bool) *SPNEGORoundTripper { + if base == nil { + base = http.DefaultTransport + } + return &SPNEGORoundTripper{ + Base: base, + Service: service, + Host: host, + UseTLSChannelBinding: useTLSChannelBinding, + MaxRounds: defaultSPNEGOMaxRounds, + useEnvironment: true, + } +} + +// NewSPNEGORoundTripperWithOptions creates an HTTP Negotiate transport with +// connection-scoped Kerberos configuration. +func NewSPNEGORoundTripperWithOptions( + base http.RoundTripper, + service string, + host string, + useTLSChannelBinding bool, + options GSSAPIOptions, +) *SPNEGORoundTripper { + if base == nil { + base = http.DefaultTransport + } + return &SPNEGORoundTripper{ + Base: base, + Service: service, + Host: host, + UseTLSChannelBinding: useTLSChannelBinding, + MaxRounds: defaultSPNEGOMaxRounds, + Options: options, + } +} + +// RoundTrip implements http.RoundTripper. It serializes authentication because +// GSSAPI contexts carry sequence state and are not safe for concurrent use. +func (transport *SPNEGORoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + transport.mutex.Lock() + defer transport.mutex.Unlock() + + body, err := replayableRequestBody(request) + if err != nil { + return nil, err + } + maxRounds := transport.MaxRounds + if maxRounds <= 0 { + maxRounds = defaultSPNEGOMaxRounds + } + var authorization string + for round := 0; round < maxRounds; round++ { + attempt, err := cloneRequestForSPNEGO(request, body) + if err != nil { + return nil, err + } + if authorization != "" { + attempt.Header.Set("Authorization", authorization) + } + response, err := transport.Base.RoundTrip(attempt) + if err != nil { + return nil, err + } + + challenge, offered, err := negotiateToken(response.Header.Values("WWW-Authenticate")) + if err != nil { + closeHTTPResponse(response) + return nil, err + } + if response.StatusCode != http.StatusUnauthorized { + if !offered || len(challenge) == 0 { + if transport.client != nil && !transport.client.Complete() { + closeHTTPResponse(response) + return nil, errors.New("HTTP Negotiate server did not return a mutual-authentication token") + } + return response, nil + } + if transport.client == nil { + closeHTTPResponse(response) + return nil, errors.New("HTTP Negotiate server returned a token before authentication started") + } + output, continueErr := transport.client.Continue(challenge) + if continueErr != nil { + closeHTTPResponse(response) + return nil, fmt.Errorf("processing HTTP Negotiate mutual-authentication token: %w", continueErr) + } + if len(output) != 0 || !transport.client.Complete() { + closeHTTPResponse(response) + return nil, errors.New("HTTP Negotiate mutual authentication did not complete in the success response") + } + return response, nil + } + if !offered { + return response, nil + } + + channelBinding := []byte(nil) + if transport.UseTLSChannelBinding { + channelBinding, err = tlsServerEndpointBinding(response.TLS) + if err != nil { + closeHTTPResponse(response) + return nil, err + } + } + closeHTTPResponse(response) + + if transport.client == nil || transport.client.Complete() { + if transport.client != nil { + _ = transport.client.Dispose() + } + if transport.useEnvironment { + transport.client, err = NewGSSAPIContextClient(transport.Service, transport.Host) + } else { + transport.client, err = NewGSSAPIContextClientWithOptions(transport.Service, transport.Host, transport.Options) + } + if err != nil { + return nil, err + } + var token []byte + token, err = transport.client.Start(channelBinding) + if err != nil { + return nil, err + } + if len(challenge) > 0 { + return nil, errors.New("HTTP Negotiate acceptor-first tokens are not supported") + } + authorization = "Negotiate " + base64.StdEncoding.EncodeToString(token) + continue + } + + token, continueErr := transport.client.Continue(challenge) + if continueErr != nil { + return nil, fmt.Errorf("continuing HTTP Negotiate authentication: %w", continueErr) + } + authorization = "Negotiate " + base64.StdEncoding.EncodeToString(token) + } + return nil, fmt.Errorf("HTTP Negotiate authentication exceeded %d rounds", maxRounds) +} + +// Close releases any active GSSAPI context. +func (transport *SPNEGORoundTripper) Close() error { + transport.mutex.Lock() + defer transport.mutex.Unlock() + if transport.client == nil { + return nil + } + err := transport.client.Dispose() + transport.client = nil + return err +} + +func replayableRequestBody(request *http.Request) ([]byte, error) { + if request.Body == nil || request.Body == http.NoBody { + return nil, nil + } + if request.GetBody != nil { + body, err := request.GetBody() + if err != nil { + return nil, fmt.Errorf("reopening HTTP Negotiate request body: %w", err) + } + defer body.Close() + return io.ReadAll(body) + } + body, err := io.ReadAll(request.Body) + if err != nil { + return nil, fmt.Errorf("buffering HTTP Negotiate request body: %w", err) + } + _ = request.Body.Close() + request.Body = io.NopCloser(strings.NewReader(string(body))) + return body, nil +} + +func cloneRequestForSPNEGO(request *http.Request, body []byte) (*http.Request, error) { + attempt := request.Clone(request.Context()) + attempt.Header = request.Header.Clone() + if body == nil { + attempt.Body = nil + attempt.GetBody = nil + return attempt, nil + } + attempt.Body = io.NopCloser(strings.NewReader(string(body))) + attempt.ContentLength = int64(len(body)) + attempt.GetBody = func() (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader(string(body))), nil + } + return attempt, nil +} + +func negotiateToken(headers []string) ([]byte, bool, error) { + for _, header := range headers { + for _, value := range strings.Split(header, ",") { + parts := strings.Fields(strings.TrimSpace(value)) + if len(parts) == 0 || !strings.EqualFold(parts[0], "Negotiate") { + continue + } + if len(parts) == 1 { + return nil, true, nil + } + if len(parts) != 2 { + return nil, true, fmt.Errorf("invalid HTTP Negotiate challenge %q", value) + } + token, err := base64.StdEncoding.DecodeString(parts[1]) + if err != nil { + return nil, true, fmt.Errorf("decoding HTTP Negotiate challenge: %w", err) + } + return token, true, nil + } + } + return nil, false, nil +} + +func closeHTTPResponse(response *http.Response) { + if response == nil || response.Body == nil { + return + } + _, _ = io.Copy(io.Discard, response.Body) + _ = response.Body.Close() +} + +func tlsServerEndpointBinding(state *tls.ConnectionState) ([]byte, error) { + if state == nil || len(state.PeerCertificates) == 0 { + return nil, errors.New("TLS channel binding requested but no peer certificate is available") + } + certificate := state.PeerCertificates[0] + var digest []byte + switch certificate.SignatureAlgorithm { + case x509.SHA384WithRSA, x509.ECDSAWithSHA384: + value := sha512.Sum384(certificate.Raw) + digest = value[:] + case x509.SHA512WithRSA, x509.ECDSAWithSHA512: + value := sha512.Sum512(certificate.Raw) + digest = value[:] + default: + value := sha256.Sum256(certificate.Raw) + digest = value[:] + } + return append([]byte("tls-server-end-point:"), digest...), nil +} diff --git a/agents/go-common/gosasl/http_spnego_test.go b/agents/go-common/gosasl/http_spnego_test.go new file mode 100644 index 0000000000..f892ca2212 --- /dev/null +++ b/agents/go-common/gosasl/http_spnego_test.go @@ -0,0 +1,182 @@ +package gosasl + +import ( + "bytes" + "crypto/sha256" + "crypto/tls" + "crypto/x509" + "encoding/base64" + "io" + "net/http" + "strings" + "testing" +) + +type scriptedRoundTripper struct { + responses []*http.Response + requests []*http.Request + bodies [][]byte +} + +func (transport *scriptedRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + body, err := io.ReadAll(request.Body) + if err != nil { + return nil, err + } + transport.requests = append(transport.requests, request) + transport.bodies = append(transport.bodies, body) + response := transport.responses[len(transport.requests)-1] + response.Request = request + return response, nil +} + +func spnegoResponse(status int, challenge string) *http.Response { + header := make(http.Header) + if challenge != "" { + header.Set("WWW-Authenticate", challenge) + } + return &http.Response{ + StatusCode: status, + Header: header, + Body: io.NopCloser(strings.NewReader("response")), + } +} + +func TestSPNEGORoundTripperPerformsMutualAuthentication(t *testing.T) { + backend := &fakeGSSAPIBackend{ + establishAfter: 2, + continueOutputs: [][]byte{[]byte("initial-token"), nil}, + } + installFakeGSSAPIBackend(t, backend) + base := &scriptedRoundTripper{responses: []*http.Response{ + spnegoResponse(http.StatusUnauthorized, "Basic realm=ignored, Negotiate"), + spnegoResponse(http.StatusOK, "Negotiate "+base64.StdEncoding.EncodeToString([]byte("mutual-token"))), + }} + transport := NewSPNEGORoundTripper(base, "HTTP", "hs2.example.com", false) + request, err := http.NewRequest(http.MethodPost, "http://hs2.example.com/cliservice", bytes.NewBufferString("thrift-request")) + if err != nil { + t.Fatal(err) + } + response, err := transport.RoundTrip(request) + if err != nil { + t.Fatal(err) + } + defer response.Body.Close() + if len(base.requests) != 2 { + t.Fatalf("expected 2 HTTP attempts, got %d", len(base.requests)) + } + if got := base.requests[0].Header.Get("Authorization"); got != "" { + t.Fatalf("first request should await the challenge, got %q", got) + } + wantAuthorization := "Negotiate " + base64.StdEncoding.EncodeToString([]byte("initial-token")) + if got := base.requests[1].Header.Get("Authorization"); got != wantAuthorization { + t.Fatalf("unexpected authorization header %q", got) + } + for index, body := range base.bodies { + if string(body) != "thrift-request" { + t.Fatalf("attempt %d changed the request body to %q", index, body) + } + } + if len(backend.continueInputs) != 2 || string(backend.continueInputs[1]) != "mutual-token" { + t.Fatalf("unexpected GSSAPI tokens %q", backend.continueInputs) + } + if !transport.client.Complete() { + t.Fatal("GSSAPI context should be complete") + } +} + +func TestSPNEGORoundTripperSupportsMultiple401Rounds(t *testing.T) { + backend := &fakeGSSAPIBackend{ + establishAfter: 3, + continueOutputs: [][]byte{[]byte("token-1"), []byte("token-2"), nil}, + } + installFakeGSSAPIBackend(t, backend) + base := &scriptedRoundTripper{responses: []*http.Response{ + spnegoResponse(http.StatusUnauthorized, "Negotiate"), + spnegoResponse(http.StatusUnauthorized, "Negotiate "+base64.StdEncoding.EncodeToString([]byte("challenge-2"))), + spnegoResponse(http.StatusOK, "Negotiate "+base64.StdEncoding.EncodeToString([]byte("challenge-3"))), + }} + transport := NewSPNEGORoundTripper(base, "HTTP", "host", false) + request, _ := http.NewRequest(http.MethodPost, "http://host/cliservice", strings.NewReader("body")) + response, err := transport.RoundTrip(request) + if err != nil { + t.Fatal(err) + } + response.Body.Close() + if len(base.requests) != 3 { + t.Fatalf("expected 3 attempts, got %d", len(base.requests)) + } + if got := base.requests[2].Header.Get("Authorization"); got != "Negotiate "+base64.StdEncoding.EncodeToString([]byte("token-2")) { + t.Fatalf("unexpected final authorization %q", got) + } +} + +func TestSPNEGORoundTripperLeavesNonNegotiate401Untouched(t *testing.T) { + base := &scriptedRoundTripper{responses: []*http.Response{ + spnegoResponse(http.StatusUnauthorized, `Basic realm="hive"`), + }} + transport := NewSPNEGORoundTripper(base, "HTTP", "host", false) + request, _ := http.NewRequest(http.MethodPost, "http://host/cliservice", strings.NewReader("body")) + response, err := transport.RoundTrip(request) + if err != nil { + t.Fatal(err) + } + defer response.Body.Close() + if response.StatusCode != http.StatusUnauthorized || transport.client != nil { + t.Fatalf("unexpected response=%d client=%v", response.StatusCode, transport.client) + } +} + +func TestSPNEGORoundTripperRejectsMissingMutualToken(t *testing.T) { + backend := &fakeGSSAPIBackend{ + establishAfter: 2, + continueOutputs: [][]byte{[]byte("initial-token")}, + } + installFakeGSSAPIBackend(t, backend) + base := &scriptedRoundTripper{responses: []*http.Response{ + spnegoResponse(http.StatusUnauthorized, "Negotiate"), + spnegoResponse(http.StatusOK, ""), + }} + transport := NewSPNEGORoundTripper(base, "HTTP", "host", false) + request, _ := http.NewRequest(http.MethodPost, "http://host/cliservice", strings.NewReader("body")) + if _, err := transport.RoundTrip(request); err == nil { + t.Fatal("expected missing mutual-authentication token error") + } +} + +func TestSPNEGORoundTripperUsesTLSChannelBinding(t *testing.T) { + certificate := &x509.Certificate{Raw: []byte("certificate"), SignatureAlgorithm: x509.SHA1WithRSA} + unauthorized := spnegoResponse(http.StatusUnauthorized, "Negotiate") + unauthorized.TLS = &tls.ConnectionState{PeerCertificates: []*x509.Certificate{certificate}} + backend := &fakeGSSAPIBackend{ + establishAfter: 1, + continueOutputs: [][]byte{[]byte("initial-token")}, + } + installFakeGSSAPIBackend(t, backend) + base := &scriptedRoundTripper{responses: []*http.Response{ + unauthorized, + spnegoResponse(http.StatusOK, ""), + }} + transport := NewSPNEGORoundTripper(base, "HTTP", "host", true) + request, _ := http.NewRequest(http.MethodPost, "https://host/cliservice", strings.NewReader("body")) + response, err := transport.RoundTrip(request) + if err != nil { + t.Fatal(err) + } + response.Body.Close() + digest := sha256.Sum256(certificate.Raw) + want := append([]byte("tls-server-end-point:"), digest[:]...) + if !bytes.Equal(backend.channelBinding, want) { + t.Fatalf("unexpected channel binding %x", backend.channelBinding) + } +} + +func TestNegotiateTokenParsing(t *testing.T) { + token, offered, err := negotiateToken([]string{"Basic realm=x", "Negotiate " + base64.StdEncoding.EncodeToString([]byte("token"))}) + if err != nil || !offered || string(token) != "token" { + t.Fatalf("token=%q offered=%v err=%v", token, offered, err) + } + if _, _, err := negotiateToken([]string{"Negotiate !!!"}); err == nil { + t.Fatal("expected invalid base64 error") + } +} diff --git a/agents/go-common/gosasl/sasl.go b/agents/go-common/gosasl/sasl.go new file mode 100644 index 0000000000..43356a3492 --- /dev/null +++ b/agents/go-common/gosasl/sasl.go @@ -0,0 +1,489 @@ +package gosasl + +import ( + "crypto/hmac" + "crypto/md5" + "crypto/rand" + "encoding/base64" + "encoding/hex" + "errors" + "fmt" + "regexp" + "strconv" + "strings" +) + +var ( + krbSPNHost = regexp.MustCompile(`\A[^/]+/(_HOST)([@/]|\z)`) +) + +// DEFAULT_MAX_LENGTH is the max length that will be requested in the negotiation +// It can be set with gssapiMechanism.MaxLength = 1000 +const DEFAULT_MAX_LENGTH = 16384000 + +// AUTH if the flag used for just basic auth, no confidentiality +var AUTH = "auth" + +// AUTH_INT is the flag for authentication and integrety +var AUTH_INT = "auth-int" + +// AUTH_CONF is the flag for authentication and confidentiality. It +// the most secure option. +var AUTH_CONF = "auth-conf" + +// QOP_TO_FLAG is a dict that translate the string flag name into the actual bit +// It can be used wiht gssapiMechanism.UserSelectQop = QOP_TO_FLAG[AUTH_CONF] | QOP_TO_FLAG[AUTH_INT] +var QOP_TO_FLAG = map[string]byte{ + AUTH: 1, + AUTH_INT: 2, + AUTH_CONF: 4, +} + +// QOP is the byte that holds the QOP flags +type QOP []byte + +// MechanismConfig is the configuration to use for mechanisms +type MechanismConfig struct { + name string + score int + complete bool + hasInitialResponse bool + allowsAnonymous bool + usesPlaintext bool + activeSafe bool + dictionarySafe bool + qop QOP + // It can be set with mechanism.getConfig().AuthorizationID = "authorizationId" + AuthorizationID string +} + +// Mechanism is the common interface for all mechanisms +type Mechanism interface { + start() ([]byte, error) + step(challenge []byte) ([]byte, error) + encode(outgoing []byte) ([]byte, error) + decode(incoming []byte) ([]byte, error) + dispose() + getConfig() *MechanismConfig +} + +// AnonymousMechanism corresponds to NONE/ Anonymous SASL mechanism +type AnonymousMechanism struct { + config *MechanismConfig +} + +// NewAnonymousMechanism returns a new AnonymousMechanism +func NewAnonymousMechanism() *AnonymousMechanism { + return &AnonymousMechanism{ + config: newDefaultConfig("Anonymous"), + } +} + +func (m *AnonymousMechanism) start() ([]byte, error) { + return m.step(nil) +} + +func (m *AnonymousMechanism) step([]byte) ([]byte, error) { + m.config.complete = true + return []byte("Anonymous, None"), nil +} + +func (m *AnonymousMechanism) encode([]byte) ([]byte, error) { + return nil, nil +} + +func (m *AnonymousMechanism) decode([]byte) ([]byte, error) { + return nil, nil +} + +func (m *AnonymousMechanism) dispose() {} + +func (m *AnonymousMechanism) getConfig() *MechanismConfig { + return m.config +} + +// PlainMechanism corresponds to PLAIN SASL mechanism +type PlainMechanism struct { + mechanismConfig *MechanismConfig + identity string + username string + password string +} + +// NewPlainMechanism returns a new PlainMechanism +func NewPlainMechanism(username string, password string) *PlainMechanism { + return &PlainMechanism{ + mechanismConfig: newDefaultConfig("PLAIN"), + username: username, + password: password, + } +} + +func (m *PlainMechanism) start() ([]byte, error) { + return m.step(nil) +} + +func (m *PlainMechanism) step(challenge []byte) ([]byte, error) { + m.mechanismConfig.complete = true + var authID string + + if m.mechanismConfig.AuthorizationID != "" { + authID = m.mechanismConfig.AuthorizationID + } else { + authID = m.identity + } + NULL := "\x00" + return []byte(fmt.Sprintf("%s%s%s%s%s", authID, NULL, m.username, NULL, m.password)), nil +} + +func (m *PlainMechanism) encode(outgoing []byte) ([]byte, error) { + return outgoing, nil +} + +func (m *PlainMechanism) decode(incoming []byte) ([]byte, error) { + return incoming, nil +} + +func (m *PlainMechanism) dispose() { + m.password = "" +} + +func (m *PlainMechanism) getConfig() *MechanismConfig { + return m.mechanismConfig +} + +// CramMD5Mechanism corresponds to PLAIN SASL mechanism +type CramMD5Mechanism struct { + *PlainMechanism +} + +// NewCramMD5Mechanism returns a new PlainMechanism +func NewCramMD5Mechanism(username string, password string) *CramMD5Mechanism { + plain := NewPlainMechanism(username, password) + return &CramMD5Mechanism{ + plain, + } +} + +func (m *CramMD5Mechanism) step(challenge []byte) ([]byte, error) { + if challenge == nil { + return nil, nil + } + m.mechanismConfig.complete = true + hash := hmac.New(md5.New, []byte(m.password)) + // hashed := make([]byte, hash.Size()) + _, err := hash.Write(challenge) + if err != nil { + return nil, err + } + return append([]byte(fmt.Sprintf("%s ", m.username)), hash.Sum(nil)...), nil +} + +// DigestMD5Mechanism corresponds to PLAIN SASL mechanism +type DigestMD5Mechanism struct { + mechanismConfig *MechanismConfig + service string + identity string + username string + password string + host string + nonceCount int + cnonce string + nonce string + keyHash string + auth string +} + +// parseChallenge turns the challenge string into a map +func parseChallenge(challenge []byte) (map[string]string, error) { + value := strings.TrimSpace(string(challenge)) + result := make(map[string]string) + for value != "" { + value = strings.TrimLeft(value, " \t,") + if value == "" { + break + } + equals := strings.IndexByte(value, '=') + comma := strings.IndexByte(value, ',') + if equals < 0 || comma >= 0 && comma < equals { + if comma < 0 { + return nil, fmt.Errorf("invalid DIGEST-MD5 challenge directive %q", value) + } + value = value[comma+1:] + continue + } + key := strings.ToLower(strings.TrimSpace(value[:equals])) + if key == "" { + return nil, errors.New("DIGEST-MD5 challenge contains an empty directive name") + } + value = strings.TrimLeft(value[equals+1:], " \t") + var parsed string + if strings.HasPrefix(value, "\"") { + value = value[1:] + var builder strings.Builder + closed := false + for index := 0; index < len(value); index++ { + switch value[index] { + case '\\': + index++ + if index >= len(value) { + return nil, errors.New("DIGEST-MD5 challenge ends with an incomplete escape") + } + builder.WriteByte(value[index]) + case '"': + parsed = builder.String() + value = strings.TrimLeft(value[index+1:], " \t") + if value != "" { + if value[0] != ',' { + return nil, fmt.Errorf("invalid data after DIGEST-MD5 directive %q", key) + } + value = value[1:] + } + closed = true + index = len(value) + default: + builder.WriteByte(value[index]) + } + } + if !closed { + return nil, fmt.Errorf("unterminated DIGEST-MD5 directive %q", key) + } + } else { + if separator := strings.IndexByte(value, ','); separator >= 0 { + parsed = strings.TrimSpace(value[:separator]) + value = value[separator+1:] + } else { + parsed = strings.TrimSpace(value) + value = "" + } + } + result[key] = parsed + } + return result, nil +} + +// NewDigestMD5Mechanism returns a new PlainMechanism +func NewDigestMD5Mechanism(service string, username string, password string) *DigestMD5Mechanism { + return &DigestMD5Mechanism{ + mechanismConfig: newDefaultConfig("DIGEST-MD5"), + service: service, + username: username, + password: password, + } +} + +func (m *DigestMD5Mechanism) start() ([]byte, error) { + return m.step(nil) +} + +func randSeq(n int) (string, error) { + if n <= 0 { + return "", nil + } + random := make([]byte, n) + if _, err := rand.Read(random); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(random)[:n], nil +} + +func (m *DigestMD5Mechanism) authenticate(digestUri string, challengeMap map[string]string) error { + a2String := ":" + digestUri + + if m.auth != "auth" { + a2String += ":00000000000000000000000000000000" + } + + if m.getHash(digestUri, a2String, challengeMap) != challengeMap["rspauth"] { + return fmt.Errorf("authenticate failed") + } + return nil +} + +func (m *DigestMD5Mechanism) getHash(digestUri string, a2String string, c map[string]string) string { + // Create a1: HEX(H(H(username:realm:password):nonce:cnonce:authid)) + if m.keyHash == "" { + x := m.username + ":" + c["realm"] + ":" + m.password + byteKeyHash := md5.Sum([]byte(x)) + m.keyHash = string(byteKeyHash[:]) + } + a1String := []string{ + m.keyHash, + m.nonce, + m.cnonce, + } + + if len(m.mechanismConfig.AuthorizationID) != 0 { + a1String = append(a1String, m.mechanismConfig.AuthorizationID) + } + + h1 := md5.Sum([]byte(strings.Join(a1String, ":"))) + a1 := hex.EncodeToString(h1[:]) + + h2 := md5.Sum([]byte(a2String)) + a2 := hex.EncodeToString(h2[:]) + + // Set nonce count nc + nc := fmt.Sprintf("%08x", m.nonceCount) + + // Create response: H(a1:nonce:nc:cnonce:qop:a2) + r := strings.ToLower(a1) + ":" + m.nonce + ":" + nc + ":" + m.cnonce + ":" + m.auth + ":" + strings.ToLower(a2) + hr := md5.Sum([]byte(r)) + + // Convert response to hex + response := strings.ToLower(hex.EncodeToString(hr[:])) + return string(response) + +} + +func (m *DigestMD5Mechanism) step(challenge []byte) ([]byte, error) { + if challenge == nil { + return nil, nil + } + + // Create map of challenge + c, err := parseChallenge(challenge) + if err != nil { + return nil, err + } + digestUri := m.service + "/" + m.host + + if _, ok := c["rspauth"]; ok { + m.mechanismConfig.complete = true + return nil, m.authenticate(digestUri, c) + } + + // Prepare response variables + m.nonce = c["nonce"] + if m.nonce == "" { + return nil, errors.New("DIGEST-MD5 challenge contains no nonce") + } + m.auth, err = selectDigestQOP(c["qop"]) + if err != nil { + return nil, err + } + if m.nonceCount == 0 { + m.cnonce, err = randSeq(14) + if err != nil { + return nil, fmt.Errorf("generating DIGEST-MD5 client nonce: %w", err) + } + } + m.nonceCount++ + + // Create a2: HEX(H(AUTHENTICATE:digest-uri-value:00000000000000000000000000000000)) + a2String := "AUTHENTICATE:" + digestUri + + maxBuf := "" + if c["qop"] != AUTH { + a2String += ":00000000000000000000000000000000" + maxBuf = ",maxbuf=16777215" + } + // Set nonce count nc + nc := fmt.Sprintf("%08x", m.nonceCount) + // 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 + + return []byte(resp), nil +} + +func selectDigestQOP(value string) (string, error) { + if strings.TrimSpace(value) == "" { + return AUTH, nil + } + for _, candidate := range strings.Split(value, ",") { + if strings.EqualFold(strings.TrimSpace(candidate), AUTH) { + return AUTH, nil + } + } + return "", fmt.Errorf("DIGEST-MD5 server does not offer supported qop=auth: %q", value) +} + +func (m *DigestMD5Mechanism) encode(outgoing []byte) ([]byte, error) { + return outgoing, nil +} + +func (m *DigestMD5Mechanism) decode(incoming []byte) ([]byte, error) { + return incoming, nil +} + +func (m *DigestMD5Mechanism) dispose() { + m.password = "" +} + +func (m *DigestMD5Mechanism) getConfig() *MechanismConfig { + return m.mechanismConfig +} + +// Client is the entry point for usage of this library +type Client struct { + host string + authorizationID string + mechanism Mechanism +} + +func newDefaultConfig(name string) *MechanismConfig { + return &MechanismConfig{ + name: name, + score: 0, + complete: false, + hasInitialResponse: false, + allowsAnonymous: true, + usesPlaintext: true, + activeSafe: false, + dictionarySafe: false, + qop: nil, + AuthorizationID: "", + } +} + +// NewSaslClient creates a new client given a host and a mechanism +func NewSaslClient(host string, mechanism Mechanism) *Client { + mech, ok := mechanism.(*GSSAPIMechanism) + if ok { + mech.host = host + } + mechDigest, ok := mechanism.(*DigestMD5Mechanism) + if ok { + mechDigest.host = host + } + return &Client{ + host: host, + mechanism: mechanism, + } +} + +// Start initializes the client and may generate the first challenge +func (client *Client) Start() ([]byte, error) { + return client.mechanism.start() +} + +// Step is used for the initial handshake +func (client *Client) Step(challenge []byte) ([]byte, error) { + return client.mechanism.step(challenge) +} + +// Complete returns true if the handshake has ended +func (client *Client) Complete() bool { + return client.mechanism.getConfig().complete +} + +// GetConfig returns the configuration of the mechanism +func (client *Client) GetConfig() *MechanismConfig { + return client.mechanism.getConfig() +} + +// Encode is applied on the outgoing bytes to secure them usually +func (client *Client) Encode(outgoing []byte) ([]byte, error) { + return client.mechanism.encode(outgoing) +} + +// Decode is used on the incoming data to produce the usable bytes +func (client *Client) Decode(incoming []byte) ([]byte, error) { + return client.mechanism.decode(incoming) +} + +// Dispose eliminates sensitive information +func (client *Client) Dispose() { + client.mechanism.dispose() +} diff --git a/agents/go-common/gosasl/sasl_gssapi_test.go b/agents/go-common/gosasl/sasl_gssapi_test.go new file mode 100644 index 0000000000..621fed5be9 --- /dev/null +++ b/agents/go-common/gosasl/sasl_gssapi_test.go @@ -0,0 +1,523 @@ +package gosasl + +import ( + "bytes" + "errors" + "fmt" + "testing" +) + +type fakeGSSAPIBackend struct { + establishAfter int + continueCalls int + established bool + integrity bool + confidentiality bool + initiatorName string + initiatedService string + channelBinding []byte + continueInputs [][]byte + continueOutputs [][]byte + securityLayerChallenge []byte + wrapInputs [][]byte + wrapConfidentiality []bool + unwrapInputs [][]byte + disposed bool +} + +func (backend *fakeGSSAPIBackend) Initiate(serviceName string, channelBinding []byte) error { + backend.initiatedService = serviceName + backend.channelBinding = append([]byte(nil), channelBinding...) + return nil +} + +func (backend *fakeGSSAPIBackend) Continue(token []byte) ([]byte, error) { + backend.continueInputs = append(backend.continueInputs, append([]byte(nil), token...)) + backend.continueCalls++ + if backend.continueCalls >= backend.establishAfter { + backend.established = true + } + if backend.continueCalls <= len(backend.continueOutputs) { + return append([]byte(nil), backend.continueOutputs[backend.continueCalls-1]...), nil + } + return nil, nil +} + +func (backend *fakeGSSAPIBackend) IsEstablished() bool { + return backend.established +} + +func (backend *fakeGSSAPIBackend) InitiatorName() string { + return backend.initiatorName +} + +func (backend *fakeGSSAPIBackend) SupportsIntegrity() bool { + return backend.integrity +} + +func (backend *fakeGSSAPIBackend) SupportsConfidentiality() bool { + return backend.confidentiality +} + +func (backend *fakeGSSAPIBackend) Wrap(payload []byte, confidentiality bool) ([]byte, error) { + backend.wrapInputs = append(backend.wrapInputs, append([]byte(nil), payload...)) + backend.wrapConfidentiality = append(backend.wrapConfidentiality, confidentiality) + return append([]byte("wrapped:"), payload...), nil +} + +func (backend *fakeGSSAPIBackend) Unwrap(token []byte) ([]byte, error) { + backend.unwrapInputs = append(backend.unwrapInputs, append([]byte(nil), token...)) + if bytes.Equal(token, []byte("security-layer")) { + return append([]byte(nil), backend.securityLayerChallenge...), nil + } + return append([]byte("unwrapped:"), token...), nil +} + +func (backend *fakeGSSAPIBackend) Dispose() error { + backend.disposed = true + return nil +} + +func installFakeGSSAPIBackend(t *testing.T, backend *fakeGSSAPIBackend) { + t.Helper() + previous := gssapiBackendFactory + gssapiBackendFactory = func() (gssapiBackend, error) { + return backend, nil + } + t.Cleanup(func() { + gssapiBackendFactory = previous + }) +} + +func installFakeGSSAPIBackendWithOptions(t *testing.T, factory func(GSSAPIOptions) (gssapiBackend, error)) { + t.Helper() + previous := gssapiBackendFactoryWithOptions + gssapiBackendFactoryWithOptions = factory + t.Cleanup(func() { + gssapiBackendFactoryWithOptions = previous + }) +} + +func newNegotiatedGSSAPIClient(t *testing.T, qop string, serverQOP byte, backend *fakeGSSAPIBackend) (*Client, *GSSAPIMechanism) { + t.Helper() + t.Setenv("DBX_KRB5_QOP", qop) + installFakeGSSAPIBackend(t, backend) + backend.establishAfter = 1 + backend.integrity = true + backend.securityLayerChallenge = []byte{serverQOP, 0, 0x10, 0} + mechanism, err := NewGSSAPIMechanism("hive") + if err != nil { + t.Fatal(err) + } + client := NewSaslClient("hs2.example.com", mechanism) + if _, err := client.Start(); err != nil { + t.Fatal(err) + } + if _, err := client.Step([]byte("security-layer")); err != nil { + t.Fatal(err) + } + return client, mechanism +} + +func TestGSSAPIMechanismNegotiatesMultipleContextTokens(t *testing.T) { + t.Setenv("DBX_KRB5_QOP", "auth-int") + t.Setenv("DBX_KRB5_AUTHORIZATION_ID", "proxy-user") + backend := &fakeGSSAPIBackend{ + establishAfter: 2, + integrity: true, + confidentiality: true, + initiatorName: "client@EXAMPLE.COM", + continueOutputs: [][]byte{[]byte("ap-req"), []byte("ap-rep-response")}, + securityLayerChallenge: []byte{QOP_TO_FLAG[AUTH] | QOP_TO_FLAG[AUTH_INT] | QOP_TO_FLAG[AUTH_CONF], 0, 0x20, 0}, + } + installFakeGSSAPIBackend(t, backend) + + mechanism, err := NewGSSAPIMechanism("hive/_HOST@EXAMPLE.COM") + if err != nil { + t.Fatal(err) + } + mechanism.MaxLength = 0x1000 + client := NewSaslClient("hs2.example.com", mechanism) + + initial, err := client.Start() + if err != nil { + t.Fatal(err) + } + if string(initial) != "ap-req" { + t.Fatalf("unexpected initial token %q", initial) + } + if backend.initiatedService != "hive/hs2.example.com@EXAMPLE.COM" { + t.Fatalf("unexpected service name %q", backend.initiatedService) + } + mutual, err := client.Step([]byte("ap-rep")) + if err != nil { + t.Fatal(err) + } + if string(mutual) != "ap-rep-response" || client.Complete() { + t.Fatalf("unexpected mutual-auth state token=%q complete=%v", mutual, client.Complete()) + } + response, err := client.Step([]byte("security-layer")) + if err != nil { + t.Fatal(err) + } + if !client.Complete() { + t.Fatal("client should be complete after security-layer negotiation") + } + if string(response) != "wrapped:\x02\x00\x10\x00proxy-user" { + t.Fatalf("unexpected security-layer response %q", response) + } + if len(backend.wrapConfidentiality) != 1 || backend.wrapConfidentiality[0] { + t.Fatalf("security-layer response must use integrity without confidentiality: %v", backend.wrapConfidentiality) + } +} + +func TestGSSAPIMechanismQOPSelectionAndDataProtection(t *testing.T) { + testCases := []struct { + qop string + serverQOP byte + confidentiality bool + expectedHeader []byte + passThrough bool + }{ + {qop: AUTH, serverQOP: 7, expectedHeader: []byte{1, 0, 0, 0}, passThrough: true}, + {qop: AUTH_INT, serverQOP: 7, expectedHeader: []byte{2, 0, 0x10, 0}}, + {qop: AUTH_CONF, serverQOP: 7, confidentiality: true, expectedHeader: []byte{4, 0, 0x10, 0}}, + } + for _, testCase := range testCases { + t.Run(testCase.qop, func(t *testing.T) { + backend := &fakeGSSAPIBackend{confidentiality: true, initiatorName: "client@EXAMPLE.COM"} + client, _ := newNegotiatedGSSAPIClient(t, testCase.qop, testCase.serverQOP, backend) + if got := backend.wrapInputs[0]; !bytes.Equal(got, append(testCase.expectedHeader, []byte("client@EXAMPLE.COM")...)) { + t.Fatalf("unexpected negotiation payload %v", got) + } + + input := []byte("payload") + encoded, err := client.Encode(input) + if err != nil { + t.Fatal(err) + } + if testCase.passThrough { + if !bytes.Equal(encoded, input) { + t.Fatalf("auth QOP changed payload: %q", encoded) + } + input[0] = 'X' + if string(encoded) != "payload" { + t.Fatal("auth QOP returned an aliased buffer") + } + } else { + if string(encoded) != "wrapped:payload" { + t.Fatalf("unexpected encoded payload %q", encoded) + } + if got := backend.wrapConfidentiality[len(backend.wrapConfidentiality)-1]; got != testCase.confidentiality { + t.Fatalf("unexpected confidentiality flag %v", got) + } + } + + decoded, err := client.Decode([]byte("reply")) + if err != nil { + t.Fatal(err) + } + expectedDecoded := "unwrapped:reply" + if testCase.passThrough { + expectedDecoded = "reply" + } + if string(decoded) != expectedDecoded { + t.Fatalf("unexpected decoded payload %q", decoded) + } + }) + } +} + +func TestGSSAPIMechanismPrefersStrongestAvailableQOP(t *testing.T) { + backend := &fakeGSSAPIBackend{confidentiality: true} + _, mechanism := newNegotiatedGSSAPIClient(t, "", QOP_TO_FLAG[AUTH]|QOP_TO_FLAG[AUTH_INT]|QOP_TO_FLAG[AUTH_CONF], backend) + if mechanism.qop != QOP_TO_FLAG[AUTH_CONF] { + t.Fatalf("expected auth-conf, got %#02x", mechanism.qop) + } +} + +func TestGSSAPIMechanismRejectsUnavailableRequestedQOP(t *testing.T) { + t.Setenv("DBX_KRB5_QOP", AUTH_CONF) + backend := &fakeGSSAPIBackend{ + establishAfter: 1, + integrity: true, + confidentiality: true, + securityLayerChallenge: []byte{QOP_TO_FLAG[AUTH] | QOP_TO_FLAG[AUTH_INT], 0, 0x10, 0}, + } + installFakeGSSAPIBackend(t, backend) + mechanism, err := NewGSSAPIMechanism("hive") + if err != nil { + t.Fatal(err) + } + client := NewSaslClient("host", mechanism) + if _, err := client.Start(); err != nil { + t.Fatal(err) + } + if _, err := client.Step([]byte("security-layer")); err == nil { + t.Fatal("expected unsupported QOP error") + } + if client.Complete() { + t.Fatal("client must not complete after unsupported QOP") + } +} + +func TestGSSAPIMechanismRejectsInvalidSecurityLayer(t *testing.T) { + testCases := []struct { + name string + integrity bool + challenge []byte + }{ + {name: "missing integrity", challenge: []byte{1, 0, 0, 0}}, + {name: "short challenge", integrity: true, challenge: []byte{1, 0, 0}}, + } + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + t.Setenv("DBX_KRB5_QOP", AUTH) + backend := &fakeGSSAPIBackend{ + establishAfter: 1, + integrity: testCase.integrity, + securityLayerChallenge: testCase.challenge, + } + installFakeGSSAPIBackend(t, backend) + mechanism, err := NewGSSAPIMechanism("hive") + if err != nil { + t.Fatal(err) + } + client := NewSaslClient("host", mechanism) + if _, err := client.Start(); err != nil { + t.Fatal(err) + } + if _, err := client.Step([]byte("security-layer")); err == nil { + t.Fatal("expected invalid security-layer error") + } + }) + } +} + +func TestGSSAPIMechanismConfigurationAndDisposal(t *testing.T) { + t.Setenv("DBX_KRB5_QOP", "invalid") + backend := &fakeGSSAPIBackend{} + installFakeGSSAPIBackend(t, backend) + if _, err := NewGSSAPIMechanism("hive"); err == nil { + t.Fatal("expected invalid QOP configuration error") + } + if !backend.disposed { + t.Fatal("backend should be disposed after constructor failure") + } + + t.Setenv("DBX_KRB5_QOP", AUTH) + backend = &fakeGSSAPIBackend{establishAfter: 1, integrity: true, securityLayerChallenge: []byte{1, 0, 0, 0}} + gssapiBackendFactory = func() (gssapiBackend, error) { return backend, nil } + mechanism, err := NewGSSAPIMechanism("hive") + if err != nil { + t.Fatal(err) + } + client := NewSaslClient("host", mechanism) + client.Dispose() + if !backend.disposed { + t.Fatal("client disposal did not dispose backend") + } +} + +func TestGSSAPIMechanismBackendFactoryError(t *testing.T) { + previous := gssapiBackendFactory + gssapiBackendFactory = func() (gssapiBackend, error) { + return nil, errors.New("credentials unavailable") + } + t.Cleanup(func() { gssapiBackendFactory = previous }) + if _, err := NewGSSAPIMechanism("hive"); err == nil || err.Error() != "credentials unavailable" { + t.Fatalf("unexpected error %v", err) + } +} + +func TestGSSAPIMechanismWithOptionsUsesConnectionScopedSettings(t *testing.T) { + t.Setenv("DBX_KRB5_PRINCIPAL", "environment@EXAMPLE.COM") + t.Setenv("DBX_KRB5_QOP", AUTH) + t.Setenv("DBX_KRB5_AUTHORIZATION_ID", "environment-user") + t.Setenv("SERVICE_HOST_QUALIFIED", "environment.example.com") + + wantOptions := GSSAPIOptions{ + ConfigPath: "/connection/krb5.conf", + CCachePath: "/connection/alice.ccache", + KeytabPath: "/connection/alice.keytab", + Principal: "alice@EXAMPLE.COM", + Password: "connection-secret", + QOP: AUTH_INT, + AuthorizationID: "proxy-user", + ServerName: "hive/_HOST@SERVER.EXAMPLE.COM", + ServiceHost: "canonical.example.com", + UseCCache: true, + UseKeytab: true, + UseSSPI: true, + CanonicalizeHost: true, + DisablePAFXFAST: true, + } + backend := &fakeGSSAPIBackend{ + establishAfter: 1, + integrity: true, + confidentiality: true, + initiatorName: "backend@EXAMPLE.COM", + continueOutputs: [][]byte{[]byte("ap-req")}, + securityLayerChallenge: []byte{QOP_TO_FLAG[AUTH] | QOP_TO_FLAG[AUTH_INT], 0, 0x20, 0}, + } + var gotOptions GSSAPIOptions + installFakeGSSAPIBackendWithOptions(t, func(options GSSAPIOptions) (gssapiBackend, error) { + gotOptions = options + return backend, nil + }) + + mechanism, err := NewGSSAPIMechanismWithOptions("ignored-service", wantOptions) + if err != nil { + t.Fatal(err) + } + mechanism.MaxLength = 0x1000 + client := NewSaslClient("alias.example.com", mechanism) + if _, err := client.Start(); err != nil { + t.Fatal(err) + } + response, err := client.Step([]byte("security-layer")) + if err != nil { + t.Fatal(err) + } + + if gotOptions != wantOptions { + t.Fatalf("backend options = %#v, want %#v", gotOptions, wantOptions) + } + if backend.initiatedService != "hive/canonical.example.com@SERVER.EXAMPLE.COM" { + t.Fatalf("service name = %q", backend.initiatedService) + } + if mechanism.qop != QOP_TO_FLAG[AUTH_INT] { + t.Fatalf("selected QOP = %#02x, want auth-int", mechanism.qop) + } + if string(response) != "wrapped:\x02\x00\x10\x00proxy-user" { + t.Fatalf("security-layer response = %q", response) + } +} + +func TestGSSAPIContextClientsKeepOptionsIsolated(t *testing.T) { + t.Setenv("DBX_KRB5_PRINCIPAL", "environment@EXAMPLE.COM") + t.Setenv("SERVICE_HOST_QUALIFIED", "environment.example.com") + + backends := []*fakeGSSAPIBackend{ + {establishAfter: 1, continueOutputs: [][]byte{[]byte("first-token")}}, + {establishAfter: 1, continueOutputs: [][]byte{[]byte("second-token")}}, + } + var received []GSSAPIOptions + installFakeGSSAPIBackendWithOptions(t, func(options GSSAPIOptions) (gssapiBackend, error) { + received = append(received, options) + return backends[len(received)-1], nil + }) + + firstOptions := GSSAPIOptions{Principal: "alice@FIRST.EXAMPLE", ServiceHost: "first.example.com"} + secondOptions := GSSAPIOptions{Principal: "bob@SECOND.EXAMPLE", ServerName: "HTTP/second.example.com@SECOND.EXAMPLE"} + first, err := NewGSSAPIContextClientWithOptions("HTTP", "alias.example.com", firstOptions) + if err != nil { + t.Fatal(err) + } + second, err := NewGSSAPIContextClientWithOptions("HTTP", "other-alias.example.com", secondOptions) + if err != nil { + t.Fatal(err) + } + if token, err := first.Start(nil); err != nil || string(token) != "first-token" { + t.Fatalf("first start token=%q err=%v", token, err) + } + if token, err := second.Start(nil); err != nil || string(token) != "second-token" { + t.Fatalf("second start token=%q err=%v", token, err) + } + + if len(received) != 2 || received[0] != firstOptions || received[1] != secondOptions { + t.Fatalf("received options = %#v", received) + } + if backends[0].initiatedService != "HTTP/first.example.com" { + t.Fatalf("first service = %q", backends[0].initiatedService) + } + if backends[1].initiatedService != "HTTP/second.example.com@SECOND.EXAMPLE" { + t.Fatalf("second service = %q", backends[1].initiatedService) + } +} + +func TestGSSAPIContextClientSupportsChannelBinding(t *testing.T) { + backend := &fakeGSSAPIBackend{ + establishAfter: 2, + continueOutputs: [][]byte{[]byte("initial"), nil}, + } + installFakeGSSAPIBackend(t, backend) + client, err := NewGSSAPIContextClient("HTTP", "hs2.example.com") + if err != nil { + t.Fatal(err) + } + token, err := client.Start([]byte("tls-server-end-point:hash")) + if err != nil { + t.Fatal(err) + } + if string(token) != "initial" { + t.Fatalf("unexpected initial token %q", token) + } + if backend.initiatedService != "HTTP/hs2.example.com" { + t.Fatalf("unexpected service %q", backend.initiatedService) + } + if string(backend.channelBinding) != "tls-server-end-point:hash" { + t.Fatalf("unexpected channel binding %q", backend.channelBinding) + } + if _, err := client.Continue([]byte("mutual")); err != nil { + t.Fatal(err) + } + if !client.Complete() { + t.Fatal("context should be complete") + } + if err := client.Dispose(); err != nil { + t.Fatal(err) + } + if !backend.disposed { + t.Fatal("context backend was not disposed") + } +} + +func TestConfiguredQOPMask(t *testing.T) { + for _, testCase := range []struct { + value string + want byte + }{ + {value: "", want: 7}, + {value: " AUTH ", want: 1}, + {value: "auth-int", want: 2}, + {value: "AUTH-CONF", want: 4}, + } { + t.Run(fmt.Sprintf("%q", testCase.value), func(t *testing.T) { + got, err := configuredQOPMask(testCase.value) + if err != nil { + t.Fatal(err) + } + if got != testCase.want { + t.Fatalf("got %#02x, want %#02x", got, testCase.want) + } + }) + } +} + +func TestQualifiedServiceNameCanonicalizesHost(t *testing.T) { + previousLookup := lookupCanonicalHostname + lookupCanonicalHostname = func(host string) (string, error) { + if host != "alias.example.com" { + t.Fatalf("unexpected lookup host: %s", host) + } + return "hs2.example.com.", nil + } + t.Cleanup(func() { lookupCanonicalHostname = previousLookup }) + t.Setenv("DBX_KRB5_CANONICALIZE_HOST", "true") + t.Setenv("SERVICE_HOST_QUALIFIED", "") + t.Setenv("DBX_KRB5_SERVER_NAME", "") + + if value := qualifiedServiceName("hive/_HOST@EXAMPLE.COM", "alias.example.com"); value != "hive/hs2.example.com@EXAMPLE.COM" { + t.Fatalf("unexpected canonical service name: %s", value) + } +} + +func TestQualifiedServiceNameHonorsExplicitHost(t *testing.T) { + t.Setenv("DBX_KRB5_CANONICALIZE_HOST", "true") + t.Setenv("SERVICE_HOST_QUALIFIED", "explicit.example.com") + t.Setenv("DBX_KRB5_SERVER_NAME", "") + if value := qualifiedServiceName("hive", "alias.example.com"); value != "hive/explicit.example.com" { + t.Fatalf("unexpected explicit service name: %s", value) + } +} diff --git a/agents/go-common/gosasl/sasl_test.go b/agents/go-common/gosasl/sasl_test.go new file mode 100644 index 0000000000..06dc7945ab --- /dev/null +++ b/agents/go-common/gosasl/sasl_test.go @@ -0,0 +1,155 @@ +//go:build !kerberos +// +build !kerberos + +package gosasl + +import ( + "fmt" + "reflect" + "strings" + "testing" +) + +func TestAnonymousMechanism(t *testing.T) { + mechanism := NewAnonymousMechanism() + client := NewSaslClient("localhost", mechanism) + client.Start() + ret, _ := client.Step(nil) + if !reflect.DeepEqual(ret, []byte("Anonymous, None")) { + t.Fatal("Unexpected response from client.process") + } + if !client.Complete() { + t.Fatal("Challenge should have completed") + } + client.Dispose() +} + +func TestPlainMechanism(t *testing.T) { + mechanism := NewPlainMechanism("user", "password") + client := NewSaslClient("localhost", mechanism) + client.Start() + response, _ := client.Step([]byte("abcd")) + if !client.Complete() { + t.Fatal("Challenge should have completed") + } + + NULL := "\x00" + expectedResponse := []byte(fmt.Sprintf("%s%s%s%s", NULL, "user", NULL, "password")) + if !reflect.DeepEqual(response, expectedResponse) { + t.Fatal("Unexpected response from client.process") + } + client.Dispose() +} + +func TestPlainMechanismWithAuthorizationId(t *testing.T) { + mechanism := NewPlainMechanism("user", "password") + client := NewSaslClient("localhost", mechanism) + client.GetConfig().AuthorizationID = "authId" + client.Start() + response, _ := client.Step([]byte("abcd")) + if !client.Complete() { + t.Fatal("Challenge should have completed") + } + + NULL := "\x00" + expectedResponse := []byte(fmt.Sprintf("%s%s%s%s%s", "authId", NULL, "user", NULL, "password")) + if !reflect.DeepEqual(response, expectedResponse) { + t.Fatal("Unexpected response from client.process") + } + client.Dispose() +} + +func TestCramMD5Mechanism(t *testing.T) { + mechanism := NewCramMD5Mechanism("user", "pass") + client := NewSaslClient("localhost", mechanism) + client.Start() + response, _ := client.Step([]byte("msg")) + if !client.Complete() { + t.Fatal("Challenge should have completed") + } + + var expected = []byte{117, 115, 101, 114, 32, 182, 240, 88, 240, 136, 183, 51, 193, 125, 1, 166, 33, 169, 193, 157, 192} + + if !reflect.DeepEqual(response, expected) { + t.Fatalf("Response expected was %x, but got %x", expected, response) + } + + client.Dispose() +} + +func mapFromString(s string) map[string]string { + entries := strings.Split(string(s), ",") + c := make(map[string]string) + for _, e := range entries { + parts := strings.SplitN(e, "=", 2) + c[parts[0]] = parts[1] + } + return c +} + +func TestDigestMD5Mechanism(t *testing.T) { + mechanism := NewDigestMD5Mechanism("imap", "chris", "secret") + client := NewSaslClient("elwood.innosoft.com", mechanism) + client.Start() + challenge := `utf-8,username="chris",realm="elwood.innosoft.com",nonce="OA6MG9tEQGm2hh",nc=00000001,cnonce="OA6MHXh6VqTrRk",digest-uri="imap/elwood.innosoft.com",response=d388dad90d4bbd760a152321f2143af7,qop=auth` + response, err := client.Step([]byte(challenge)) + if err != nil { + t.Fatal(err) + } + if client.Complete() { + t.Fatal("Challenge should not have completed") + } + + var expected = `username="chris",realm="elwood.innosoft.com",qop=auth,nonce="OA6MG9tEQGm2hh",nc=00000001,digest-uri="imap/elwood.innosoft.com"` + expectedMap := mapFromString(expected) + actualMap := mapFromString(string(response)) + delete(actualMap, "cnonce") + delete(actualMap, "response") + + if !reflect.DeepEqual(actualMap, expectedMap) { + t.Fatalf("Response expected was %s, but got %s", expectedMap, actualMap) + } + + serverResponse := []byte("rspauth=ea40f60335c427b5527b84dbabcdfffd") + response, err = client.Step([]byte(serverResponse)) + if !client.Complete() { + t.Fatal("Challenge should have completed") + } + if response != nil { + t.Fatalf("Response should be nil, instead: %s", response) + } + + client.Dispose() +} + +func TestDigestMD5ChallengeParserHandlesQuotedCommasAndEscapes(t *testing.T) { + values, err := parseChallenge([]byte(`realm="EXAMPLE, INC",nonce="a\\\"b",qop="auth,auth-int"`)) + if err != nil { + t.Fatal(err) + } + if values["realm"] != "EXAMPLE, INC" || values["nonce"] != `a\"b` || values["qop"] != "auth,auth-int" { + t.Fatalf("unexpected challenge values: %#v", values) + } +} + +func TestDigestMD5RejectsMalformedChallengeWithoutPanicking(t *testing.T) { + mechanism := NewDigestMD5Mechanism("hive", "user", "password") + client := NewSaslClient("hs2.example.com", mechanism) + if _, err := client.Start(); err != nil { + t.Fatal(err) + } + if _, err := client.Step([]byte(`nonce="unterminated`)); err == nil { + t.Fatal("expected malformed challenge error") + } +} + +func TestDigestMD5RejectsUnsupportedQOP(t *testing.T) { + mechanism := NewDigestMD5Mechanism("hive", "user", "password") + client := NewSaslClient("hs2.example.com", mechanism) + if _, err := client.Start(); err != nil { + t.Fatal(err) + } + if _, err := client.Step([]byte(`realm="EXAMPLE.COM",nonce="nonce",qop="auth-conf"`)); err == nil || !strings.Contains(err.Error(), "qop=auth") { + t.Fatalf("expected unsupported QOP error, got %v", err) + } +} diff --git a/agents/scripts/driver_release_packages_test.py b/agents/scripts/driver_release_packages_test.py index a9e1b78f62..be89f9ae9d 100644 --- a/agents/scripts/driver_release_packages_test.py +++ b/agents/scripts/driver_release_packages_test.py @@ -41,6 +41,7 @@ class DriverReleasePackagesTest(unittest.TestCase): "duckdb": "0.1.0", "rabbitmq": "0.1.0", "cassandra": "0.1.37", + "hive": "0.1.43", "tdengine": "0.1.0", } @@ -261,6 +262,21 @@ class DriverReleasePackagesTest(unittest.TestCase): self.assertFalse(source.exists()) self.assertEqual(versioned.read_bytes(), b"\x7fELFtest-iotdb-agent") + def test_versions_hive_native_artifacts(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + release_dir = Path(temp_dir) + source = release_dir / "dbx-agent-hive-windows-x64.exe" + source.write_bytes(b"MZtest-hive-agent") + versions = {driver: "0.1.0" for driver in NATIVE_DRIVERS} + versions["hive"] = "0.1.44" + + renamed = version_agent_artifacts(release_dir, versions) + versioned = release_dir / "dbx-agent-hive-0.1.44-windows-x64.exe" + + self.assertEqual(renamed, [versioned]) + self.assertFalse(source.exists()) + self.assertEqual(versioned.read_bytes(), b"MZtest-hive-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 8a77a431cc..e4c51fa0b1 100644 --- a/agents/scripts/validate_agents.py +++ b/agents/scripts/validate_agents.py @@ -15,6 +15,7 @@ NON_JDBC_AGENT_MODULES = {"mongodb", "etcd", "zookeeper", "kafka", "rocketmq", " NATIVE_ONLY_AGENT_MODULES = { "cassandra": "drivers/cassandra-go", "duckdb": "drivers/duckdb", + "hive": "drivers/hive-go", "oracle": "drivers/oracle-go", "kingbase": "drivers/kingbase-go", "iotdb": "drivers/iotdb", diff --git a/agents/scripts/validate_agents_test.py b/agents/scripts/validate_agents_test.py index 849d464388..402ecbf02d 100644 --- a/agents/scripts/validate_agents_test.py +++ b/agents/scripts/validate_agents_test.py @@ -138,6 +138,7 @@ class ValidateAgentsTest(unittest.TestCase): ) for driver in ( "cassandra-go", + "hive-go", "oracle-go", "kingbase-go", "iotdb", @@ -150,7 +151,7 @@ class ValidateAgentsTest(unittest.TestCase): ): (root / "drivers" / driver).mkdir(parents=True) (root / "versions.json").write_text( - json.dumps({"h2": "0.1.0", "cassandra": "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", "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", "tdengine": "0.1.0"}), encoding="utf-8", ) diff --git a/agents/scripts/version_agent_artifacts.py b/agents/scripts/version_agent_artifacts.py index f272f6586a..c4ee4f4eae 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", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "duckdb", "rabbitmq", "tdengine") +NATIVE_DRIVERS = ("cassandra", "hive", "oracle", "xugu", "kingbase", "iotdb", "neo4j", "vastbase", "duckdb", "rabbitmq", "tdengine") PLATFORMS = ( "macos-aarch64", "macos-x64", diff --git a/agents/settings.gradle b/agents/settings.gradle index ca1f6648c4..cf73b32c8d 100644 --- a/agents/settings.gradle +++ b/agents/settings.gradle @@ -4,7 +4,7 @@ def infrastructureModules = ['common', 'test-support'] def driverModules = [ 'access', 'dameng', 'goldendb', 'databend', 'databricks', 'saphana', 'teradata', 'vertica', 'firebird', 'exasol', 'oceanbase-oracle', 'gbase8a', 'gbase8s', - 'bigquery', 'kylin', 'sundb', 'h2', 'h2-legacy', 'snowflake', 'trino', 'hive', 'spark', + 'bigquery', 'kylin', 'sundb', 'h2', 'h2-legacy', 'snowflake', 'trino', 'spark', 'db2', 'informix', 'mongodb', 'highgo', 'uxdb', 'yashandb', 'oscar', 'iris', 'etcd', 'zookeeper', 'kafka', 'rocketmq', 'sqlserver-legacy' ] diff --git a/crates/dbx-core/src/agent_manager.rs b/crates/dbx-core/src/agent_manager.rs index 8969a4d2e1..8b6a30fbc4 100644 --- a/crates/dbx-core/src/agent_manager.rs +++ b/crates/dbx-core/src/agent_manager.rs @@ -275,6 +275,22 @@ mod tests { assert_eq!(launch.working_dir.as_deref(), Some(manager.driver_dir("dameng").as_path())); } + #[test] + fn hive_native_launch_ignores_stale_legacy_jar_without_jre() { + let manager = test_manager("hive-native-over-legacy-jar"); + let native = manager.driver_native_path("hive"); + touch(&native); + touch(&manager.driver_jar_path("hive")); + + let launch = manager + .resolve_agent_launch_spec(&AgentState::default(), "hive", DEFAULT_JRE_KEY) + .expect("Hive native launch should not require a JRE"); + + assert_eq!(launch.program, native); + assert_eq!(launch.args, Vec::::new()); + assert_eq!(launch.working_dir.as_deref(), Some(manager.driver_dir("hive").as_path())); + } + #[test] fn resolves_relative_native_agent_launch_with_absolute_paths() { let base_dir = diff --git a/crates/dbx-core/src/agent_service.rs b/crates/dbx-core/src/agent_service.rs index 1681c0652d..31e9f73b56 100644 --- a/crates/dbx-core/src/agent_service.rs +++ b/crates/dbx-core/src/agent_service.rs @@ -2930,14 +2930,16 @@ mod agent_registry_install_tests { } #[tokio::test] - async fn registry_install_accepts_native_driver_with_legacy_jar_fallback() { - let manager = test_manager("native-with-jar-fallback"); - let db_type = "oracle"; + async fn registry_install_replaces_hive_legacy_jar_with_native_driver() { + let manager = test_manager("hive-native-replaces-legacy-jar"); + let db_type = "hive"; let version = "0.1.31"; - let native_url = "https://example.com/dbx-agent-oracle"; + let native_url = "https://example.com/dbx-agent-hive"; let native_bytes = b"native-agent"; let registry = registry_with_native_and_legacy_jar(db_type, version, native_url, native_bytes.len() as u64); let native_path = manager.driver_native_path(db_type); + std::fs::create_dir_all(manager.driver_dir(db_type)).unwrap(); + write_test_agent_jar(&manager.driver_jar_path(db_type)); let cache_path = write_cached_driver_download(&manager, db_type, version, native_url, &native_path, native_bytes); let events = std::sync::Mutex::new(Vec::new());