mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-08-04 11:43:57 +08:00
fix(pi): harden gateway runtime boundaries
This commit is contained in:
@@ -99,6 +99,7 @@ windows-sys = { version = "0.61", features = [
|
|||||||
"Win32_Globalization",
|
"Win32_Globalization",
|
||||||
"Win32_Security",
|
"Win32_Security",
|
||||||
"Win32_Storage_FileSystem",
|
"Win32_Storage_FileSystem",
|
||||||
|
"Win32_System_Diagnostics_ToolHelp",
|
||||||
"Win32_System_JobObjects",
|
"Win32_System_JobObjects",
|
||||||
"Win32_System_IO",
|
"Win32_System_IO",
|
||||||
"Win32_System_Threading",
|
"Win32_System_Threading",
|
||||||
|
|||||||
@@ -277,7 +277,15 @@ impl CandidateHeaderPlan {
|
|||||||
Some(_) => {}
|
Some(_) => {}
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut protocol_identity_predictable = true;
|
// Anthropic's credential kind changes the pinned wire protocol:
|
||||||
|
// OAuth adds a different beta profile. A command-valued credential
|
||||||
|
// therefore cannot be pre-classified without executing the command,
|
||||||
|
// and commands are deliberately executed only for the one real
|
||||||
|
// network attempt.
|
||||||
|
let mut protocol_identity_predictable = !matches!(
|
||||||
|
(family, credential),
|
||||||
|
(PiGatewayApiFamily::AnthropicMessages, Some(value)) if value.starts_with('!')
|
||||||
|
);
|
||||||
let provider_headers = plan_configured_headers(
|
let provider_headers = plan_configured_headers(
|
||||||
&model.provider_headers,
|
&model.provider_headers,
|
||||||
&mut reasons,
|
&mut reasons,
|
||||||
@@ -1208,6 +1216,47 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn anthropic_command_credentials_are_single_direct_attempts_after_materialization() {
|
||||||
|
for (resolved, expected_auth, oauth) in [
|
||||||
|
("ordinary-secret", "ordinary-secret", false),
|
||||||
|
(
|
||||||
|
"prefix-sk-ant-oat01-token-suffix",
|
||||||
|
"Bearer prefix-sk-ant-oat01-token-suffix",
|
||||||
|
true,
|
||||||
|
),
|
||||||
|
] {
|
||||||
|
let composition = composed(json!({
|
||||||
|
"api": "anthropic-messages",
|
||||||
|
"baseUrl": "https://candidate.example",
|
||||||
|
"apiKey": "!credential-command",
|
||||||
|
"models": [{"id": "m"}]
|
||||||
|
}));
|
||||||
|
let mut assessment = assess_composition_for_runtime(&composition);
|
||||||
|
assert_eq!(assessment.capability, PiGatewayCapability::Proxyable);
|
||||||
|
let plan = assessment.plans.remove(0);
|
||||||
|
assert!(
|
||||||
|
!plan.protocol_identity_is_predictable(),
|
||||||
|
"credential commands affecting Anthropic beta must not be pre-executed"
|
||||||
|
);
|
||||||
|
let materialized = plan
|
||||||
|
.materialize_for_runtime(&|expression: &str| {
|
||||||
|
(expression == "!credential-command").then(|| resolved.to_string())
|
||||||
|
})
|
||||||
|
.expect("materialize command result once");
|
||||||
|
let auth_name = if oauth { "authorization" } else { "x-api-key" };
|
||||||
|
assert_eq!(
|
||||||
|
materialized.headers[&HeaderName::from_static(auth_name)],
|
||||||
|
expected_auth
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
materialized.headers.get("x-api-key").is_none(),
|
||||||
|
oauth,
|
||||||
|
"OAuth must never be proxied through x-api-key"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn deferred_values_replay_actual_pinned_pi_transport_results() {
|
fn deferred_values_replay_actual_pinned_pi_transport_results() {
|
||||||
let oracle: Value =
|
let oracle: Value =
|
||||||
|
|||||||
@@ -79,14 +79,12 @@ pub(crate) async fn handle_pi_native(
|
|||||||
drop(admission);
|
drop(admission);
|
||||||
|
|
||||||
let is_streaming = request_is_streaming(&uri, &incoming_headers, &request_json);
|
let is_streaming = request_is_streaming(&uri, &incoming_headers, &request_json);
|
||||||
let max_attempts = (route.app_config.max_retries as usize)
|
// Retry policy counts actual upstream sends. Circuit-open candidates,
|
||||||
.saturating_add(1)
|
// protocol-ineligible failovers, and materialization failures must not
|
||||||
.min(route.candidates.len());
|
// consume the budget or hide a later eligible candidate.
|
||||||
let attempts = route
|
let mut network_budget =
|
||||||
.candidates
|
NetworkAttemptBudget::new((route.app_config.max_retries as usize).saturating_add(1));
|
||||||
.into_iter()
|
let attempts = route.candidates;
|
||||||
.take(max_attempts)
|
|
||||||
.collect::<Vec<_>>();
|
|
||||||
let request_headers = filtered_incoming_headers(&incoming_headers);
|
let request_headers = filtered_incoming_headers(&incoming_headers);
|
||||||
let started = Instant::now();
|
let started = Instant::now();
|
||||||
let session_id =
|
let session_id =
|
||||||
@@ -96,7 +94,7 @@ pub(crate) async fn handle_pi_native(
|
|||||||
record_request_start(&state).await;
|
record_request_start(&state).await;
|
||||||
|
|
||||||
let mut index = 0;
|
let mut index = 0;
|
||||||
while index < attempts.len() {
|
while index < attempts.len() && network_budget.has_remaining() {
|
||||||
let provider_id = attempts[index].provider_id.clone();
|
let provider_id = attempts[index].provider_id.clone();
|
||||||
let provider_end = attempts[index..]
|
let provider_end = attempts[index..]
|
||||||
.iter()
|
.iter()
|
||||||
@@ -114,8 +112,10 @@ pub(crate) async fn handle_pi_native(
|
|||||||
|
|
||||||
let mut saw_health_failure = false;
|
let mut saw_health_failure = false;
|
||||||
let mut reached_network = false;
|
let mut reached_network = false;
|
||||||
for (offset, candidate) in attempts[index..provider_end].iter().cloned().enumerate() {
|
for candidate in attempts[index..provider_end].iter().cloned() {
|
||||||
let attempt_index = index + offset;
|
if !network_budget.has_remaining() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
let Some(single_direct_attempt) =
|
let Some(single_direct_attempt) =
|
||||||
begin_protocol_materialization(&mut protocol_anchor, candidate.is_failover)
|
begin_protocol_materialization(&mut protocol_anchor, candidate.is_failover)
|
||||||
else {
|
else {
|
||||||
@@ -149,6 +149,9 @@ pub(crate) async fn handle_pi_native(
|
|||||||
} else {
|
} else {
|
||||||
route.app_config.non_streaming_timeout
|
route.app_config.non_streaming_timeout
|
||||||
};
|
};
|
||||||
|
if !network_budget.begin_send() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
let send = crate::proxy::http_client::get()
|
let send = crate::proxy::http_client::get()
|
||||||
.request(method.clone(), materialized.url.clone())
|
.request(method.clone(), materialized.url.clone())
|
||||||
.headers(outgoing_headers)
|
.headers(outgoing_headers)
|
||||||
@@ -180,8 +183,7 @@ pub(crate) async fn handle_pi_native(
|
|||||||
};
|
};
|
||||||
|
|
||||||
let status = response.status();
|
let status = response.status();
|
||||||
let is_last_attempt = attempt_index + 1 == attempts.len();
|
if retryable_status(status) && network_budget.has_remaining() {
|
||||||
if retryable_status(status) && !is_last_attempt {
|
|
||||||
saw_health_failure = true;
|
saw_health_failure = true;
|
||||||
last_error = Some(format!("Pi upstream returned retryable status {status}"));
|
last_error = Some(format!("Pi upstream returned retryable status {status}"));
|
||||||
continue;
|
continue;
|
||||||
@@ -268,6 +270,32 @@ pub(crate) async fn handle_pi_native(
|
|||||||
Err(ProxyError::ForwardFailed(error))
|
Err(ProxyError::ForwardFailed(error))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
struct NetworkAttemptBudget {
|
||||||
|
remaining: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl NetworkAttemptBudget {
|
||||||
|
fn new(max_attempts: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
remaining: max_attempts,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn has_remaining(&self) -> bool {
|
||||||
|
self.remaining > 0
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Consume budget only immediately before an actual upstream send.
|
||||||
|
fn begin_send(&mut self) -> bool {
|
||||||
|
if self.remaining == 0 {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
self.remaining -= 1;
|
||||||
|
true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Return whether this attempt consumes the one direct-only grant, or `None`
|
/// Return whether this attempt consumes the one direct-only grant, or `None`
|
||||||
/// when the candidate must not even be materialized.
|
/// when the candidate must not even be materialized.
|
||||||
fn begin_protocol_materialization(anchor: &mut ProtocolAnchor, is_failover: bool) -> Option<bool> {
|
fn begin_protocol_materialization(anchor: &mut ProtocolAnchor, is_failover: bool) -> Option<bool> {
|
||||||
@@ -278,10 +306,7 @@ fn begin_protocol_materialization(anchor: &mut ProtocolAnchor, is_failover: bool
|
|||||||
*anchor = ProtocolAnchor::Ineligible;
|
*anchor = ProtocolAnchor::Ineligible;
|
||||||
Some(true)
|
Some(true)
|
||||||
}
|
}
|
||||||
ProtocolAnchor::FailoverPending if is_failover => Some(false),
|
ProtocolAnchor::DirectOnlyPending | ProtocolAnchor::Ineligible => None,
|
||||||
ProtocolAnchor::DirectOnlyPending
|
|
||||||
| ProtocolAnchor::FailoverPending
|
|
||||||
| ProtocolAnchor::Ineligible => None,
|
|
||||||
ProtocolAnchor::Unset | ProtocolAnchor::Predictable(_) => Some(false),
|
ProtocolAnchor::Unset | ProtocolAnchor::Predictable(_) => Some(false),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -289,7 +314,6 @@ fn begin_protocol_materialization(anchor: &mut ProtocolAnchor, is_failover: bool
|
|||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
enum ProtocolAnchor {
|
enum ProtocolAnchor {
|
||||||
Unset,
|
Unset,
|
||||||
FailoverPending,
|
|
||||||
Predictable((String, HeaderMap)),
|
Predictable((String, HeaderMap)),
|
||||||
DirectOnlyPending,
|
DirectOnlyPending,
|
||||||
Ineligible,
|
Ineligible,
|
||||||
@@ -306,11 +330,10 @@ impl ProtocolAnchor {
|
|||||||
match primary.planned_protocol_identity() {
|
match primary.planned_protocol_identity() {
|
||||||
Ok(Some(identity)) => Self::Predictable(identity),
|
Ok(Some(identity)) => Self::Predictable(identity),
|
||||||
Ok(None) => Self::DirectOnlyPending,
|
Ok(None) => Self::DirectOnlyPending,
|
||||||
// A provider-level deferred credential/header may be unavailable
|
// If the primary's protocol identity cannot be established, a
|
||||||
// while a compatible backup is healthy. Skip every endpoint of
|
// backup must not self-declare compatibility. Give the primary
|
||||||
// that provider and let the first materialized failover establish
|
// exactly one direct materialization; failure remains fail-closed.
|
||||||
// the request's protocol anchor.
|
Err(_) => Self::DirectOnlyPending,
|
||||||
Err(_) => Self::FailoverPending,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -320,11 +343,11 @@ fn protocol_identity_allows_attempt(
|
|||||||
candidate: Option<(String, HeaderMap)>,
|
candidate: Option<(String, HeaderMap)>,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
match (&*anchor, candidate) {
|
match (&*anchor, candidate) {
|
||||||
(ProtocolAnchor::Unset | ProtocolAnchor::FailoverPending, Some(candidate)) => {
|
(ProtocolAnchor::Unset, Some(candidate)) => {
|
||||||
*anchor = ProtocolAnchor::Predictable(candidate);
|
*anchor = ProtocolAnchor::Predictable(candidate);
|
||||||
true
|
true
|
||||||
}
|
}
|
||||||
(ProtocolAnchor::Unset | ProtocolAnchor::FailoverPending, None) => false,
|
(ProtocolAnchor::Unset, None) => false,
|
||||||
(ProtocolAnchor::Predictable(primary), Some(candidate)) => primary == &candidate,
|
(ProtocolAnchor::Predictable(primary), Some(candidate)) => primary == &candidate,
|
||||||
(ProtocolAnchor::Predictable(_), None)
|
(ProtocolAnchor::Predictable(_), None)
|
||||||
| (ProtocolAnchor::DirectOnlyPending, _)
|
| (ProtocolAnchor::DirectOnlyPending, _)
|
||||||
@@ -1270,22 +1293,29 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn first_compatible_failover_can_anchor_when_primary_identity_is_unavailable() {
|
fn unavailable_primary_identity_is_direct_only_and_cannot_self_anchor_from_failover() {
|
||||||
let mut anchor = ProtocolAnchor::FailoverPending;
|
let mut anchor = ProtocolAnchor::DirectOnlyPending;
|
||||||
let mut candidate_headers = HeaderMap::new();
|
|
||||||
candidate_headers.insert(
|
|
||||||
"anthropic-version",
|
|
||||||
http::HeaderValue::from_static("candidate-only"),
|
|
||||||
);
|
|
||||||
let candidate = Some(("anthropic-messages".to_string(), candidate_headers));
|
|
||||||
|
|
||||||
assert_eq!(begin_protocol_materialization(&mut anchor, false), None);
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
begin_protocol_materialization(&mut anchor, true),
|
begin_protocol_materialization(&mut anchor, false),
|
||||||
Some(false)
|
Some(true)
|
||||||
);
|
);
|
||||||
assert!(protocol_identity_allows_attempt(&mut anchor, candidate));
|
assert_eq!(begin_protocol_materialization(&mut anchor, true), None);
|
||||||
assert!(matches!(anchor, ProtocolAnchor::Predictable(_)));
|
assert!(matches!(anchor, ProtocolAnchor::Ineligible));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn skipped_candidates_do_not_reduce_the_network_retry_budget() {
|
||||||
|
let mut budget = NetworkAttemptBudget::new(2);
|
||||||
|
|
||||||
|
// Circuit, protocol, and materialization skips never call begin_send.
|
||||||
|
for _ in 0..4 {
|
||||||
|
assert!(budget.has_remaining());
|
||||||
|
}
|
||||||
|
assert!(budget.begin_send());
|
||||||
|
assert!(budget.has_remaining());
|
||||||
|
assert!(budget.begin_send());
|
||||||
|
assert!(!budget.has_remaining());
|
||||||
|
assert!(!budget.begin_send());
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -23,10 +23,7 @@ use sha2::{Digest, Sha256};
|
|||||||
use std::collections::{BTreeMap, HashMap};
|
use std::collections::{BTreeMap, HashMap};
|
||||||
use std::io::Read;
|
use std::io::Read;
|
||||||
use std::process::{Command, Stdio};
|
use std::process::{Command, Stdio};
|
||||||
use std::sync::{
|
use std::sync::{mpsc, Arc, RwLock};
|
||||||
atomic::{AtomicU64, Ordering},
|
|
||||||
mpsc, Arc, RwLock,
|
|
||||||
};
|
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
use tokio::sync::{OwnedRwLockReadGuard, RwLock as AsyncRwLock};
|
use tokio::sync::{OwnedRwLockReadGuard, RwLock as AsyncRwLock};
|
||||||
use url::Url;
|
use url::Url;
|
||||||
@@ -577,26 +574,39 @@ fn pi_route_token(provider_id: &str, provider_key: &str) -> String {
|
|||||||
/// catalog while a replacement is prepared.
|
/// catalog while a replacement is prepared.
|
||||||
#[derive(Debug, Default)]
|
#[derive(Debug, Default)]
|
||||||
pub(crate) struct PiRuntimeStore {
|
pub(crate) struct PiRuntimeStore {
|
||||||
current: RwLock<Option<Arc<PiRuntimeSnapshot>>>,
|
publication: RwLock<PiRuntimePublication>,
|
||||||
catalog_epoch: AtomicU64,
|
|
||||||
epoch_gate: Arc<AsyncRwLock<()>>,
|
epoch_gate: Arc<AsyncRwLock<()>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Default)]
|
||||||
|
struct PiRuntimePublication {
|
||||||
|
current: Option<Arc<PiRuntimeSnapshot>>,
|
||||||
|
catalog_epoch: u64,
|
||||||
|
}
|
||||||
|
|
||||||
impl PiRuntimeStore {
|
impl PiRuntimeStore {
|
||||||
pub(crate) async fn begin_mutation(&self) -> u64 {
|
pub(crate) async fn begin_mutation(&self) -> u64 {
|
||||||
let _guard = self.epoch_gate.write().await;
|
let _guard = self.epoch_gate.write().await;
|
||||||
let current = self.catalog_epoch.load(Ordering::Acquire);
|
let mut publication = self
|
||||||
|
.publication
|
||||||
|
.write()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
|
let current = publication.catalog_epoch;
|
||||||
let odd = if current % 2 == 0 {
|
let odd = if current % 2 == 0 {
|
||||||
current.saturating_add(1)
|
current.saturating_add(1)
|
||||||
} else {
|
} else {
|
||||||
current
|
current
|
||||||
};
|
};
|
||||||
self.catalog_epoch.store(odd, Ordering::Release);
|
publication.catalog_epoch = odd;
|
||||||
odd.saturating_add(1)
|
odd.saturating_add(1)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn next_even_epoch(&self) -> Result<u64, AppError> {
|
pub(crate) fn next_even_epoch(&self) -> Result<u64, AppError> {
|
||||||
let current = self.catalog_epoch.load(Ordering::Acquire);
|
let current = self
|
||||||
|
.publication
|
||||||
|
.read()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||||
|
.catalog_epoch;
|
||||||
if current % 2 != 0 {
|
if current % 2 != 0 {
|
||||||
return Err(AppError::Conflict(
|
return Err(AppError::Conflict(
|
||||||
"cannot publish a sorted Pi runtime while catalog admission is fenced".to_string(),
|
"cannot publish a sorted Pi runtime while catalog admission is fenced".to_string(),
|
||||||
@@ -618,12 +628,12 @@ impl PiRuntimeStore {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
let _guard = self.epoch_gate.write().await;
|
let _guard = self.epoch_gate.write().await;
|
||||||
let epoch = snapshot.catalog_epoch;
|
let mut publication = self
|
||||||
*self
|
.publication
|
||||||
.current
|
|
||||||
.write()
|
.write()
|
||||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(snapshot);
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
self.catalog_epoch.store(epoch, Ordering::Release);
|
publication.catalog_epoch = snapshot.catalog_epoch;
|
||||||
|
publication.current = Some(snapshot);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -634,46 +644,54 @@ impl PiRuntimeStore {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
let _guard = self.epoch_gate.write().await;
|
let _guard = self.epoch_gate.write().await;
|
||||||
*self
|
let mut publication = self
|
||||||
.current
|
.publication
|
||||||
.write()
|
.write()
|
||||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = None;
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
self.catalog_epoch.store(even_epoch, Ordering::Release);
|
publication.current = None;
|
||||||
|
publication.catalog_epoch = even_epoch;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn republish_current(&self, even_epoch: u64) -> Result<bool, AppError> {
|
pub(crate) async fn republish_current(&self, even_epoch: u64) -> Result<bool, AppError> {
|
||||||
let current = self
|
if even_epoch % 2 != 0 {
|
||||||
.current
|
return Err(AppError::Config(
|
||||||
.read()
|
"Pi runtime re-publication requires an even epoch".to_string(),
|
||||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
));
|
||||||
.as_ref()
|
}
|
||||||
.cloned();
|
let _guard = self.epoch_gate.write().await;
|
||||||
|
let mut publication = self
|
||||||
|
.publication
|
||||||
|
.write()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
|
let current = publication.current.as_ref().cloned();
|
||||||
let Some(current) = current else {
|
let Some(current) = current else {
|
||||||
self.close(even_epoch).await?;
|
publication.catalog_epoch = even_epoch;
|
||||||
return Ok(false);
|
return Ok(false);
|
||||||
};
|
};
|
||||||
let mut next = (*current).clone();
|
let mut next = (*current).clone();
|
||||||
next.catalog_epoch = even_epoch;
|
next.catalog_epoch = even_epoch;
|
||||||
self.publish(Arc::new(next)).await?;
|
publication.current = Some(Arc::new(next));
|
||||||
|
publication.catalog_epoch = even_epoch;
|
||||||
Ok(true)
|
Ok(true)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn lease(&self, server_generation: u64) -> Option<Arc<PiRuntimeSnapshot>> {
|
pub(crate) fn lease(&self, server_generation: u64) -> Option<Arc<PiRuntimeSnapshot>> {
|
||||||
let epoch = self.catalog_epoch.load(Ordering::Acquire);
|
let publication = self
|
||||||
if epoch % 2 != 0 {
|
.publication
|
||||||
|
.read()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
|
if publication.catalog_epoch % 2 != 0 {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
let snapshot = self
|
publication
|
||||||
.current
|
.current
|
||||||
.read()
|
|
||||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.filter(|snapshot| {
|
.filter(|snapshot| {
|
||||||
snapshot.server_generation == server_generation && snapshot.catalog_epoch == epoch
|
snapshot.server_generation == server_generation
|
||||||
|
&& snapshot.catalog_epoch == publication.catalog_epoch
|
||||||
})
|
})
|
||||||
.cloned()?;
|
.cloned()
|
||||||
(self.catalog_epoch.load(Ordering::Acquire) == epoch).then_some(snapshot)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn is_admitting(&self, server_generation: u64) -> bool {
|
pub(crate) fn is_admitting(&self, server_generation: u64) -> bool {
|
||||||
@@ -686,17 +704,16 @@ impl PiRuntimeStore {
|
|||||||
snapshot: &Arc<PiRuntimeSnapshot>,
|
snapshot: &Arc<PiRuntimeSnapshot>,
|
||||||
) -> Option<OwnedRwLockReadGuard<()>> {
|
) -> Option<OwnedRwLockReadGuard<()>> {
|
||||||
let guard = self.epoch_gate.clone().read_owned().await;
|
let guard = self.epoch_gate.clone().read_owned().await;
|
||||||
let current = self
|
let publication = self
|
||||||
.current
|
.publication
|
||||||
.read()
|
.read()
|
||||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
.as_ref()
|
let current = publication.current.as_ref().is_some_and(|current| {
|
||||||
.is_some_and(|current| {
|
snapshot.catalog_epoch % 2 == 0
|
||||||
snapshot.catalog_epoch % 2 == 0
|
&& publication.catalog_epoch == snapshot.catalog_epoch
|
||||||
&& self.catalog_epoch.load(Ordering::Acquire) == snapshot.catalog_epoch
|
&& current.server_generation == server_generation
|
||||||
&& current.server_generation == server_generation
|
&& Arc::ptr_eq(current, snapshot)
|
||||||
&& Arc::ptr_eq(current, snapshot)
|
});
|
||||||
});
|
|
||||||
current.then_some(guard)
|
current.then_some(guard)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -705,8 +722,12 @@ impl PiRuntimeStore {
|
|||||||
expected_epoch: u64,
|
expected_epoch: u64,
|
||||||
) -> Option<OwnedRwLockReadGuard<()>> {
|
) -> Option<OwnedRwLockReadGuard<()>> {
|
||||||
let guard = self.epoch_gate.clone().read_owned().await;
|
let guard = self.epoch_gate.clone().read_owned().await;
|
||||||
(expected_epoch % 2 == 0 && self.catalog_epoch.load(Ordering::Acquire) == expected_epoch)
|
let current = self
|
||||||
.then_some(guard)
|
.publication
|
||||||
|
.read()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||||
|
.catalog_epoch;
|
||||||
|
(expected_epoch % 2 == 0 && current == expected_epoch).then_some(guard)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -835,6 +856,15 @@ fn execute_config_command(script: &str) -> Result<String, String> {
|
|||||||
}
|
}
|
||||||
command
|
command
|
||||||
};
|
};
|
||||||
|
#[cfg(windows)]
|
||||||
|
{
|
||||||
|
use std::os::windows::process::CommandExt;
|
||||||
|
use windows_sys::Win32::System::Threading::CREATE_SUSPENDED;
|
||||||
|
// The shell must not execute user code before it belongs to the
|
||||||
|
// kill-on-close Job Object. Its primary thread is resumed only after
|
||||||
|
// CommandTree::attach succeeds.
|
||||||
|
command.creation_flags(CREATE_SUSPENDED);
|
||||||
|
}
|
||||||
let mut child = command
|
let mut child = command
|
||||||
.stdin(Stdio::null())
|
.stdin(Stdio::null())
|
||||||
.stdout(Stdio::piped())
|
.stdout(Stdio::piped())
|
||||||
@@ -842,6 +872,12 @@ fn execute_config_command(script: &str) -> Result<String, String> {
|
|||||||
.spawn()
|
.spawn()
|
||||||
.map_err(|error| format!("failed to start Pi config command: {error}"))?;
|
.map_err(|error| format!("failed to start Pi config command: {error}"))?;
|
||||||
let command_tree = CommandTree::attach(&mut child)?;
|
let command_tree = CommandTree::attach(&mut child)?;
|
||||||
|
#[cfg(windows)]
|
||||||
|
if let Err(error) = resume_suspended_process(child.id()) {
|
||||||
|
command_tree.terminate(&mut child);
|
||||||
|
let _ = child.wait();
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
let stdout = child
|
let stdout = child
|
||||||
.stdout
|
.stdout
|
||||||
.take()
|
.take()
|
||||||
@@ -1013,6 +1049,59 @@ impl CommandTree {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(windows)]
|
||||||
|
fn resume_suspended_process(process_id: u32) -> Result<(), String> {
|
||||||
|
use std::mem::size_of;
|
||||||
|
use windows_sys::Win32::Foundation::{CloseHandle, INVALID_HANDLE_VALUE};
|
||||||
|
use windows_sys::Win32::System::Diagnostics::ToolHelp::{
|
||||||
|
CreateToolhelp32Snapshot, Thread32First, Thread32Next, TH32CS_SNAPTHREAD, THREADENTRY32,
|
||||||
|
};
|
||||||
|
use windows_sys::Win32::System::Threading::{OpenThread, ResumeThread, THREAD_SUSPEND_RESUME};
|
||||||
|
|
||||||
|
// SAFETY: the snapshot and thread handles are checked before use and
|
||||||
|
// closed on every path. CREATE_SUSPENDED prevents the target from adding
|
||||||
|
// threads while this enumeration runs.
|
||||||
|
unsafe {
|
||||||
|
let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
|
||||||
|
if snapshot == INVALID_HANDLE_VALUE {
|
||||||
|
return Err(format!(
|
||||||
|
"failed to enumerate the suspended Pi config command: {}",
|
||||||
|
std::io::Error::last_os_error()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let mut entry = THREADENTRY32 {
|
||||||
|
dwSize: size_of::<THREADENTRY32>() as u32,
|
||||||
|
..THREADENTRY32::default()
|
||||||
|
};
|
||||||
|
let mut has_entry = Thread32First(snapshot, &mut entry);
|
||||||
|
while has_entry != 0 {
|
||||||
|
if entry.th32OwnerProcessID == process_id {
|
||||||
|
let thread = OpenThread(THREAD_SUSPEND_RESUME, 0, entry.th32ThreadID);
|
||||||
|
if thread.is_null() {
|
||||||
|
let error = std::io::Error::last_os_error();
|
||||||
|
CloseHandle(snapshot);
|
||||||
|
return Err(format!(
|
||||||
|
"failed to open the suspended Pi config command thread: {error}"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let result = ResumeThread(thread);
|
||||||
|
let error = (result == u32::MAX).then(std::io::Error::last_os_error);
|
||||||
|
CloseHandle(thread);
|
||||||
|
CloseHandle(snapshot);
|
||||||
|
return match error {
|
||||||
|
Some(error) => Err(format!(
|
||||||
|
"failed to resume the suspended Pi config command: {error}"
|
||||||
|
)),
|
||||||
|
None => Ok(()),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
has_entry = Thread32Next(snapshot, &mut entry);
|
||||||
|
}
|
||||||
|
CloseHandle(snapshot);
|
||||||
|
}
|
||||||
|
Err("the suspended Pi config command had no primary thread".to_string())
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(windows)]
|
#[cfg(windows)]
|
||||||
impl Drop for CommandTree {
|
impl Drop for CommandTree {
|
||||||
fn drop(&mut self) {
|
fn drop(&mut self) {
|
||||||
@@ -1065,6 +1154,23 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(windows)]
|
||||||
|
#[test]
|
||||||
|
fn command_job_assignment_precedes_descendant_execution() {
|
||||||
|
let started = Instant::now();
|
||||||
|
for _ in 0..8 {
|
||||||
|
let output = execute_config_command(
|
||||||
|
"start \"\" /b cmd /D /S /C \"ping 127.0.0.1 -n 30 >nul\" & <nul set /p =ready",
|
||||||
|
)
|
||||||
|
.expect("command output");
|
||||||
|
assert_eq!(output, "ready");
|
||||||
|
}
|
||||||
|
assert!(
|
||||||
|
started.elapsed() < Duration::from_secs(5),
|
||||||
|
"descendants must remain in the job and release inherited pipes"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn family_url_builders_preserve_candidate_origin_and_base_path() {
|
fn family_url_builders_preserve_candidate_origin_and_base_path() {
|
||||||
let base = Url::parse("https://candidate.example:8443/root/v1").unwrap();
|
let base = Url::parse("https://candidate.example:8443/root/v1").unwrap();
|
||||||
|
|||||||
Reference in New Issue
Block a user