mirror of
https://github.com/block/buzz.git
synced 2026-10-27 04:30:52 +00:00
feat(buzz-agent): add opt-in workspace-scoped Databricks reuse (#7606)
Posted by Brain, Wes Billman’s AI assistant, on behalf of `wesbillman`. ## Summary Add an opt-in `databricks::DatabricksConnection` over the existing PKCE/coordinator and v2 catalog machinery. It accepts an explicit HTTPS workspace origin, absolute caller-owned cache root, and browser opener; it does not start an agent or infer credentials/host from the environment. - Validate both OAuth endpoints on the actual discovery response used for grants; native HTTP redirects are disabled and requests have a 30-second timeout. - Keep strict cache/single-flight state in a separate namespace. Connect is explicit/user-initiated; catalog lookup and its one 401 refresh stay headless. Tokens are not exposed by the connection API. - Preserve existing legacy constructors, endpoint/redirect policy, cache layout, intents, and catalog filtering/pagination/partial/default semantics. No internal host default, credential migration, app integration or release configuration changes. ### Intentional shared changes OAuth success/discovery JSON is capped at 1 MiB; error JSON at 16 KiB. Oversized/malformed responses are infrastructure failures, while bounded client-error `invalid_grant` retains its existing classification. OAuth diagnostics omit raw bodies/URLs/opener/callback details. Overflowing token expiry returns an infrastructure failure. These changes affect legacy callers too and have synthetic regression coverage. See [the reuse contract](crates/buzz-agent/DATABRICKS_REUSE.md) for API, compatibility and caller obligations. Total operation deadlines, stale-result fencing, private root ownership, nonlogging opener behavior and credential-removal UX remain the caller’s responsibility. ### Related issue Related discovery work: #6918. Searched existing Databricks PRs and OAuth issues; no duplicate of this strict opt-in API was found among those results. This is a new reuse boundary, not a claim of a live credential leak. ### Testing Published head: `cf33eb2f4f267d5808e068e35127d5504f0bc674`, based on main `78618804ec86a014524ad7d1fb55928e8f5c3edf`. Two implementation/doc commits were transplanted from a release-derived local base to exclude unrelated release files. `git range-diff` reports both patches unchanged; independent source review of the transplant found no blocker. **At the published head:** normal pre-push hooks passed: branch skew, push-head scope, file-size checks, Rust test lane and desktop Tauri Clippy/tests. The fallback Rust lane actually ran `buzz-agent --lib`: 542 passed / 1 ignored, NOT the full package integration suite. Hosted full package/repository CI remains pending. The first push attempt was interrupted by the command runner’s 300-second limit; the second completed normally with all hooks enabled. **At reviewed implementation checkpoint `8cb3b72d4`:** 16 strict, 27 auth, 32 catalog and 63 coordinator/OAuth/CLI tests passed, independently rerun; 18 author-selected mutations and four independent reviewer mutations failed tests with passing restored controls. The subsequent `539319411` delta changed only two doc-comment blocks. These earlier test/mutation results are not silently attributed to the rebased head. Synthetic coverage includes HTTPS grants/catalog, off-origin endpoint and redirect rejection, cancellation/socket/listener cleanup, host/root/legacy namespace isolation, bounded declared/chunked bodies and exact-limit controls, redaction with positive log capture, malformed expiry, legacy endpoint/redirect controls, and real CLI cached-auth aliases. Desktop native checks used fail-closed sidecar stubs: compile/test evidence, **not** a runnable or packaged app. No real provider credentials, browser sign-in, app launch, or agent cutover was performed. Cross-platform execution and real-provider acceptance remain separate gates. No merge or downstream adoption is requested by this PR. --------- Signed-off-by: Brain <1a02c72794dcd0f07058a353bc3a81f4028b8c77c92c87fce6d5c8b85970a20b@buzz.block.builderlab.xyz> Co-authored-by: Brain <1a02c72794dcd0f07058a353bc3a81f4028b8c77c92c87fce6d5c8b85970a20b@buzz.block.builderlab.xyz>
This commit is contained in:
Wes
123
Cargo.lock
generated
123
Cargo.lock
generated
@@ -209,6 +209,45 @@ version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
|
||||
|
||||
[[package]]
|
||||
name = "asn1-rs"
|
||||
version = "0.7.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b7f43a50ac4fdca5df8e885c21b835997f0a1cdee65494a6847694a98652d9d8"
|
||||
dependencies = [
|
||||
"asn1-rs-derive",
|
||||
"asn1-rs-impl",
|
||||
"displaydoc",
|
||||
"nom",
|
||||
"num-traits",
|
||||
"rusticata-macros",
|
||||
"thiserror 2.0.18",
|
||||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "asn1-rs-derive"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3109e49b1e4909e9db6515a30c633684d68cdeaa252f215214cb4fa1a5bfee2c"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
"synstructure",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "asn1-rs-impl"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7b18050c2cd6fe86c3a76584ef5e0baf286d038cda203eb6223df2cc413565f7"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-broadcast"
|
||||
version = "0.7.2"
|
||||
@@ -646,6 +685,15 @@ version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
|
||||
|
||||
[[package]]
|
||||
name = "bit-vec"
|
||||
version = "0.9.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b71798fca2c1fe1086445a7258a4bc81e6e49dcd24c8d0dd9a1e57395b603f51"
|
||||
dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bitcoin-consensus-encoding"
|
||||
version = "1.1.0"
|
||||
@@ -898,6 +946,7 @@ dependencies = [
|
||||
"getrandom 0.4.3",
|
||||
"hex",
|
||||
"nix 0.31.3",
|
||||
"rcgen",
|
||||
"reqwest 0.13.4",
|
||||
"rmcp",
|
||||
"serde",
|
||||
@@ -906,6 +955,7 @@ dependencies = [
|
||||
"sha2 0.11.0",
|
||||
"tempfile",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"url",
|
||||
@@ -2399,6 +2449,20 @@ dependencies = [
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "der-parser"
|
||||
version = "10.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "07da5016415d5a3c4dd39b11ed26f915f52fc4e0dc197d87908bc916e51bc1a6"
|
||||
dependencies = [
|
||||
"asn1-rs",
|
||||
"displaydoc",
|
||||
"nom",
|
||||
"num-bigint",
|
||||
"num-traits",
|
||||
"rusticata-macros",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "deranged"
|
||||
version = "0.5.8"
|
||||
@@ -6273,6 +6337,15 @@ dependencies = [
|
||||
"objc2-security",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "oid-registry"
|
||||
version = "0.8.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "12f40cff3dde1b6087cc5d5f5d4d65712f34016a03ed60e9c08dcc392736b5b7"
|
||||
dependencies = [
|
||||
"asn1-rs",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "once_cell"
|
||||
version = "1.21.4"
|
||||
@@ -7697,6 +7770,19 @@ dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rcgen"
|
||||
version = "0.14.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8774e05a7d0de114588e6a28fe7e71694b82614ed569d86d8b389dfbc98b8ad8"
|
||||
dependencies = [
|
||||
"ring",
|
||||
"rustls-pki-types",
|
||||
"time",
|
||||
"x509-parser",
|
||||
"yasna",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redb"
|
||||
version = "3.1.3"
|
||||
@@ -8080,6 +8166,15 @@ dependencies = [
|
||||
"semver",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rusticata-macros"
|
||||
version = "4.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "faf0c4a6ece9950b9abdb62b1cfcf2a68b3b67a10ba445b3bb85be2a293d0632"
|
||||
dependencies = [
|
||||
"nom",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustix"
|
||||
version = "0.38.44"
|
||||
@@ -11435,6 +11530,24 @@ dependencies = [
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "x509-parser"
|
||||
version = "0.18.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d43b0f71ce057da06bc0851b23ee24f3f86190b07203dd8f567d0b706a185202"
|
||||
dependencies = [
|
||||
"asn1-rs",
|
||||
"data-encoding",
|
||||
"der-parser",
|
||||
"lazy_static",
|
||||
"nom",
|
||||
"oid-registry",
|
||||
"ring",
|
||||
"rusticata-macros",
|
||||
"thiserror 2.0.18",
|
||||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "xattr"
|
||||
version = "1.6.1"
|
||||
@@ -11617,6 +11730,16 @@ version = "0.8.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3"
|
||||
|
||||
[[package]]
|
||||
name = "yasna"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b5f6765e852b9b4dc8e2a76843e4d64d1cea8e79bcde0b6901aea8e7c7f08282"
|
||||
dependencies = [
|
||||
"bit-vec 0.9.1",
|
||||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "yoke"
|
||||
version = "0.8.2"
|
||||
|
||||
@@ -73,6 +73,9 @@ fs2 = "0.4"
|
||||
nix = { version = "0.31", default-features = false, features = ["signal", "process"] }
|
||||
|
||||
[dev-dependencies]
|
||||
# Synthetic TLS OAuth/catalog servers exercise the strict HTTPS production path.
|
||||
rcgen = { version = "0.14", default-features = false, features = ["crypto", "ring"] }
|
||||
tokio-rustls = { version = "0.26", default-features = false, features = ["ring", "tls12"] }
|
||||
tokio = { workspace = true, features = ["test-util", "rt-multi-thread", "macros", "io-std", "io-util", "sync", "process", "time", "net"] }
|
||||
nix = { version = "0.31", default-features = false, features = ["signal", "process"] }
|
||||
axum = { workspace = true }
|
||||
|
||||
70
crates/buzz-agent/DATABRICKS_REUSE.md
Normal file
70
crates/buzz-agent/DATABRICKS_REUSE.md
Normal file
@@ -0,0 +1,70 @@
|
||||
# Opt-in Databricks connection reuse
|
||||
|
||||
`databricks::DatabricksConnection` reuses the PKCE coordinator and v2 catalog
|
||||
without running an agent, ACP, a CLI subprocess or any runtime configuration.
|
||||
|
||||
- Construct with an explicit HTTPS workspace origin, an absolute app-owned cache
|
||||
root and a `BrowserOpener`. No host default or environment credential lookup.
|
||||
Construction can read/create that cache, but never networks or opens a browser.
|
||||
- `connect()` is user-initiated authentication only; it does not fetch models.
|
||||
- `discover_models(filter)` is always headless, including one 401-driven refresh.
|
||||
It never selects a model. Existing filtering, partial-success and labelled
|
||||
authenticated-empty/no-filter fallback remain; a vector is not proof of a
|
||||
complete catalog.
|
||||
- Strict discovery validates both OAuth endpoints from the very response used
|
||||
for refresh/code exchange. They must be same-origin HTTPS with no userinfo,
|
||||
query or fragment. Native discovery/token/catalog redirects are disabled.
|
||||
Browser SSO navigation is separate from the native client's redirect policy.
|
||||
- The caller chooses and exclusively owns its cache root BEFORE construction.
|
||||
`databricks-strict` separates strict state from legacy `databricks` caches and
|
||||
single-flight coordination even if roots coincide. It is not credential
|
||||
migration. Existing host/client/scopes hash and Unix owner-only atomic token
|
||||
persistence remain. Non-Unix token persistence remains disabled.
|
||||
- Keep secrets native. The supplied opener must not log authorization URLs or
|
||||
its failure details. Drop the operation future to cancel actual requests and
|
||||
the callback listener. Set an app-level overall deadline, generation-fence
|
||||
results, and never await under the agent-controller lock. A synchronous opener
|
||||
must return promptly; the library cannot cancel a blocking caller callback.
|
||||
- A cache root is an application trust boundary, not protection against a
|
||||
malicious process with the same OS user or a caller lending out that root.
|
||||
|
||||
## Legacy compatibility and intentional shared changes
|
||||
|
||||
Existing public OAuth config/constructors, auth intents, cache defaults, static
|
||||
bearer selection, CLI auth, desktop discovery and v1/v2 catalog APIs remain.
|
||||
They DO NOT inherit strict same-origin or no-redirect policy.
|
||||
|
||||
Shared OAuth response handling has three intentional observable differences:
|
||||
|
||||
1. Discovery and successful token responses are stream-capped at **1 MiB**;
|
||||
OAuth error responses at **16 KiB**. Exact-boundary JSON succeeds; oversized,
|
||||
malformed or unreadable responses are infrastructure failures, never grant
|
||||
rejection. Only a client-error `error: invalid_grant` within the bound retains
|
||||
refresh/code rejection classification. Non-success discovery bodies are not
|
||||
consumed. Existing catalog limits/retries are unchanged.
|
||||
2. OAuth logs no longer include raw response bodies, transport/decode details,
|
||||
callback error text or opener error text. Status/category logs remain. The
|
||||
legacy default opener still prints its manual sign-in URL deliberately; the
|
||||
strict API never selects that opener automatically. Strict catalog errors
|
||||
project fixed auth/infrastructure messages; legacy catalog diagnostics remain.
|
||||
3. A token response whose `expires_in` overflows epoch arithmetic now fails as an
|
||||
infrastructure error instead of panicking/wrapping.
|
||||
|
||||
No old cache migration/deletion, credential rewrite, app launch, internal host,
|
||||
internal release settings, or production dependency on a developer checkout is
|
||||
part of this patch.
|
||||
|
||||
## Validation boundaries
|
||||
|
||||
Synthetic HTTPS servers use ephemeral test-only certificates; all cache contents,
|
||||
refresh tokens and callback codes are fabricated. Tests cover new and legacy
|
||||
paths, positive destination controls, refresh/code exchange, redirects, response
|
||||
limits, classifications, captured diagnostics, cancellation/retry, cache roots
|
||||
and namespace isolation. A real CLI child tests legacy auth aliases/cache reuse;
|
||||
existing coordinator tests include real subprocess single-flight/locking.
|
||||
|
||||
Before integration, require independent compatibility review at an exact commit,
|
||||
full `buzz-agent` package-suite evidence, affected desktop checks and repository
|
||||
gates. Focused tests are not a substitute. App orchestration, internal build-input
|
||||
wiring, credential-removal UX, native preview acceptance and real sign-in are
|
||||
separate subsequent work, not claims made by this helper patch.
|
||||
@@ -31,7 +31,7 @@ use serde_json::Value;
|
||||
use sha2::Digest;
|
||||
use tokio::sync::{watch, Mutex};
|
||||
|
||||
use crate::types::AgentError;
|
||||
use crate::{auth_http::read_oauth_json, databricks::StrictWorkspace, types::AgentError};
|
||||
|
||||
/// Buffer before `expires_at` to consider a cached token "still good".
|
||||
/// Keeps us off the cliff if the clock or the server's clock drifts.
|
||||
@@ -370,6 +370,7 @@ enum RefreshOutcome {
|
||||
pub struct PkceOAuthTokenSource {
|
||||
cfg: PkceOAuthConfig,
|
||||
http: Client,
|
||||
workspace: Option<StrictWorkspace>,
|
||||
cache_path: PathBuf,
|
||||
/// Injected browser launcher, called inside [`browser_pkce_flow`] while the
|
||||
/// localhost listener is live. Production uses [`DefaultBrowserOpener`];
|
||||
@@ -412,26 +413,44 @@ impl PkceOAuthTokenSource {
|
||||
cfg: PkceOAuthConfig,
|
||||
opener: Arc<dyn BrowserOpener>,
|
||||
http_timeout: Duration,
|
||||
) -> Result<Arc<Self>, AgentError> {
|
||||
let http = Client::builder()
|
||||
.timeout(http_timeout)
|
||||
.build()
|
||||
.map_err(|_| AgentError::Llm("oauth http client construction failed".into()))?;
|
||||
Self::from_parts(cfg, opener, http, None)
|
||||
}
|
||||
|
||||
pub(crate) fn for_workspace(
|
||||
workspace: StrictWorkspace,
|
||||
cache_root: &Path,
|
||||
opener: Arc<dyn BrowserOpener>,
|
||||
http: Client,
|
||||
) -> Result<Arc<Self>, AgentError> {
|
||||
let mut cfg =
|
||||
crate::llm::databricks_pkce_config(workspace.as_str(), Some(cache_root.to_path_buf()));
|
||||
// Strict and legacy sources must not share single-flight/cache state,
|
||||
// even if a caller accidentally supplies the same root to both APIs.
|
||||
cfg.cache_namespace = "databricks-strict".into();
|
||||
Self::from_parts(cfg, opener, http, Some(workspace))
|
||||
}
|
||||
|
||||
fn from_parts(
|
||||
cfg: PkceOAuthConfig,
|
||||
opener: Arc<dyn BrowserOpener>,
|
||||
http: Client,
|
||||
workspace: Option<StrictWorkspace>,
|
||||
) -> Result<Arc<Self>, AgentError> {
|
||||
let cache_path = cache_path_for(&cfg)?;
|
||||
if let Some(parent) = cache_path.parent() {
|
||||
fs::create_dir_all(parent)
|
||||
.map_err(|e| AgentError::Llm(format!("oauth cache dir {parent:?}: {e}")))?;
|
||||
}
|
||||
// Every OAuth HTTP call inherits this timeout so a hung provider can
|
||||
// never stall the caller — nor the same-key callers waiting on the
|
||||
// cross-process lock this holder owns. Construction is fallible, so a
|
||||
// build failure propagates rather than silently falling back to an
|
||||
// untimed client — an untimed client would restore exactly the
|
||||
// unbounded-HTTP-under-lock failure the timeout exists to prevent.
|
||||
let http = Client::builder()
|
||||
.timeout(http_timeout)
|
||||
.build()
|
||||
.map_err(|e| AgentError::Llm(format!("oauth http client: {e}")))?;
|
||||
let initial = read_cache(&cache_path);
|
||||
Ok(Arc::new(Self {
|
||||
cfg,
|
||||
http,
|
||||
workspace,
|
||||
cache_path,
|
||||
opener,
|
||||
state: Mutex::new(initial),
|
||||
@@ -460,17 +479,18 @@ impl PkceOAuthTokenSource {
|
||||
|
||||
/// Discover authorization + token endpoints from the well-known URL.
|
||||
async fn endpoints(&self) -> Result<OidcEndpoints, AgentError> {
|
||||
let v: Value = self
|
||||
let response = self
|
||||
.http
|
||||
.get(&self.cfg.discovery_url)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth discovery: {e}")))?
|
||||
.error_for_status()
|
||||
.map_err(|e| AgentError::Llm(format!("oauth discovery status: {e}")))?
|
||||
.json()
|
||||
.map_err(|_| AgentError::Llm("oauth discovery request failed".into()))?;
|
||||
if !response.status().is_success() {
|
||||
return Err(AgentError::Llm("oauth discovery status failed".into()));
|
||||
}
|
||||
let v = read_oauth_json(response)
|
||||
.await
|
||||
.map_err(|e| AgentError::Llm(format!("oauth discovery json: {e}")))?;
|
||||
.map_err(|_| AgentError::Llm("oauth discovery response invalid or too large".into()))?;
|
||||
let auth = v
|
||||
.get("authorization_endpoint")
|
||||
.and_then(Value::as_str)
|
||||
@@ -483,6 +503,10 @@ impl PkceOAuthTokenSource {
|
||||
.and_then(Value::as_str)
|
||||
.ok_or_else(|| AgentError::Llm("oauth discovery: token_endpoint missing".into()))?
|
||||
.to_string();
|
||||
if let Some(workspace) = &self.workspace {
|
||||
workspace.validate_endpoint(&auth)?;
|
||||
workspace.validate_endpoint(&token)?;
|
||||
}
|
||||
Ok(OidcEndpoints {
|
||||
authorization_endpoint: auth,
|
||||
token_endpoint: token,
|
||||
@@ -664,14 +688,14 @@ impl PkceOAuthTokenSource {
|
||||
Ok(resp) => resp,
|
||||
// Transport error or the per-request timeout elapsed: no verdict
|
||||
// from the provider, so this is infrastructural, not a rejection.
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "oauth refresh transport failure");
|
||||
Err(_) => {
|
||||
tracing::warn!("oauth refresh transport failure");
|
||||
return RefreshOutcome::Network;
|
||||
}
|
||||
};
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
let body = read_oauth_json(resp).await.ok();
|
||||
// Per RFC 6749 §5.2 only `error == "invalid_grant"` means the
|
||||
// refresh token itself is dead (expired/revoked) — the one failure
|
||||
// a browser sign-in can repair. Every other 4xx (`invalid_request`,
|
||||
@@ -681,22 +705,21 @@ impl PkceOAuthTokenSource {
|
||||
// they stay in the non-credential bucket and surface as
|
||||
// `NetworkUnavailable` without ever popping a browser.
|
||||
if status.is_client_error()
|
||||
&& serde_json::from_str::<Value>(&body)
|
||||
.ok()
|
||||
&& body
|
||||
.and_then(|v| v.get("error").and_then(Value::as_str).map(str::to_owned))
|
||||
.as_deref()
|
||||
== Some("invalid_grant")
|
||||
{
|
||||
tracing::warn!(status = %status, body = %body, "oauth refresh grant rejected");
|
||||
tracing::warn!(status = %status, "oauth refresh grant rejected");
|
||||
return RefreshOutcome::Rejected;
|
||||
}
|
||||
tracing::warn!(status = %status, body = %body, "oauth refresh not repairable by browser");
|
||||
tracing::warn!(status = %status, "oauth refresh not repairable by browser");
|
||||
return RefreshOutcome::Network;
|
||||
}
|
||||
let v: Value = match resp.json().await {
|
||||
let v: Value = match read_oauth_json(resp).await {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "oauth refresh response decode failure");
|
||||
Err(_) => {
|
||||
tracing::warn!("oauth refresh response decode failure");
|
||||
return RefreshOutcome::Network;
|
||||
}
|
||||
};
|
||||
@@ -1952,13 +1975,18 @@ fn token_from_response(
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
.or_else(|| fallback_refresh.map(str::to_string));
|
||||
let expires_at = v.get("expires_in").and_then(Value::as_u64).map(|secs| {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
+ secs
|
||||
});
|
||||
let expires_at = v
|
||||
.get("expires_in")
|
||||
.and_then(Value::as_u64)
|
||||
.map(|secs| {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
.checked_add(secs)
|
||||
.ok_or_else(|| AgentError::Llm("oauth: token expiry overflow".into()))
|
||||
})
|
||||
.transpose()?;
|
||||
Ok(CachedToken {
|
||||
access_token,
|
||||
refresh_token,
|
||||
@@ -1990,8 +2018,8 @@ fn random_state() -> Result<String, AgentError> {
|
||||
/// the *static* HTML shown in the browser. The page never embeds any request
|
||||
/// parameter — the `error` query value is attacker-influenceable, so
|
||||
/// reflecting it would be an XSS sink on the localhost callback. Failure
|
||||
/// detail travels only through `result`, which surfaces in the process error
|
||||
/// and logs, never in the served markup.
|
||||
/// detail travels only through `result`; the waiting flow discards it and
|
||||
/// reports a typed error and fixed diagnostic. The page stays static.
|
||||
fn callback_outcome(
|
||||
params: &std::collections::HashMap<String, String>,
|
||||
expected_state: &str,
|
||||
@@ -2012,10 +2040,9 @@ fn callback_outcome(
|
||||
(result, page)
|
||||
}
|
||||
|
||||
/// Neutralize an attacker-controllable OAuth `error` value before it enters
|
||||
/// an error string that later reaches the logs. Control characters (CR/LF in
|
||||
/// particular) enable log-line injection, and an unbounded value could flood
|
||||
/// the logs — replace control chars with spaces and cap the length.
|
||||
/// Bound and normalize an attacker-controllable OAuth `error` value carried
|
||||
/// internally through the callback result. The waiting flow discards this detail;
|
||||
/// it must not be reflected in browser markup, outward errors or logs.
|
||||
fn sanitize_callback_detail(raw: &str) -> String {
|
||||
const MAX: usize = 200;
|
||||
raw.chars()
|
||||
@@ -2105,8 +2132,8 @@ async fn browser_pkce_flow(
|
||||
|
||||
// Launch the browser while the listener is live. A launch failure aborts
|
||||
// before we wait on a redirect nobody can send.
|
||||
opener.open(&auth_url).map_err(|e| {
|
||||
tracing::warn!(error = %e, "oauth browser launch failed");
|
||||
opener.open(&auth_url).map_err(|_| {
|
||||
tracing::warn!("oauth browser launch failed");
|
||||
AuthError::BrowserOpenFailed
|
||||
})?;
|
||||
|
||||
@@ -2116,8 +2143,8 @@ async fn browser_pkce_flow(
|
||||
// Callback task dropped the sender without sending — treat as timeout.
|
||||
Ok(Err(_)) => return Err(AuthError::TimedOut),
|
||||
// Provider/user reported an error (denial, state mismatch, missing code).
|
||||
Ok(Ok(Err(detail))) => {
|
||||
tracing::warn!(detail = %detail, "oauth callback reported failure");
|
||||
Ok(Ok(Err(_))) => {
|
||||
tracing::warn!("oauth callback reported failure");
|
||||
return Err(AuthError::Denied);
|
||||
}
|
||||
Ok(Ok(Ok(code))) => code,
|
||||
@@ -2141,7 +2168,7 @@ async fn browser_pkce_flow(
|
||||
.map_err(|_| AuthError::NetworkUnavailable)?;
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
let body = read_oauth_json(resp).await.ok();
|
||||
// Only a 4xx `invalid_grant` (RFC 6749 §6.4.1) establishes the
|
||||
// authorization code itself was rejected — the terminal, cooldown-worthy
|
||||
// `ExchangeFailed`. A 429, any 5xx, and any other/unparseable 4xx are a
|
||||
@@ -2150,24 +2177,22 @@ async fn browser_pkce_flow(
|
||||
// refresh classifier, which likewise keys on the body `error`, not the
|
||||
// bare status class.
|
||||
if status.is_client_error()
|
||||
&& serde_json::from_str::<Value>(&body)
|
||||
.ok()
|
||||
&& body
|
||||
.and_then(|v| v.get("error").and_then(Value::as_str).map(str::to_owned))
|
||||
.as_deref()
|
||||
== Some("invalid_grant")
|
||||
{
|
||||
tracing::warn!(status = %status, body = %body, "oauth code exchange rejected");
|
||||
tracing::warn!(status = %status, "oauth code exchange rejected");
|
||||
return Err(AuthError::ExchangeFailed);
|
||||
}
|
||||
tracing::warn!(status = %status, body = %body, "oauth code exchange not a grant rejection");
|
||||
tracing::warn!(status = %status, "oauth code exchange not a grant rejection");
|
||||
return Err(AuthError::NetworkUnavailable);
|
||||
}
|
||||
// A 2xx whose body is missing/malformed or lacks an access token is a
|
||||
// provider fault, not a rejected grant: it never establishes that the code
|
||||
// was refused, so it stays in the transient bucket rather than poisoning a
|
||||
// 5-minute cooldown.
|
||||
let v: Value = resp
|
||||
.json()
|
||||
let v = read_oauth_json(resp)
|
||||
.await
|
||||
.map_err(|_| AuthError::NetworkUnavailable)?;
|
||||
token_from_response(&v, None).map_err(|_| AuthError::NetworkUnavailable)
|
||||
|
||||
30
crates/buzz-agent/src/auth_http.rs
Normal file
30
crates/buzz-agent/src/auth_http.rs
Normal file
@@ -0,0 +1,30 @@
|
||||
//! Bounded OAuth response decoding shared by legacy and strict callers.
|
||||
//!
|
||||
//! Deliberate compatibility change: oversized discovery/token documents are
|
||||
//! infrastructure failures, never grant rejections. Raw provider diagnostics
|
||||
//! are not returned or logged (they can echo codes, verifiers and tokens).
|
||||
|
||||
use reqwest::Response;
|
||||
use serde_json::Value;
|
||||
|
||||
pub(crate) const MAX_OAUTH_RESPONSE_BYTES: usize = 1024 * 1024;
|
||||
pub(crate) const MAX_OAUTH_ERROR_BYTES: usize = 16 * 1024;
|
||||
|
||||
pub(crate) async fn read_oauth_json(mut response: Response) -> Result<Value, ()> {
|
||||
let limit = if response.status().is_success() {
|
||||
MAX_OAUTH_RESPONSE_BYTES
|
||||
} else {
|
||||
MAX_OAUTH_ERROR_BYTES
|
||||
};
|
||||
if response.content_length().is_some_and(|n| n > limit as u64) {
|
||||
return Err(());
|
||||
}
|
||||
let mut bytes = Vec::with_capacity(limit.min(16 * 1024));
|
||||
while let Some(chunk) = response.chunk().await.map_err(|_| ())? {
|
||||
if chunk.len() > limit - bytes.len() {
|
||||
return Err(());
|
||||
}
|
||||
bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
serde_json::from_slice(&bytes).map_err(|_| ())
|
||||
}
|
||||
@@ -158,20 +158,27 @@ pub async fn discover_databricks_models_with_cache_dir(
|
||||
async fn discover_databricks_models_with_token_source(
|
||||
cfg: &Config,
|
||||
token_source: Arc<dyn TokenSource>,
|
||||
) -> Result<Vec<ModelEntry>, AgentError> {
|
||||
discover_databricks_models_with_client(cfg, token_source, &Client::new()).await
|
||||
}
|
||||
|
||||
pub(crate) async fn discover_databricks_models_with_client(
|
||||
cfg: &Config,
|
||||
token_source: Arc<dyn TokenSource>,
|
||||
http: &Client,
|
||||
) -> Result<Vec<ModelEntry>, AgentError> {
|
||||
let mut bearer = token_source.bearer_no_browser().await?;
|
||||
let http = Client::new();
|
||||
let host = cfg.base_url.trim_end_matches('/');
|
||||
let mut refreshed = false;
|
||||
|
||||
loop {
|
||||
let result = match cfg.provider {
|
||||
Provider::Databricks => fetch_v1_models(&http, host, &bearer)
|
||||
Provider::Databricks => fetch_v1_models(http, host, &bearer)
|
||||
.await
|
||||
.map(|models| apply_model_filter(models, cfg.databricks_model_filter.as_ref())),
|
||||
Provider::DatabricksV2 => {
|
||||
fetch_v2_models(
|
||||
&http,
|
||||
http,
|
||||
host,
|
||||
&bearer,
|
||||
cfg.databricks_model_filter.as_ref(),
|
||||
|
||||
186
crates/buzz-agent/src/databricks.rs
Normal file
186
crates/buzz-agent/src/databricks.rs
Normal file
@@ -0,0 +1,186 @@
|
||||
//! Opt-in, workspace-scoped Databricks OAuth and catalog reuse.
|
||||
//!
|
||||
//! Legacy OAuth/catalog entry points retain their destination and intent policy.
|
||||
//! This entry requires an explicit HTTPS workspace, app-owned cache root and
|
||||
//! opener. Construction reads only that cache; it performs no network or browser
|
||||
//! work. The caller must not lend this root to legacy APIs or other applications.
|
||||
|
||||
use std::{path::Path, sync::Arc, time::Duration};
|
||||
|
||||
use reqwest::{redirect::Policy, Client, ClientBuilder};
|
||||
use url::Url;
|
||||
|
||||
use crate::{
|
||||
auth::{AuthError, AuthIntent, BrowserOpener, PkceOAuthTokenSource},
|
||||
catalog::{discover_databricks_models_with_client, ModelEntry},
|
||||
config::{Config, DatabricksModelFilter, Provider},
|
||||
AgentError,
|
||||
};
|
||||
|
||||
/// An explicitly scoped connection, not an agent runtime or a sign-in side effect.
|
||||
///
|
||||
/// Only [`connect`](Self::connect) permits browser opening. Discovery is always
|
||||
/// headless, including its one 401 refresh. Dropping either operation's future
|
||||
/// cancels its actual HTTP/callback work; callers should impose their own total
|
||||
/// deadline and fence stale results. Do not hold a controller lock while awaiting.
|
||||
/// Tokens remain inside the connection; never serialize this object to a webview.
|
||||
pub struct DatabricksConnection {
|
||||
workspace: StrictWorkspace,
|
||||
source: Arc<PkceOAuthTokenSource>,
|
||||
http: Client,
|
||||
}
|
||||
|
||||
impl DatabricksConnection {
|
||||
/// Construct without network/browser work or ambient credential/host lookup.
|
||||
///
|
||||
/// `workspace` must be an HTTPS origin (optional trailing slash), without
|
||||
/// userinfo, query or fragment. `cache_root` must be absolute and owned by
|
||||
/// this app; a separate `databricks-strict` namespace prevents accidental
|
||||
/// coalescing with legacy auth. Unix token files retain owner-only atomic
|
||||
/// persistence; on non-Unix the existing engine keeps credentials in memory.
|
||||
/// The opener must not log the authorization URL or its own error details.
|
||||
pub fn new(
|
||||
workspace: &str,
|
||||
cache_root: &Path,
|
||||
opener: Arc<dyn BrowserOpener>,
|
||||
) -> Result<Self, AgentError> {
|
||||
Self::build(workspace, cache_root, opener, Client::builder())
|
||||
}
|
||||
|
||||
fn build(
|
||||
workspace: &str,
|
||||
cache_root: &Path,
|
||||
opener: Arc<dyn BrowserOpener>,
|
||||
http: ClientBuilder,
|
||||
) -> Result<Self, AgentError> {
|
||||
let workspace = StrictWorkspace::new(workspace)?;
|
||||
if !cache_root.is_absolute() {
|
||||
return Err(AgentError::InvalidParams(
|
||||
"an absolute app-owned OAuth cache root is required".into(),
|
||||
));
|
||||
}
|
||||
// Apply policy AFTER the test transport's trust/DNS settings. No caller
|
||||
// can inject a client that follows redirects or disable production TLS.
|
||||
let http = http
|
||||
.redirect(Policy::none())
|
||||
.timeout(Duration::from_secs(30))
|
||||
.build()
|
||||
.map_err(|_| AgentError::Llm("Databricks HTTP client construction failed".into()))?;
|
||||
let source = PkceOAuthTokenSource::for_workspace(
|
||||
workspace.clone(),
|
||||
cache_root,
|
||||
opener,
|
||||
http.clone(),
|
||||
)
|
||||
.map_err(|_| AgentError::Llm("Databricks private OAuth cache unavailable".into()))?;
|
||||
Ok(Self {
|
||||
workspace,
|
||||
source,
|
||||
http,
|
||||
})
|
||||
}
|
||||
|
||||
/// Authenticate after an explicit user action. May call the supplied opener.
|
||||
/// A valid cached credential short-circuits; this does not fetch models.
|
||||
pub async fn connect(&self) -> Result<(), AuthError> {
|
||||
self.source
|
||||
.acquire_with_intent(AuthIntent::UserInitiated, None)
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
/// Discover v2 models without a browser, using this connection's workspace
|
||||
/// and credentials only. Preserves filtering, partial success and labelled
|
||||
/// authenticated-empty/no-filter fallback; this is not proof of completeness.
|
||||
/// Does not choose or mutate a model. Errors contain no raw provider content.
|
||||
pub async fn discover_models(
|
||||
&self,
|
||||
filter: Option<DatabricksModelFilter>,
|
||||
) -> Result<Vec<ModelEntry>, AgentError> {
|
||||
let cfg = Config::for_discovery(
|
||||
Provider::DatabricksV2,
|
||||
String::new(),
|
||||
self.workspace.as_str().into(),
|
||||
filter,
|
||||
);
|
||||
discover_databricks_models_with_client(&cfg, self.source.clone(), &self.http)
|
||||
.await
|
||||
.map_err(|error| match error {
|
||||
AgentError::LlmAuth(_) => {
|
||||
AgentError::LlmAuth("Databricks authentication required".into())
|
||||
}
|
||||
_ => AgentError::Llm("Databricks model discovery unavailable".into()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Unforgeable outside this module: every production value passed to auth was
|
||||
/// validated as an explicit HTTPS origin. Endpoint validation occurs inside the
|
||||
/// engine on the SAME discovery response subsequently used for both grants.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct StrictWorkspace(Url);
|
||||
|
||||
impl StrictWorkspace {
|
||||
fn new(raw: &str) -> Result<Self, AgentError> {
|
||||
let rest = raw.strip_prefix("https://").ok_or_else(invalid_workspace)?;
|
||||
let authority = rest.strip_suffix('/').unwrap_or(rest);
|
||||
if authority.contains(['/', '@']) {
|
||||
return Err(invalid_workspace());
|
||||
}
|
||||
let url = Url::parse(raw).map_err(|_| invalid_workspace())?;
|
||||
if raw.trim() != raw
|
||||
|| raw.chars().any(char::is_control)
|
||||
|| raw.contains('\\')
|
||||
|| url.scheme() != "https"
|
||||
|| url.host_str().is_none()
|
||||
|| !url.username().is_empty()
|
||||
|| url.password().is_some()
|
||||
|| url.query().is_some()
|
||||
|| url.fragment().is_some()
|
||||
|| url.path() != "/"
|
||||
{
|
||||
return Err(invalid_workspace());
|
||||
}
|
||||
Ok(Self(url))
|
||||
}
|
||||
|
||||
pub(crate) fn as_str(&self) -> &str {
|
||||
self.0.as_str().trim_end_matches('/')
|
||||
}
|
||||
|
||||
pub(crate) fn validate_endpoint(&self, raw: &str) -> Result<(), AgentError> {
|
||||
let authority = raw
|
||||
.strip_prefix("https://")
|
||||
.ok_or_else(invalid_endpoint)?
|
||||
.split('/')
|
||||
.next()
|
||||
.ok_or_else(invalid_endpoint)?;
|
||||
if authority.contains('@') {
|
||||
return Err(invalid_endpoint());
|
||||
}
|
||||
let endpoint = Url::parse(raw).map_err(|_| invalid_endpoint())?;
|
||||
if endpoint.origin() != self.0.origin()
|
||||
|| !endpoint.username().is_empty()
|
||||
|| endpoint.password().is_some()
|
||||
|| endpoint.query().is_some()
|
||||
|| endpoint.fragment().is_some()
|
||||
|| raw.trim() != raw
|
||||
|| raw.chars().any(char::is_control)
|
||||
|| raw.contains('\\')
|
||||
{
|
||||
return Err(invalid_endpoint());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn invalid_workspace() -> AgentError {
|
||||
AgentError::InvalidParams("Databricks workspace must be an explicit HTTPS origin".into())
|
||||
}
|
||||
fn invalid_endpoint() -> AgentError {
|
||||
AgentError::Llm("Databricks OAuth endpoint outside workspace policy".into())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "databricks_tests.rs"]
|
||||
mod tests;
|
||||
284
crates/buzz-agent/src/databricks_test_support.rs
Normal file
284
crates/buzz-agent/src/databricks_test_support.rs
Normal file
@@ -0,0 +1,284 @@
|
||||
//! Local-only synthetic OAuth/catalog transport. No browser, ambient cache or credentials.
|
||||
use super::*;
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
Mutex,
|
||||
},
|
||||
};
|
||||
use tokio::{
|
||||
io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
|
||||
task::JoinSet,
|
||||
};
|
||||
use tokio_rustls::{rustls, TlsAcceptor};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(super) struct Request {
|
||||
pub path: String,
|
||||
pub body: String,
|
||||
pub authorization: Option<String>,
|
||||
}
|
||||
impl Request {
|
||||
pub fn form(&self) -> HashMap<String, String> {
|
||||
url::form_urlencoded::parse(self.body.as_bytes())
|
||||
.into_owned()
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
pub(super) enum Reply {
|
||||
Json(u16, serde_json::Value),
|
||||
#[cfg(unix)]
|
||||
Raw(String),
|
||||
#[cfg(unix)]
|
||||
Stall,
|
||||
}
|
||||
pub(super) struct Server {
|
||||
pub base: String,
|
||||
pub cert: Option<reqwest::Certificate>,
|
||||
pub requests: Arc<Mutex<Vec<Request>>>,
|
||||
#[cfg_attr(not(unix), allow(dead_code))]
|
||||
pub active: Arc<AtomicUsize>,
|
||||
task: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
impl Drop for Server {
|
||||
fn drop(&mut self) {
|
||||
self.task.abort();
|
||||
}
|
||||
}
|
||||
impl Server {
|
||||
pub async fn start(
|
||||
tls: bool,
|
||||
handler: impl Fn(&str, &Request) -> Reply + Send + Sync + 'static,
|
||||
) -> Self {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base = format!(
|
||||
"{}://localhost:{}",
|
||||
if tls { "https" } else { "http" },
|
||||
listener.local_addr().unwrap().port()
|
||||
);
|
||||
let (acceptor, cert) = if tls {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
|
||||
let der = cert.cert.der().clone();
|
||||
let config = rustls::ServerConfig::builder_with_provider(Arc::new(
|
||||
rustls::crypto::ring::default_provider(),
|
||||
))
|
||||
.with_safe_default_protocol_versions()
|
||||
.unwrap()
|
||||
.with_no_client_auth()
|
||||
.with_single_cert(
|
||||
vec![der.clone()],
|
||||
rustls::pki_types::PrivatePkcs8KeyDer::from(cert.signing_key.serialize_der())
|
||||
.into(),
|
||||
)
|
||||
.unwrap();
|
||||
(
|
||||
Some(TlsAcceptor::from(Arc::new(config))),
|
||||
Some(reqwest::Certificate::from_der(der.as_ref()).unwrap()),
|
||||
)
|
||||
} else {
|
||||
(None, None)
|
||||
};
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let active = Arc::new(AtomicUsize::new(0));
|
||||
let handler = Arc::new(handler);
|
||||
let (base2, requests2, active2) = (base.clone(), requests.clone(), active.clone());
|
||||
let task = tokio::spawn(async move {
|
||||
let mut children = JoinSet::new();
|
||||
loop {
|
||||
tokio::select! {
|
||||
accepted = listener.accept() => {
|
||||
let (stream, _) = accepted.unwrap();
|
||||
let (handler, base, requests, active, acceptor) = (handler.clone(), base2.clone(), requests2.clone(), active2.clone(), acceptor.clone());
|
||||
children.spawn(async move {
|
||||
if let Some(acceptor) = acceptor {
|
||||
if let Ok(stream) = acceptor.accept(stream).await { serve(stream, base, handler, requests, active).await; }
|
||||
} else { serve(stream, base, handler, requests, active).await; }
|
||||
});
|
||||
}
|
||||
_ = children.join_next(), if !children.is_empty() => {}
|
||||
}
|
||||
}
|
||||
});
|
||||
Self {
|
||||
base,
|
||||
cert,
|
||||
requests,
|
||||
active,
|
||||
task,
|
||||
}
|
||||
}
|
||||
pub fn builder(&self) -> ClientBuilder {
|
||||
let builder = Client::builder().no_proxy();
|
||||
match &self.cert {
|
||||
Some(cert) => builder.add_root_certificate(cert.clone()),
|
||||
None => builder,
|
||||
}
|
||||
}
|
||||
pub fn count(&self, path: &str) -> usize {
|
||||
self.requests
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.filter(|r| r.path.starts_with(path))
|
||||
.count()
|
||||
}
|
||||
#[cfg(unix)]
|
||||
pub async fn wait_for(&self, path: &str) {
|
||||
tokio::time::timeout(Duration::from_secs(3), async {
|
||||
while self.count(path) == 0 {
|
||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
struct Active(Arc<AtomicUsize>);
|
||||
impl Drop for Active {
|
||||
fn drop(&mut self) {
|
||||
self.0.fetch_sub(1, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
async fn serve<S: AsyncRead + AsyncWrite + Unpin>(
|
||||
mut stream: S,
|
||||
base: String,
|
||||
handler: Arc<impl Fn(&str, &Request) -> Reply>,
|
||||
requests: Arc<Mutex<Vec<Request>>>,
|
||||
active: Arc<AtomicUsize>,
|
||||
) {
|
||||
active.fetch_add(1, Ordering::SeqCst);
|
||||
let _guard = Active(active);
|
||||
let mut bytes = Vec::new();
|
||||
let mut buf = [0; 4096];
|
||||
let header_end = loop {
|
||||
let Ok(n) = stream.read(&mut buf).await else {
|
||||
return;
|
||||
};
|
||||
if n == 0 {
|
||||
return;
|
||||
}
|
||||
bytes.extend_from_slice(&buf[..n]);
|
||||
if let Some(pos) = bytes.windows(4).position(|w| w == b"\r\n\r\n") {
|
||||
break pos + 4;
|
||||
}
|
||||
assert!(bytes.len() < 32 * 1024);
|
||||
};
|
||||
let header = String::from_utf8(bytes[..header_end].to_vec()).unwrap();
|
||||
let length = header
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
line.to_ascii_lowercase()
|
||||
.strip_prefix("content-length: ")
|
||||
.map(|v| v.parse::<usize>().unwrap())
|
||||
})
|
||||
.unwrap_or(0);
|
||||
while bytes.len() < header_end + length {
|
||||
let n = stream.read(&mut buf).await.unwrap();
|
||||
if n == 0 {
|
||||
return;
|
||||
}
|
||||
bytes.extend_from_slice(&buf[..n]);
|
||||
}
|
||||
let request = Request {
|
||||
path: header
|
||||
.lines()
|
||||
.next()
|
||||
.unwrap()
|
||||
.split_whitespace()
|
||||
.nth(1)
|
||||
.unwrap()
|
||||
.into(),
|
||||
body: String::from_utf8(bytes[header_end..header_end + length].to_vec()).unwrap(),
|
||||
authorization: header.lines().find_map(|line| {
|
||||
line.split_once(':')
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
.map(|(_, value)| value.trim().into())
|
||||
}),
|
||||
};
|
||||
requests.lock().unwrap().push(request.clone());
|
||||
let response = match handler(&base, &request) {
|
||||
Reply::Json(status, body) => {
|
||||
let body = body.to_string();
|
||||
format!("HTTP/1.1 {status} Test\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len())
|
||||
}
|
||||
#[cfg(unix)]
|
||||
Reply::Raw(raw) => raw,
|
||||
#[cfg(unix)]
|
||||
Reply::Stall => {
|
||||
// Send headers, then stall the body until the client cancels.
|
||||
let _ = stream
|
||||
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 1000000\r\n\r\n")
|
||||
.await;
|
||||
let _ = stream.flush().await;
|
||||
let _ = stream.read(&mut buf).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
let _ = stream.write_all(response.as_bytes()).await;
|
||||
let _ = stream.shutdown().await;
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct Opener {
|
||||
pub urls: Mutex<Vec<String>>,
|
||||
pub fail: bool,
|
||||
pub callback: bool,
|
||||
pub denial: bool,
|
||||
}
|
||||
impl BrowserOpener for Opener {
|
||||
fn open(&self, raw: &str) -> Result<(), String> {
|
||||
self.urls.lock().unwrap().push(raw.into());
|
||||
if self.fail {
|
||||
return Err(format!("secret-opener-error {raw}"));
|
||||
}
|
||||
if self.callback {
|
||||
let url = Url::parse(raw).unwrap();
|
||||
let params: HashMap<_, _> = url.query_pairs().into_owned().collect();
|
||||
assert_eq!(params["code_challenge_method"], "S256");
|
||||
assert_eq!(params["client_id"], "databricks-cli");
|
||||
let mut redirect = Url::parse(¶ms["redirect_uri"]).unwrap();
|
||||
if self.denial {
|
||||
redirect
|
||||
.query_pairs_mut()
|
||||
.append_pair("error", "secret-denial");
|
||||
} else {
|
||||
redirect
|
||||
.query_pairs_mut()
|
||||
.append_pair("code", "synthetic-code")
|
||||
.append_pair("state", ¶ms["state"]);
|
||||
}
|
||||
tokio::spawn(async move {
|
||||
let _ = Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.unwrap()
|
||||
.get(redirect)
|
||||
.send()
|
||||
.await;
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
pub(super) fn discovery(base: &str) -> Reply {
|
||||
Reply::Json(
|
||||
200,
|
||||
serde_json::json!({"authorization_endpoint": format!("{base}/authorize"), "token_endpoint": format!("{base}/token")}),
|
||||
)
|
||||
}
|
||||
#[cfg(unix)]
|
||||
pub(super) fn seed(root: &Path, base: &str, namespace: &str, token: &str, expired: bool) {
|
||||
use sha2::{Digest, Sha256};
|
||||
let key = format!(
|
||||
"{base}/oidc/.well-known/oauth-authorization-server|databricks-cli|all-apis,offline_access"
|
||||
);
|
||||
let dir = root.join(namespace);
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
std::fs::write(dir.join(format!("{}.json", hex::encode(Sha256::digest(key)))), serde_json::json!({"access_token": token, "refresh_token": "synthetic-refresh", "expires_at": if expired { 1 } else { 4_000_000_000u64 }}).to_string()).unwrap();
|
||||
}
|
||||
#[cfg(unix)]
|
||||
pub(super) fn chunked(status: u16, body: &str) -> Reply {
|
||||
Reply::Raw(format!("HTTP/1.1 {status} Test\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n{:x}\r\n{body}\r\n0\r\n\r\n", body.len()))
|
||||
}
|
||||
817
crates/buzz-agent/src/databricks_tests.rs
Normal file
817
crates/buzz-agent/src/databricks_tests.rs
Normal file
@@ -0,0 +1,817 @@
|
||||
use super::*;
|
||||
#[cfg(unix)]
|
||||
use crate::auth::TokenSource;
|
||||
use serde_json::json;
|
||||
#[cfg(unix)]
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
#[path = "databricks_test_support.rs"]
|
||||
mod support;
|
||||
use support::*;
|
||||
|
||||
fn strict(server: &Server, root: &Path, opener: Arc<Opener>) -> DatabricksConnection {
|
||||
DatabricksConnection::build(&server.base, root, opener, server.builder()).unwrap()
|
||||
}
|
||||
#[cfg(unix)]
|
||||
fn legacy(server: &Server, root: &Path, opener: Arc<Opener>) -> Arc<PkceOAuthTokenSource> {
|
||||
PkceOAuthTokenSource::new_with(
|
||||
crate::llm::databricks_pkce_config(&server.base, Some(root.into())),
|
||||
opener,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn public_constructor_requires_explicit_https_origin_and_root() {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
for host in [
|
||||
"",
|
||||
"localhost",
|
||||
"http://localhost",
|
||||
"https://user:pass@localhost",
|
||||
"https://localhost/path",
|
||||
"https://localhost?secret",
|
||||
"https://localhost#secret",
|
||||
" https://localhost",
|
||||
"https://local\nhost",
|
||||
"https://localhost\\evil",
|
||||
"https:localhost",
|
||||
"https://@localhost",
|
||||
"https://localhost/path/..",
|
||||
] {
|
||||
assert!(
|
||||
DatabricksConnection::new(host, root.path(), Arc::new(Opener::default())).is_err(),
|
||||
"{host:?}"
|
||||
);
|
||||
}
|
||||
assert!(DatabricksConnection::new(
|
||||
"https://workspace.invalid",
|
||||
Path::new("relative"),
|
||||
Arc::new(Opener::default())
|
||||
)
|
||||
.is_err());
|
||||
let connection = DatabricksConnection::new(
|
||||
"https://WORKSPACE.invalid:443/",
|
||||
root.path(),
|
||||
Arc::new(Opener::default()),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(connection.workspace.as_str(), "https://workspace.invalid");
|
||||
assert!(root.path().join("databricks-strict").exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn strict_headless_empty_never_networks_or_browses() {
|
||||
let server = Server::start(true, |_, _| panic!("unexpected network")).await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let opener = Arc::new(Opener::default());
|
||||
let c = strict(&server, root.path(), opener.clone());
|
||||
assert!(matches!(
|
||||
c.discover_models(None).await,
|
||||
Err(AgentError::LlmAuth(_))
|
||||
));
|
||||
assert!(server.requests.lock().unwrap().is_empty());
|
||||
assert!(opener.urls.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn strict_rejects_discovered_endpoints_before_grants_or_browser() {
|
||||
let sink = Server::start(true, |_, _| {
|
||||
Reply::Json(200, json!({"access_token":"wrong"}))
|
||||
})
|
||||
.await;
|
||||
// Same host/different port, different host, insecure URL, userinfo, query,
|
||||
// fragment and malformed URLs must all fail before either grant or opener.
|
||||
for endpoint in [
|
||||
format!("{}/token", sink.base),
|
||||
"https://other.invalid/token".into(),
|
||||
"http://localhost/token".into(),
|
||||
"https://user:secret@localhost/token".into(),
|
||||
"not a URL".into(),
|
||||
] {
|
||||
for field in ["authorization_endpoint", "token_endpoint"] {
|
||||
let endpoint = endpoint.clone();
|
||||
let server = Server::start(true, move |base, _| {
|
||||
let mut body = json!({"authorization_endpoint":format!("{base}/authorize"),"token_endpoint":format!("{base}/token")});
|
||||
body[field] = json!(endpoint);
|
||||
Reply::Json(200, body)
|
||||
}).await;
|
||||
for expired in [false, true] {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
if expired {
|
||||
seed(root.path(), &server.base, "databricks-strict", "old", true);
|
||||
}
|
||||
let opener = Arc::new(Opener {
|
||||
fail: true,
|
||||
..Default::default()
|
||||
});
|
||||
let c = strict(&server, root.path(), opener.clone());
|
||||
assert_eq!(c.connect().await, Err(AuthError::NetworkUnavailable));
|
||||
assert_eq!(server.count("/token"), 0);
|
||||
assert!(opener.urls.lock().unwrap().is_empty());
|
||||
}
|
||||
}
|
||||
}
|
||||
assert!(sink.requests.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn strict_rejects_same_origin_endpoint_credentials_query_fragment() {
|
||||
for suffix in [
|
||||
"/token?secret=echo",
|
||||
"/token#fragment",
|
||||
"/token\n",
|
||||
"\\token",
|
||||
] {
|
||||
for field in ["authorization_endpoint", "token_endpoint"] {
|
||||
let server = Server::start(true, move |base, _| {
|
||||
let mut body = json!({"authorization_endpoint":format!("{base}/authorize"),"token_endpoint":format!("{base}/token")});
|
||||
body[field] = json!(format!("{base}{suffix}"));
|
||||
Reply::Json(200, body)
|
||||
}).await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let c = strict(
|
||||
&server,
|
||||
root.path(),
|
||||
Arc::new(Opener {
|
||||
fail: true,
|
||||
..Default::default()
|
||||
}),
|
||||
);
|
||||
assert_eq!(c.connect().await, Err(AuthError::NetworkUnavailable));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn strict_connect_code_exchange_catalog_and_headless_401_refresh() {
|
||||
let server = Server::start(true, |base, request| match request.path.as_str() {
|
||||
"/oidc/.well-known/oauth-authorization-server" => discovery(base),
|
||||
"/token" => {
|
||||
let form = request.form();
|
||||
assert_eq!(form["client_id"], "databricks-cli");
|
||||
let token = if form["grant_type"] == "authorization_code" {
|
||||
assert_eq!(form["code"], "synthetic-code");
|
||||
assert!(form["code_verifier"].len() >= 43);
|
||||
"rejected"
|
||||
} else {
|
||||
assert_eq!(form["refresh_token"], "synthetic-refresh");
|
||||
"fresh"
|
||||
};
|
||||
Reply::Json(
|
||||
200,
|
||||
json!({"access_token":token,"refresh_token":"synthetic-refresh","expires_in":3600}),
|
||||
)
|
||||
}
|
||||
path if path.starts_with("/api/") => {
|
||||
if request.authorization.as_deref() == Some("Bearer rejected") {
|
||||
return Reply::Json(401, json!({"secret":"echo"}));
|
||||
}
|
||||
assert_eq!(request.authorization.as_deref(), Some("Bearer fresh"));
|
||||
if path.starts_with("/api/ai-gateway") {
|
||||
Reply::Json(200, json!({"endpoints":[{"name":"custom-chat"}]}))
|
||||
} else {
|
||||
Reply::Json(200, json!({"model_services":[]}))
|
||||
}
|
||||
}
|
||||
_ => panic!("unexpected request {request:?}"),
|
||||
})
|
||||
.await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let opener = Arc::new(Opener {
|
||||
callback: true,
|
||||
..Default::default()
|
||||
});
|
||||
let c = strict(&server, root.path(), opener.clone());
|
||||
c.connect().await.unwrap();
|
||||
assert_eq!(server.count("/api/"), 0, "connect must not discover models");
|
||||
let models = c.discover_models(None).await.unwrap();
|
||||
assert_eq!(
|
||||
models,
|
||||
vec![ModelEntry {
|
||||
id: "custom-chat".into(),
|
||||
name: "custom-chat".into()
|
||||
}]
|
||||
);
|
||||
assert_eq!(server.count("/token"), 2);
|
||||
assert_eq!(opener.urls.lock().unwrap().len(), 1);
|
||||
c.connect().await.unwrap();
|
||||
assert_eq!(
|
||||
opener.urls.lock().unwrap().len(),
|
||||
1,
|
||||
"fresh cache must not reopen browser"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn strict_never_follows_discovery_token_or_catalog_redirects() {
|
||||
let sink = Server::start(true, |_, _| {
|
||||
Reply::Json(200, json!({"access_token":"escaped"}))
|
||||
})
|
||||
.await;
|
||||
for status in [301, 302, 303, 307, 308] {
|
||||
for stage in ["discovery", "refresh", "exchange", "catalog"] {
|
||||
let location = format!("{}/sink", sink.base);
|
||||
let server = Server::start(true, move |base, r| {
|
||||
let redirect = match stage {
|
||||
"discovery" => true,
|
||||
"refresh" | "exchange" => r.path == "/token",
|
||||
_ => r.path.starts_with("/api/"),
|
||||
};
|
||||
if redirect { return Reply::Raw(format!("HTTP/1.1 {status} Redirect\r\nLocation: {location}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")); }
|
||||
discovery(base)
|
||||
}).await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
if stage == "refresh" || stage == "catalog" {
|
||||
seed(
|
||||
root.path(),
|
||||
&server.base,
|
||||
"databricks-strict",
|
||||
"synthetic-access",
|
||||
stage == "refresh",
|
||||
);
|
||||
}
|
||||
let opener = Arc::new(Opener {
|
||||
callback: true,
|
||||
..Default::default()
|
||||
});
|
||||
// Trust the second server too: lack of trust cannot falsely prove
|
||||
// redirect containment. The no-redirect policy must do the work.
|
||||
let c = DatabricksConnection::build(
|
||||
&server.base,
|
||||
root.path(),
|
||||
opener.clone(),
|
||||
server
|
||||
.builder()
|
||||
.add_root_certificate(sink.cert.clone().unwrap()),
|
||||
)
|
||||
.unwrap();
|
||||
if stage == "catalog" {
|
||||
assert!(matches!(
|
||||
c.discover_models(None).await,
|
||||
Err(AgentError::Llm(_))
|
||||
));
|
||||
} else {
|
||||
assert_eq!(c.connect().await, Err(AuthError::NetworkUnavailable));
|
||||
}
|
||||
assert_eq!(
|
||||
opener.urls.lock().unwrap().len(),
|
||||
usize::from(stage == "exchange")
|
||||
);
|
||||
assert_eq!(sink.count("/sink"), 0, "redirect escaped: {stage}/{status}");
|
||||
}
|
||||
}
|
||||
// Positive control: the sink is reachable with the same trust settings.
|
||||
sink.builder()
|
||||
.build()
|
||||
.unwrap()
|
||||
.get(format!("{}/control", sink.base))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(sink.count("/control"), 1);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn strict_catalog_preserves_empty_filter_partial_and_paging_semantics() {
|
||||
for scenario in ["empty", "filtered-empty", "partial", "pages"] {
|
||||
let server = Server::start(true, move |_, r| {
|
||||
if scenario == "partial" && r.path.starts_with("/api/2.1/") { return Reply::Json(403, json!({"message":"unavailable"})); }
|
||||
if r.path.starts_with("/api/2.1/") { return Reply::Json(200, json!({"model_services":[]})); }
|
||||
assert_eq!(r.authorization.as_deref(), Some("Bearer synthetic-access"));
|
||||
match scenario {
|
||||
"empty" | "filtered-empty" => Reply::Json(200, json!({"endpoints":[]})),
|
||||
"pages" if !r.path.contains("page_token=") => Reply::Json(200, json!({"endpoints":[{"name":"one"}],"next_page_token":"https://other.invalid/?secret"})),
|
||||
_ => Reply::Json(200, json!({"endpoints":[{"name":"two"}]})),
|
||||
}
|
||||
}).await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
seed(
|
||||
root.path(),
|
||||
&server.base,
|
||||
"databricks-strict",
|
||||
"synthetic-access",
|
||||
false,
|
||||
);
|
||||
let c = strict(&server, root.path(), Arc::new(Opener::default()));
|
||||
let filter = if scenario == "filtered-empty" {
|
||||
DatabricksModelFilter::parse(Some("no-match")).unwrap()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let models = c.discover_models(filter).await.unwrap();
|
||||
match scenario {
|
||||
"empty" => assert!(
|
||||
models.len() > 1
|
||||
&& models
|
||||
.iter()
|
||||
.all(|m| m.name.ends_with(" (default catalog)"))
|
||||
),
|
||||
"filtered-empty" => assert!(models.is_empty()),
|
||||
"partial" => assert_eq!(
|
||||
models.iter().map(|m| m.id.as_str()).collect::<Vec<_>>(),
|
||||
["two"]
|
||||
),
|
||||
_ => {
|
||||
assert_eq!(models.len(), 2);
|
||||
assert_eq!(server.count("/api/ai-gateway"), 2);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn strict_cancellation_drops_real_http_work_and_allows_retry() {
|
||||
for stage in ["discovery", "refresh", "exchange", "catalog"] {
|
||||
let stalled = Arc::new(std::sync::atomic::AtomicBool::new(true));
|
||||
let stalled2 = stalled.clone();
|
||||
let server = Server::start(true, move |base, r| {
|
||||
let target = match stage { "discovery" => r.path.contains("well-known"), "refresh"|"exchange" => r.path == "/token", _ => r.path.starts_with("/api/") };
|
||||
if target && stalled2.load(Ordering::SeqCst) { return Reply::Stall; }
|
||||
if r.path.contains("well-known") { discovery(base) }
|
||||
else if r.path == "/token" { Reply::Json(200, json!({"access_token":"fresh", "refresh_token":"synthetic-refresh","expires_in":3600})) }
|
||||
else if r.path.starts_with("/api/ai-gateway") { Reply::Json(200, json!({"endpoints":[]})) }
|
||||
else { Reply::Json(200, json!({"model_services":[]})) }
|
||||
}).await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
if stage == "refresh" || stage == "catalog" {
|
||||
seed(
|
||||
root.path(),
|
||||
&server.base,
|
||||
"databricks-strict",
|
||||
"old",
|
||||
stage == "refresh",
|
||||
);
|
||||
}
|
||||
let opener = Arc::new(Opener {
|
||||
callback: true,
|
||||
..Default::default()
|
||||
});
|
||||
let c = Arc::new(strict(&server, root.path(), opener.clone()));
|
||||
let c2 = c.clone();
|
||||
let task = tokio::spawn(async move {
|
||||
if stage == "catalog" {
|
||||
c2.discover_models(None).await.map(|_| ())
|
||||
} else {
|
||||
c2.connect().await.map_err(AgentError::from)
|
||||
}
|
||||
});
|
||||
server
|
||||
.wait_for(match stage {
|
||||
"discovery" => "/oidc/",
|
||||
"refresh" | "exchange" => "/token",
|
||||
_ => "/api/",
|
||||
})
|
||||
.await;
|
||||
task.abort();
|
||||
assert!(task.await.unwrap_err().is_cancelled());
|
||||
tokio::time::timeout(Duration::from_secs(3), async {
|
||||
while server.active.load(Ordering::SeqCst) != 0 {
|
||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
stalled.store(false, Ordering::SeqCst);
|
||||
tokio::time::timeout(Duration::from_secs(3), async {
|
||||
if stage == "catalog" {
|
||||
c.discover_models(None).await.unwrap();
|
||||
} else {
|
||||
c.connect().await.unwrap();
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn strict_browser_wait_cancellation_releases_callback_listener() {
|
||||
let server = Server::start(true, |base, _| discovery(base)).await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let opener = Arc::new(Opener::default());
|
||||
let c = Arc::new(strict(&server, root.path(), opener.clone()));
|
||||
let c2 = c.clone();
|
||||
let task = tokio::spawn(async move { c2.connect().await });
|
||||
tokio::time::timeout(Duration::from_secs(3), async {
|
||||
while opener.urls.lock().unwrap().is_empty() {
|
||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let url = Url::parse(&opener.urls.lock().unwrap()[0]).unwrap();
|
||||
let callback = url
|
||||
.query_pairs()
|
||||
.find(|(k, _)| k == "redirect_uri")
|
||||
.unwrap()
|
||||
.1
|
||||
.into_owned();
|
||||
let client = Client::builder().no_proxy().build().unwrap();
|
||||
// Connect-only positive control leaves the one-shot callback unconsumed.
|
||||
let port = Url::parse(&callback).unwrap().port().unwrap();
|
||||
let socket = tokio::net::TcpStream::connect(("127.0.0.1", port))
|
||||
.await
|
||||
.unwrap();
|
||||
drop(socket);
|
||||
task.abort();
|
||||
assert!(task.await.unwrap_err().is_cancelled());
|
||||
tokio::time::timeout(Duration::from_secs(3), async {
|
||||
loop {
|
||||
if client.get(&callback).send().await.is_err() {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn legacy_different_origin_endpoints_and_redirects_remain_supported() {
|
||||
let sink = Server::start(false, |_, r| {
|
||||
assert_eq!(r.form()["refresh_token"], "synthetic-refresh");
|
||||
Reply::Json(
|
||||
200,
|
||||
json!({"access_token":"legacy-fresh","expires_in":3600}),
|
||||
)
|
||||
})
|
||||
.await;
|
||||
for redirected in [false, true] {
|
||||
let endpoint = format!("{}/token", sink.base);
|
||||
let server = Server::start(false, move |base, r| {
|
||||
if r.path == "/token" { Reply::Raw(format!("HTTP/1.1 307 Redirect\r\nLocation: {endpoint}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")) }
|
||||
else { Reply::Json(200, json!({"authorization_endpoint":format!("{base}/authorize"),"token_endpoint":if redirected {format!("{base}/token")} else {endpoint.clone()}})) }
|
||||
}).await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
seed(root.path(), &server.base, "databricks", "old", true);
|
||||
assert_eq!(
|
||||
legacy(&server, root.path(), Arc::new(Opener::default()))
|
||||
.bearer_no_browser()
|
||||
.await
|
||||
.unwrap(),
|
||||
"legacy-fresh"
|
||||
);
|
||||
}
|
||||
assert_eq!(sink.count("/token"), 2);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn strict_host_root_and_legacy_namespace_isolation() {
|
||||
let server = Server::start(true, |_, _| panic!("should not reach network")).await;
|
||||
let a = tempfile::tempdir().unwrap();
|
||||
let b = tempfile::tempdir().unwrap();
|
||||
seed(a.path(), &server.base, "databricks", "legacy", false);
|
||||
let c = strict(&server, a.path(), Arc::new(Opener::default()));
|
||||
assert!(matches!(
|
||||
c.discover_models(None).await,
|
||||
Err(AgentError::LlmAuth(_))
|
||||
));
|
||||
seed(
|
||||
a.path(),
|
||||
&server.base,
|
||||
"databricks-strict",
|
||||
"strict-a",
|
||||
false,
|
||||
);
|
||||
let c = strict(&server, b.path(), Arc::new(Opener::default()));
|
||||
assert_eq!(
|
||||
c.source.bearer_no_browser().await.unwrap_err().to_string(),
|
||||
AgentError::from(AuthError::NoCredential).to_string()
|
||||
);
|
||||
let c = DatabricksConnection::new(
|
||||
"https://different.invalid",
|
||||
a.path(),
|
||||
Arc::new(Opener::default()),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
c.source.bearer_no_browser().await,
|
||||
Err(AgentError::LlmAuth(_))
|
||||
));
|
||||
let c = strict(&server, a.path(), Arc::new(Opener::default()));
|
||||
assert_eq!(c.source.bearer_no_browser().await.unwrap(), "strict-a");
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn shared_oauth_limits_bound_chunked_and_declared_bodies_in_both_paths() {
|
||||
use crate::auth_http::{MAX_OAUTH_ERROR_BYTES, MAX_OAUTH_RESPONSE_BYTES};
|
||||
for strict_path in [false, true] {
|
||||
for stage in ["discovery", "refresh", "exchange"] {
|
||||
for (status, limit) in [
|
||||
(200, MAX_OAUTH_RESPONSE_BYTES),
|
||||
(400, MAX_OAUTH_ERROR_BYTES),
|
||||
] {
|
||||
for declared in [false, true] {
|
||||
let server = Server::start(strict_path, move |base, request| {
|
||||
let target = if stage == "discovery" { request.path.contains("well-known") } else { request.path == "/token" };
|
||||
if !target { return discovery(base); }
|
||||
if declared {
|
||||
// Headers alone exceed the bound; never wait on the body.
|
||||
Reply::Raw(format!("HTTP/1.1 {status} Test\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", limit+1))
|
||||
} else {
|
||||
// Valid JSON + padding means a decoder without the cap
|
||||
// would succeed (or misclassify an oversized invalid_grant).
|
||||
let value = if status == 400 { json!({"error":"invalid_grant"}) }
|
||||
else if stage == "discovery" { json!({"authorization_endpoint":format!("{base}/authorize"),"token_endpoint":format!("{base}/token")}) }
|
||||
else { json!({"access_token":"must-not-cache","expires_in":3600}) };
|
||||
let mut body = value.to_string();
|
||||
body.push_str(&" ".repeat(limit+1-body.len()));
|
||||
chunked(status, &body)
|
||||
}
|
||||
}).await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
if stage == "refresh" {
|
||||
seed(
|
||||
root.path(),
|
||||
&server.base,
|
||||
if strict_path {
|
||||
"databricks-strict"
|
||||
} else {
|
||||
"databricks"
|
||||
},
|
||||
"old",
|
||||
true,
|
||||
);
|
||||
}
|
||||
let opener = Arc::new(Opener {
|
||||
callback: true,
|
||||
..Default::default()
|
||||
});
|
||||
let result = tokio::time::timeout(Duration::from_secs(3), async {
|
||||
if strict_path {
|
||||
strict(&server, root.path(), opener.clone()).connect().await
|
||||
} else {
|
||||
legacy(&server, root.path(), opener.clone())
|
||||
.acquire_with_intent(AuthIntent::UserInitiated, None)
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
result,
|
||||
Err(AuthError::NetworkUnavailable),
|
||||
"strict={strict_path} {stage} {status} declared={declared}"
|
||||
);
|
||||
assert_eq!(
|
||||
opener.urls.lock().unwrap().len(),
|
||||
usize::from(stage == "exchange")
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn shared_exact_limit_responses_and_grant_classifications_are_preserved() {
|
||||
use crate::auth_http::{MAX_OAUTH_ERROR_BYTES, MAX_OAUTH_RESPONSE_BYTES};
|
||||
for strict_path in [false, true] {
|
||||
for stage in ["refresh", "exchange"] {
|
||||
for (status, error, expected) in [
|
||||
(200, "", Ok(())),
|
||||
(
|
||||
400,
|
||||
"invalid_grant",
|
||||
Err(if stage == "refresh" {
|
||||
AuthError::RefreshRejected
|
||||
} else {
|
||||
AuthError::ExchangeFailed
|
||||
}),
|
||||
),
|
||||
(400, "invalid_client", Err(AuthError::NetworkUnavailable)),
|
||||
(429, "invalid_request", Err(AuthError::NetworkUnavailable)),
|
||||
(503, "invalid_grant", Err(AuthError::NetworkUnavailable)),
|
||||
] {
|
||||
let server = Server::start(strict_path, move |base, request| {
|
||||
let mut body = if request.path.contains("well-known") { json!({"authorization_endpoint":format!("{base}/authorize"),"token_endpoint":format!("{base}/token")}).to_string() }
|
||||
else if status == 200 { json!({"access_token":"valid","refresh_token":"synthetic-refresh","expires_in":3600}).to_string() }
|
||||
else { json!({"error":error, "error_description":"secret-echo"}).to_string() };
|
||||
let status = if request.path.contains("well-known") { 200 } else { status };
|
||||
let limit = if status == 200 { MAX_OAUTH_RESPONSE_BYTES } else { MAX_OAUTH_ERROR_BYTES };
|
||||
body.push_str(&" ".repeat(limit-body.len()));
|
||||
chunked(status, &body)
|
||||
}).await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
if stage == "refresh" {
|
||||
seed(
|
||||
root.path(),
|
||||
&server.base,
|
||||
if strict_path {
|
||||
"databricks-strict"
|
||||
} else {
|
||||
"databricks"
|
||||
},
|
||||
"old",
|
||||
true,
|
||||
);
|
||||
}
|
||||
let opener = Arc::new(Opener {
|
||||
callback: true,
|
||||
..Default::default()
|
||||
});
|
||||
let c;
|
||||
let source = if strict_path {
|
||||
c = strict(&server, root.path(), opener.clone());
|
||||
c.source.clone()
|
||||
} else {
|
||||
legacy(&server, root.path(), opener.clone())
|
||||
};
|
||||
let intent = if stage == "refresh" {
|
||||
AuthIntent::Headless
|
||||
} else {
|
||||
AuthIntent::UserInitiated
|
||||
};
|
||||
assert_eq!(
|
||||
source.acquire_with_intent(intent, None).await.map(|_| ()),
|
||||
expected
|
||||
);
|
||||
assert_eq!(
|
||||
opener.urls.lock().unwrap().len(),
|
||||
usize::from(stage == "exchange")
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn shared_diagnostics_never_log_oauth_bodies_urls_opener_or_callback_details() {
|
||||
use std::{io::Write, sync::Mutex};
|
||||
// Isolate tracing's global callsite interest cache from concurrently running
|
||||
// lib tests. The child runs ONLY this test, never an agent or real browser.
|
||||
const MARKER: &str = "BUZZ_TEST_OAUTH_LOG_CAPTURE_CHILD";
|
||||
if std::env::var_os(MARKER).is_none() {
|
||||
let output = tokio::time::timeout(Duration::from_secs(15),
|
||||
tokio::process::Command::new(std::env::current_exe().unwrap())
|
||||
.arg("--exact")
|
||||
.arg("databricks::tests::shared_diagnostics_never_log_oauth_bodies_urls_opener_or_callback_details")
|
||||
.arg("--nocapture")
|
||||
.env(MARKER, "1")
|
||||
.kill_on_drop(true)
|
||||
.output()).await.unwrap().unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{}\n{}",
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
return;
|
||||
}
|
||||
#[derive(Clone, Default)]
|
||||
struct Capture(Arc<Mutex<Vec<u8>>>);
|
||||
impl Write for Capture {
|
||||
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
|
||||
self.0.lock().unwrap().extend_from_slice(bytes);
|
||||
Ok(bytes.len())
|
||||
}
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for Capture {
|
||||
type Writer = Self;
|
||||
fn make_writer(&'a self) -> Self {
|
||||
self.clone()
|
||||
}
|
||||
}
|
||||
let logs = Capture::default();
|
||||
tracing::subscriber::set_global_default(
|
||||
tracing_subscriber::fmt()
|
||||
.with_ansi(false)
|
||||
.without_time()
|
||||
.with_writer(logs.clone())
|
||||
.finish(),
|
||||
)
|
||||
.unwrap();
|
||||
for strict_path in [false, true] {
|
||||
for scenario in [
|
||||
"refresh-rejected",
|
||||
"refresh-network",
|
||||
"exchange-rejected",
|
||||
"exchange-network",
|
||||
"opener",
|
||||
"denial",
|
||||
] {
|
||||
logs.0.lock().unwrap().clear();
|
||||
let server = Server::start(strict_path, move |base, request| {
|
||||
if request.path.contains("well-known") { return discovery(base); }
|
||||
let rejected = scenario.ends_with("rejected");
|
||||
Reply::Json(if rejected {400} else {503}, json!({"error":if rejected {"invalid_grant"} else {"server_error"},"error_description":format!("secret-provider-body {} {base}", request.body)}))
|
||||
}).await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
if scenario.starts_with("refresh") {
|
||||
seed(
|
||||
root.path(),
|
||||
&server.base,
|
||||
if strict_path {
|
||||
"databricks-strict"
|
||||
} else {
|
||||
"databricks"
|
||||
},
|
||||
"old",
|
||||
true,
|
||||
);
|
||||
}
|
||||
let opener = Arc::new(Opener {
|
||||
callback: true,
|
||||
fail: scenario == "opener",
|
||||
denial: scenario == "denial",
|
||||
..Default::default()
|
||||
});
|
||||
let source = if strict_path {
|
||||
strict(&server, root.path(), opener).source
|
||||
} else {
|
||||
legacy(&server, root.path(), opener)
|
||||
};
|
||||
let intent = if scenario.starts_with("refresh") {
|
||||
AuthIntent::Headless
|
||||
} else {
|
||||
AuthIntent::UserInitiated
|
||||
};
|
||||
assert!(source.acquire_with_intent(intent, None).await.is_err());
|
||||
let output = String::from_utf8(logs.0.lock().unwrap().clone()).unwrap();
|
||||
assert!(
|
||||
output.contains("oauth"),
|
||||
"capture must contain positive log evidence: strict={strict_path} scenario={scenario}"
|
||||
);
|
||||
for secret in [
|
||||
"secret-provider-body",
|
||||
"secret-opener-error",
|
||||
"secret-denial",
|
||||
"synthetic-refresh",
|
||||
"synthetic-code",
|
||||
"code_verifier",
|
||||
"code_challenge",
|
||||
&server.base,
|
||||
] {
|
||||
assert!(
|
||||
!output.contains(secret),
|
||||
"diagnostic leaked {secret}: {output}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn strict_malformed_token_expiry_is_a_safe_infrastructure_failure() {
|
||||
let server = Server::start(true, |base, r| {
|
||||
if r.path.contains("well-known") {
|
||||
discovery(base)
|
||||
} else {
|
||||
Reply::Json(
|
||||
200,
|
||||
json!({"access_token":"synthetic", "expires_in":u64::MAX}),
|
||||
)
|
||||
}
|
||||
})
|
||||
.await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let opener = Arc::new(Opener {
|
||||
callback: true,
|
||||
..Default::default()
|
||||
});
|
||||
assert_eq!(
|
||||
strict(&server, root.path(), opener).connect().await,
|
||||
Err(AuthError::NetworkUnavailable)
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn strict_catalog_errors_redact_provider_text_and_urls() {
|
||||
let server = Server::start(true, |_, _| {
|
||||
Reply::Json(
|
||||
403,
|
||||
json!({"message":"secret-provider-diagnostic https://private.invalid/?token=secret"}),
|
||||
)
|
||||
})
|
||||
.await;
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
seed(
|
||||
root.path(),
|
||||
&server.base,
|
||||
"databricks-strict",
|
||||
"synthetic-access",
|
||||
false,
|
||||
);
|
||||
let c = strict(&server, root.path(), Arc::new(Opener::default()));
|
||||
assert_eq!(
|
||||
c.discover_models(None).await.unwrap_err().to_string(),
|
||||
AgentError::Llm("Databricks model discovery unavailable".into()).to_string()
|
||||
);
|
||||
assert_eq!(server.count("/api/"), 2);
|
||||
}
|
||||
@@ -1,9 +1,11 @@
|
||||
#![forbid(unsafe_code)]
|
||||
mod agent;
|
||||
pub mod auth;
|
||||
mod auth_http;
|
||||
mod builtin;
|
||||
pub mod catalog;
|
||||
pub mod config;
|
||||
pub mod databricks;
|
||||
mod handoff;
|
||||
mod hints;
|
||||
mod llm;
|
||||
|
||||
48
crates/buzz-agent/tests/databricks_cli_auth.rs
Normal file
48
crates/buzz-agent/tests/databricks_cli_auth.rs
Normal file
@@ -0,0 +1,48 @@
|
||||
//! CLI compatibility: the real auth subcommand uses the legacy cache layout.
|
||||
//! No ACP, browser, external network, or user cache. Every child has a fresh root.
|
||||
#![cfg(unix)]
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::time::Duration;
|
||||
|
||||
#[tokio::test]
|
||||
async fn cli_signin_aliases_reuse_legacy_cache_without_runtime_configuration() {
|
||||
for provider in ["databricks", "databricks_v2", "databricks-v2"] {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let host = "https://workspace.invalid";
|
||||
let key = format!("{host}/oidc/.well-known/oauth-authorization-server|databricks-cli|all-apis,offline_access");
|
||||
let dir = root.path().join("buzz-agent/oauth/databricks");
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let cache = dir.join(format!("{}.json", hex::encode(Sha256::digest(key))));
|
||||
let content = serde_json::json!({"access_token":"synthetic-cli-cache", "refresh_token":null, "expires_at":4_000_000_000u64}).to_string();
|
||||
std::fs::write(&cache, &content).unwrap();
|
||||
let output = tokio::time::timeout(
|
||||
Duration::from_secs(5),
|
||||
tokio::process::Command::new(env!("CARGO_BIN_EXE_buzz-agent"))
|
||||
.args(["auth", provider])
|
||||
.env_clear()
|
||||
.env("BUZZ_AGENT_CONFIG_DIR", root.path())
|
||||
.env("HOME", root.path())
|
||||
.env("DATABRICKS_HOST", host)
|
||||
.kill_on_drop(true)
|
||||
.output(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
assert!(output.stdout.is_empty(), "auth must not emit ACP frames");
|
||||
let stderr = String::from_utf8(output.stderr).unwrap();
|
||||
assert!(stderr.contains("Authenticated. Token cached"));
|
||||
assert!(!stderr.contains("Opening browser"));
|
||||
assert_eq!(std::fs::read_to_string(&cache).unwrap(), content);
|
||||
assert!(!root
|
||||
.path()
|
||||
.join("buzz-agent/oauth/databricks-strict")
|
||||
.exists());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user