use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, OnceLock};
use chio_core::crypto::{Ed25519Backend, SigningBackend};
use chio_log_redact::redacted;
use tokio::runtime::Handle;
use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit, Semaphore};
use tokio::task::JoinHandle;
use tracing::{debug, warn};
use crate::{ChioReceipt, ChioReceiptBody, KernelError, Keypair, DEFAULT_MAX_STREAM_TOTAL_BYTES};
pub const DEFAULT_SIGNING_CHANNEL_CAPACITY: usize = 256;
#[allow(dead_code)]
pub const DEFAULT_MAX_SIGNING_CONTENT_BYTES: usize =
clamp_u64_to_usize(DEFAULT_MAX_STREAM_TOTAL_BYTES);
const PER_REQUEST_BUDGET_UNLIMITED: usize = 0;
pub(crate) const DEFAULT_MAX_SIGNING_QUEUED_BYTES: usize =
clamp_u64_to_usize(DEFAULT_MAX_STREAM_TOTAL_BYTES);
const fn clamp_aggregate_permits(budget: usize) -> u32 {
let ceiling = u32::MAX as usize;
let clamped = if budget > ceiling { ceiling } else { budget };
if clamped == 0 {
1
} else {
clamped as u32
}
}
#[allow(dead_code)]
const fn clamp_u64_to_usize(value: u64) -> usize {
if value > usize::MAX as u64 {
usize::MAX
} else {
value as usize
}
}
#[cfg_attr(test, allow(unused_imports))]
pub use chio_metrics_spec::CHIO_SIGNING_QUEUE_BLOCK_TOTAL as METRIC_CHIO_SIGNING_QUEUE_BLOCK_TOTAL;
fn record_signing_queue_block(reason: &str) {
chio_metrics_spec::runtime::families::SIGNING_QUEUE_BLOCK.incr(&[reason]);
}
#[allow(dead_code)]
type TrySignOutcome =
Result<oneshot::Receiver<Result<ChioReceipt, KernelError>>, (ChioReceiptBody, Vec<u8>)>;
enum EnqueueOutcome {
Enqueued(oneshot::Receiver<Result<ChioReceipt, KernelError>>),
Backpressure(ChioReceiptBody, Vec<u8>),
Closed(ChioReceiptBody, Vec<u8>),
}
pub(crate) struct SignRequest {
pub(crate) body: ChioReceiptBody,
pub(crate) canonical_content: Vec<u8>,
pub(crate) reply: oneshot::Sender<Result<ChioReceipt, KernelError>>,
pub(crate) _aggregate_permit: Option<OwnedSemaphorePermit>,
}
struct SigningTaskInner {
sender: Mutex<Option<mpsc::Sender<SignRequest>>>,
join: Mutex<Option<JoinHandle<()>>>,
}
impl SigningTaskInner {
fn sender_clone(&self) -> Option<mpsc::Sender<SignRequest>> {
match self.sender.lock() {
Ok(slot) => slot.as_ref().cloned(),
Err(poisoned) => poisoned.into_inner().as_ref().cloned(),
}
}
}
pub(crate) struct SigningTaskHandle {
inner: OnceLock<SigningTaskInner>,
keypair: Keypair,
capacity: usize,
max_content_bytes: usize,
aggregate_byte_budget: Arc<Semaphore>,
aggregate_budget_permits: u32,
spawn_gate: Mutex<()>,
closed: AtomicBool,
}
impl SigningTaskHandle {
#[allow(dead_code)]
pub(crate) fn spawn(keypair: Keypair) -> Self {
Self::with_capacity(keypair, DEFAULT_SIGNING_CHANNEL_CAPACITY)
}
#[allow(dead_code)]
pub(crate) fn with_capacity(keypair: Keypair, capacity: usize) -> Self {
Self::with_capacity_and_max_content_bytes(
keypair,
capacity,
DEFAULT_MAX_SIGNING_CONTENT_BYTES,
)
}
pub(crate) fn with_capacity_and_max_content_bytes(
keypair: Keypair,
capacity: usize,
max_content_bytes: usize,
) -> Self {
Self::with_capacity_max_content_and_queued_bytes(
keypair,
capacity,
max_content_bytes,
if max_content_bytes == PER_REQUEST_BUDGET_UNLIMITED {
DEFAULT_MAX_SIGNING_QUEUED_BYTES
} else {
max_content_bytes
},
)
}
#[allow(dead_code)]
pub(crate) fn with_capacity_max_content_and_queued_bytes(
keypair: Keypair,
capacity: usize,
max_content_bytes: usize,
max_queued_bytes: usize,
) -> Self {
let capacity = capacity.max(1);
let aggregate_budget_permits = clamp_aggregate_permits(max_queued_bytes);
Self {
inner: OnceLock::new(),
keypair,
capacity,
max_content_bytes,
aggregate_byte_budget: Arc::new(Semaphore::new(aggregate_budget_permits as usize)),
aggregate_budget_permits,
spawn_gate: Mutex::new(()),
closed: AtomicBool::new(false),
}
}
fn lock_spawn_gate(&self) -> MutexGuard<'_, ()> {
match self.spawn_gate.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
}
}
fn shutdown_error() -> KernelError {
KernelError::Internal("receipt signing task already shut down".to_string())
}
fn oversized_content_error(&self, len: usize) -> KernelError {
let budget = self.max_content_bytes;
KernelError::ReceiptSigningFailed(format!(
"receipt signing refused: canonical content is {len} bytes, over the \
{budget}-byte per-request queue budget (sign oversized \
receipts inline rather than through the async queue)"
))
}
fn exceeds_per_request_cap(&self, len: usize) -> bool {
self.max_content_bytes != PER_REQUEST_BUDGET_UNLIMITED && len > self.max_content_bytes
}
fn permits_for(&self, len: usize) -> u32 {
debug_assert!(
len <= self.aggregate_budget_permits as usize,
"permits_for must only run for a preimage that fits the aggregate budget",
);
len as u32
}
fn exceeds_aggregate_budget(&self, len: usize) -> bool {
len > self.aggregate_budget_permits as usize
}
fn ensure_spawned(&self) -> Result<&SigningTaskInner, KernelError> {
let _spawn_guard = self.lock_spawn_gate();
if self.closed.load(Ordering::Acquire) {
return Err(Self::shutdown_error());
}
if let Some(inner) = self.inner.get() {
return Ok(inner);
}
let handle = Handle::try_current().map_err(|_| {
KernelError::Internal(
"receipt signing task requires an active tokio runtime".to_string(),
)
})?;
let (sender, receiver) = mpsc::channel::<SignRequest>(self.capacity);
let join = handle.spawn(run_signing_task(self.keypair.clone(), receiver));
let candidate = SigningTaskInner {
sender: Mutex::new(Some(sender)),
join: Mutex::new(Some(join)),
};
match self.inner.set(candidate) {
Ok(()) => self.inner.get().ok_or_else(|| {
KernelError::Internal("receipt signing task failed to initialize".to_string())
}),
Err(_orphan) => {
self.inner.get().ok_or_else(|| {
KernelError::Internal("receipt signing task failed to initialize".to_string())
})
}
}
}
fn try_enqueue_if_open(
&self,
body: ChioReceiptBody,
canonical_content: Vec<u8>,
) -> EnqueueOutcome {
let _spawn_guard = self.lock_spawn_gate();
if self.closed.load(Ordering::Acquire) {
return EnqueueOutcome::Closed(body, canonical_content);
}
let Some(sender) = self.inner.get().and_then(SigningTaskInner::sender_clone) else {
return EnqueueOutcome::Closed(body, canonical_content);
};
let permits = self.permits_for(canonical_content.len());
let permit = match Arc::clone(&self.aggregate_byte_budget).try_acquire_many_owned(permits) {
Ok(permit) => permit,
Err(_) => {
record_signing_queue_block("byte_budget");
return EnqueueOutcome::Backpressure(body, canonical_content);
}
};
let (reply_tx, reply_rx) = oneshot::channel();
let request = SignRequest {
body,
canonical_content,
reply: reply_tx,
_aggregate_permit: Some(permit),
};
match sender.try_send(request) {
Ok(()) => EnqueueOutcome::Enqueued(reply_rx),
Err(mpsc::error::TrySendError::Full(rejected)) => {
record_signing_queue_block("channel_full");
EnqueueOutcome::Backpressure(rejected.body, rejected.canonical_content)
}
Err(mpsc::error::TrySendError::Closed(rejected)) => {
EnqueueOutcome::Closed(rejected.body, rejected.canonical_content)
}
}
}
fn sign_inline(
&self,
body: ChioReceiptBody,
canonical_content: Vec<u8>,
) -> Result<ChioReceipt, KernelError> {
sign_one(&self.keypair, body, canonical_content)
}
pub(crate) async fn sign(
&self,
body: ChioReceiptBody,
canonical_content: Vec<u8>,
) -> Result<ChioReceipt, KernelError> {
let len = canonical_content.len();
if self.exceeds_per_request_cap(len) {
return Err(self.oversized_content_error(len));
}
if self.exceeds_aggregate_budget(len) {
record_signing_queue_block("oversized");
return self.sign_inline_if_open(body, canonical_content);
}
self.ensure_spawned()?;
match self.try_enqueue_if_open(body, canonical_content) {
EnqueueOutcome::Enqueued(reply_rx) => match reply_rx.await {
Ok(result) => result,
Err(_) => Err(KernelError::Internal(
"receipt signing task dropped reply channel".to_string(),
)),
},
EnqueueOutcome::Backpressure(body, canonical_content) => {
self.sign_inline(body, canonical_content)
}
EnqueueOutcome::Closed(_body, _canonical_content) => Err(KernelError::Internal(
"receipt signing task already shut down".to_string(),
)),
}
}
fn sign_inline_if_open(
&self,
body: ChioReceiptBody,
canonical_content: Vec<u8>,
) -> Result<ChioReceipt, KernelError> {
{
let _spawn_guard = self.lock_spawn_gate();
if self.closed.load(Ordering::Acquire) {
return Err(KernelError::Internal(
"receipt signing task already shut down".to_string(),
));
}
}
self.sign_inline(body, canonical_content)
}
#[allow(dead_code, clippy::result_large_err)]
pub(crate) fn try_sign(
&self,
body: ChioReceiptBody,
canonical_content: Vec<u8>,
) -> TrySignOutcome {
if self.exceeds_per_request_cap(canonical_content.len()) {
return Err((body, canonical_content));
}
if self.exceeds_aggregate_budget(canonical_content.len()) {
return Err((body, canonical_content));
}
let inner = match self.ensure_spawned() {
Ok(inner) => inner,
Err(_) => return Err((body, canonical_content)),
};
let Some(sender) = inner.sender_clone() else {
return Err((body, canonical_content));
};
let permits = self.permits_for(canonical_content.len());
let permit = match Arc::clone(&self.aggregate_byte_budget).try_acquire_many_owned(permits) {
Ok(permit) => permit,
Err(_) => return Err((body, canonical_content)),
};
let (reply_tx, reply_rx) = oneshot::channel();
let request = SignRequest {
body,
canonical_content,
reply: reply_tx,
_aggregate_permit: Some(permit),
};
match sender.try_send(request) {
Ok(()) => Ok(reply_rx),
Err(mpsc::error::TrySendError::Full(rejected)) => {
Err((rejected.body, rejected.canonical_content))
}
Err(mpsc::error::TrySendError::Closed(rejected)) => {
Err((rejected.body, rejected.canonical_content))
}
}
}
#[allow(dead_code)]
pub(crate) fn capacity(&self) -> usize {
self.capacity
}
#[allow(dead_code)]
pub(crate) fn is_spawned(&self) -> bool {
self.inner.get().is_some()
}
#[allow(dead_code)]
pub(crate) fn abort_for_crash_recovery_test(&self) {
let Some(inner) = self.inner.get() else {
return;
};
let dropped_sender = match inner.sender.lock() {
Ok(mut slot) => slot.take(),
Err(poisoned) => poisoned.into_inner().take(),
};
drop(dropped_sender);
let guard = match inner.join.lock() {
Ok(slot) => slot,
Err(poisoned) => poisoned.into_inner(),
};
if let Some(join) = guard.as_ref() {
join.abort();
}
}
pub(crate) async fn shutdown(&self) {
let join = {
let _spawn_guard = self.lock_spawn_gate();
self.closed.store(true, Ordering::Release);
let Some(inner) = self.inner.get() else {
return;
};
let dropped_sender = match inner.sender.lock() {
Ok(mut slot) => slot.take(),
Err(poisoned) => poisoned.into_inner().take(),
};
drop(dropped_sender);
match inner.join.lock() {
Ok(mut slot) => slot.take(),
Err(poisoned) => poisoned.into_inner().take(),
}
};
let Some(join) = join else {
return;
};
match join.await {
Ok(()) => {}
Err(err) if err.is_cancelled() => {
debug!("signing task cancelled before shutdown completed");
}
Err(err) => {
warn!(error = %redacted!(&err), "signing task join failed (panic)");
}
}
}
}
impl Drop for SigningTaskHandle {
fn drop(&mut self) {
}
}
async fn run_signing_task(keypair: Keypair, mut receiver: mpsc::Receiver<SignRequest>) {
debug!("signing task started");
while let Some(request) = receiver.recv().await {
let SignRequest {
body,
canonical_content,
reply,
_aggregate_permit,
} = request;
let result = sign_one(&keypair, body, canonical_content);
drop(_aggregate_permit);
let _ = reply.send(result);
}
debug!("signing task exited (channel closed)");
}
fn sign_one(
keypair: &Keypair,
body: ChioReceiptBody,
canonical_content: Vec<u8>,
) -> Result<ChioReceipt, KernelError> {
let backend = Ed25519Backend::new(keypair.clone());
sign_one_with_backend(body, &backend, canonical_content)
}
fn sign_one_with_backend(
body: ChioReceiptBody,
backend: &dyn SigningBackend,
canonical_content: Vec<u8>,
) -> Result<ChioReceipt, KernelError> {
let handle =
chio_core::receipt::signing::ReceiptSigningHandle::from_content_preimage(canonical_content);
chio_kernel_core::sign_receipt_with_handle(body, backend, handle).map_err(|error| {
use chio_kernel_core::ReceiptSigningError;
let message = match error {
ReceiptSigningError::KernelKeyMismatch => {
"kernel signing key does not match receipt body kernel_key".to_string()
}
ReceiptSigningError::ContentHashMismatch {
recomputed,
claimed,
} => format!(
"receipt content_hash mismatch: body claimed {claimed} but signer \
recomputed {recomputed} over the canonical content (WYSIWYS refused)"
),
ReceiptSigningError::SigningFailed(reason) => reason,
};
KernelError::ReceiptSigningFailed(message)
})
}