use std::error::Error;
use std::fmt;
use std::fs;
use std::future::Future;
use std::io;
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::{Duration, Instant, SystemTime};
use bytes::Bytes;
use http_body::{Body as HttpBody, Frame};
use http_body_util::BodyExt;
use tonic::body::Body;
use tonic::codegen::{BoxFuture, Service, http};
use tonic::metadata::MetadataValue;
use tonic::service::Interceptor;
use tonic::service::interceptor::{InterceptedService, ResponseBody};
use tonic::transport::Channel;
use tonic::{Request, Status};
type RejectedTokenRefreshFuture<'a> =
Pin<Box<dyn Future<Output = Result<String, CredentialError>> + Send + 'a>>;
#[cfg(feature = "api-key-credential")]
#[path = "credential/api_key_credential.rs"]
mod api_key_credential;
#[cfg(feature = "api-key-credential")]
pub use api_key_credential::{ApiKeyCredential, ApiKeyCredentialBuilder, ApiKeyCredentialError};
pub trait Credential: Send + Sync {
fn token(&self) -> Result<String, CredentialError>;
fn discard_rejected_token(&self, _rejected_token: &str) {}
fn token_after_rejection(&self) -> Option<RejectedTokenRefreshFuture<'_>> {
None
}
}
impl fmt::Debug for dyn Credential {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("Credential")
}
}
#[derive(Debug, Clone)]
pub struct CredentialError {
message: String,
}
impl CredentialError {
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl fmt::Display for CredentialError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl Error for CredentialError {}
pub struct FileCredential {
path: PathBuf,
state: Mutex<FileCredentialState>,
}
struct FileCredentialState {
token: Result<String, CredentialError>,
modified: SystemTime,
size: u64,
checked_at: Instant,
}
impl fmt::Debug for FileCredential {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("FileCredential")
.field("path", &self.path)
.field("token", &"<redacted>")
.finish()
}
}
impl FileCredential {
pub fn new(path: impl AsRef<Path>) -> io::Result<Self> {
let path = path.as_ref().to_path_buf();
let metadata = fs::metadata(&path)?;
let token = read_token(&path)?;
Ok(Self {
path,
state: Mutex::new(FileCredentialState {
token: Ok(token),
modified: metadata.modified()?,
size: metadata.len(),
checked_at: Instant::now(),
}),
})
}
}
impl Credential for FileCredential {
fn token(&self) -> Result<String, CredentialError> {
let mut state = self.state.lock().expect("credential file lock poisoned");
if state.checked_at.elapsed() >= Duration::from_secs(1) {
state.checked_at = Instant::now();
state.token = match fs::metadata(&self.path) {
Ok(metadata) => match metadata.modified() {
Ok(modified)
if state.token.is_ok()
&& modified == state.modified
&& metadata.len() == state.size =>
{
state.token.clone()
}
Ok(modified) => {
let token = read_token(&self.path).map_err(credential_file_error);
if token.is_ok() {
state.modified = modified;
state.size = metadata.len();
}
token
}
Err(error) => Err(credential_file_error(error)),
},
Err(error) => Err(credential_file_error(error)),
};
}
state.token.clone()
}
}
fn read_token(path: &Path) -> io::Result<String> {
let token = fs::read_to_string(path)?.trim().to_string();
if token.is_empty() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("credential file {} is empty", path.display()),
));
}
Ok(token)
}
fn credential_file_error(error: io::Error) -> CredentialError {
CredentialError::new(format!("read credential file: {error}"))
}
#[derive(Clone)]
pub(crate) struct CredentialChannel {
service: InterceptedService<RetryChannel, CredentialInterceptor>,
credential: Option<Arc<dyn Credential>>,
}
impl CredentialChannel {
pub(crate) fn new(channel: Channel, interceptor: CredentialInterceptor) -> Self {
let credential = interceptor.credential.clone();
let retry_channel = RetryChannel {
channel,
credential: credential.clone(),
};
Self {
service: InterceptedService::new(retry_channel, interceptor),
credential,
}
}
pub(crate) fn credential(&self) -> Option<Arc<dyn Credential>> {
self.credential.clone()
}
}
impl Service<http::Request<Body>> for CredentialChannel {
type Response = http::Response<ResponseBody<Body>>;
type Error = tonic::transport::Error;
type Future = BoxFuture<Self::Response, Self::Error>;
fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.service.poll_ready(context)
}
fn call(&mut self, request: http::Request<Body>) -> Self::Future {
Box::pin(self.service.call(request))
}
}
#[derive(Clone)]
struct RetryChannel {
channel: Channel,
credential: Option<Arc<dyn Credential>>,
}
impl Service<http::Request<Body>> for RetryChannel {
type Response = http::Response<Body>;
type Error = tonic::transport::Error;
type Future = BoxFuture<Self::Response, Self::Error>;
fn poll_ready(&mut self, _context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, request: http::Request<Body>) -> Self::Future {
let mut channel = self.channel.clone();
let credential = self.credential.clone();
let deadline = request_deadline(&request);
Box::pin(async move {
let attempt = async {
let mut request = request;
let mut rejected_token = match request.headers().get("authorization") {
Some(value) => match value.to_str() {
Ok(value) => value.strip_prefix("Bearer ").map(str::to_owned),
Err(_) => None,
},
None => None,
};
if rejected_token.is_none()
&& let Some(credential) = &credential
&& let Some(fresh_token) = credential.token_after_rejection()
{
let fresh_token = match fresh_token.await {
Ok(token) => token,
Err(error) => {
return Ok(Status::unauthenticated(error.to_string()).into_http());
}
};
let authorization = match bearer_header(&fresh_token) {
Ok(value) => value,
Err(status) => return Ok(status.into_http()),
};
request.headers_mut().insert("authorization", authorization);
rejected_token = Some(fresh_token);
}
let is_poll_stream =
request.uri().path() == "/dex.WorkerTaskService/PollWorkerTask";
if is_poll_stream || credential.is_none() {
let response = send_channel_request(&mut channel, request).await?;
let Some(credential) = credential else {
return Ok(response);
};
return Ok(observe_stream_rejection(
response,
credential,
rejected_token,
));
}
let (parts, body) = request.into_parts();
let request_body = match body.collect().await {
Ok(body) => body,
Err(status) => return Ok(status.into_http()),
};
let request_trailers = request_body.trailers().cloned();
let request_bytes = request_body.to_bytes();
let retry_method = parts.method.clone();
let retry_uri = parts.uri.clone();
let retry_version = parts.version;
let retry_headers = parts.headers.clone();
let request = http::Request::from_parts(
parts,
buffered_body(request_bytes.clone(), request_trailers.clone()),
);
let response = send_channel_request(&mut channel, request).await?;
let (response_parts, response_body) = response.into_parts();
let response_body = match response_body.collect().await {
Ok(body) => body,
Err(status) => return Ok(status.into_http()),
};
let response_trailers = response_body.trailers().cloned();
let is_rejected = is_unauthenticated(&response_parts.headers)
|| response_trailers.as_ref().is_some_and(is_unauthenticated);
if !is_rejected {
return Ok(http::Response::from_parts(
response_parts,
buffered_body(response_body.to_bytes(), response_trailers),
));
}
let Some(credential) = credential else {
return Ok(http::Response::from_parts(
response_parts,
buffered_body(response_body.to_bytes(), response_trailers),
));
};
let Some(rejected_token) = rejected_token else {
return Ok(http::Response::from_parts(
response_parts,
buffered_body(response_body.to_bytes(), response_trailers),
));
};
credential.discard_rejected_token(&rejected_token);
let Some(fresh_token) = credential.token_after_rejection() else {
return Ok(http::Response::from_parts(
response_parts,
buffered_body(response_body.to_bytes(), response_trailers),
));
};
let fresh_token = match fresh_token.await {
Ok(token) => token,
Err(error) => return Ok(Status::unauthenticated(error.to_string()).into_http()),
};
let authorization = match bearer_header(&fresh_token) {
Ok(value) => value,
Err(status) => return Ok(status.into_http()),
};
let mut retry = http::Request::new(buffered_body(request_bytes, request_trailers));
*retry.method_mut() = retry_method;
*retry.uri_mut() = retry_uri;
*retry.version_mut() = retry_version;
*retry.headers_mut() = retry_headers;
retry.headers_mut().insert("authorization", authorization);
if let Some(deadline) = deadline
&& !set_remaining_timeout(&mut retry, deadline)
{
return Ok(Status::deadline_exceeded("request deadline exceeded").into_http());
}
let response = send_channel_request(&mut channel, retry).await?;
let (parts, body) = response.into_parts();
let response = observe_stream_rejection(
http::Response::from_parts(parts, body),
credential,
Some(fresh_token),
);
Ok(response)
};
if let Some(deadline) = deadline {
match tokio::time::timeout_at(deadline.into(), attempt).await {
Ok(result) => result,
Err(_) => {
Ok(Status::deadline_exceeded("request deadline exceeded").into_http())
}
}
} else {
attempt.await
}
})
}
}
fn bearer_header(token: &str) -> Result<http::HeaderValue, Status> {
http::HeaderValue::from_str(&format!("Bearer {token}"))
.map_err(|_| Status::unauthenticated("invalid bearer token metadata"))
}
fn request_deadline(request: &http::Request<Body>) -> Option<Instant> {
let timeout = request.headers().get("grpc-timeout")?.to_str().ok()?;
let (amount, unit) = timeout.split_at(timeout.len().checked_sub(1)?);
if amount.is_empty() || amount.len() > 8 {
return None;
}
let amount = amount.parse::<u64>().ok()?;
let duration = match unit {
"H" => Duration::from_secs(amount.checked_mul(3_600)?),
"M" => Duration::from_secs(amount.checked_mul(60)?),
"S" => Duration::from_secs(amount),
"m" => Duration::from_millis(amount),
"u" => Duration::from_micros(amount),
"n" => Duration::from_nanos(amount),
_ => return None,
};
Instant::now().checked_add(duration)
}
fn set_remaining_timeout(request: &mut http::Request<Body>, deadline: Instant) -> bool {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return false;
}
let nanoseconds = remaining.as_nanos();
let timeout = [
(1, "n"),
(1_000, "u"),
(1_000_000, "m"),
(1_000_000_000, "S"),
(60_000_000_000, "M"),
(3_600_000_000_000, "H"),
]
.into_iter()
.find_map(|(unit, suffix)| {
let amount = nanoseconds.div_ceil(unit);
(amount <= 99_999_999).then(|| format!("{amount}{suffix}"))
})
.expect("gRPC timeout fits an eight-digit hour value");
request.headers_mut().insert(
"grpc-timeout",
http::HeaderValue::from_str(&timeout).expect("numeric gRPC timeout header"),
);
true
}
async fn send_channel_request(
channel: &mut Channel,
request: http::Request<Body>,
) -> Result<http::Response<Body>, tonic::transport::Error> {
std::future::poll_fn(|context| channel.poll_ready(context)).await?;
channel.call(request).await
}
fn is_unauthenticated(headers: &http::HeaderMap) -> bool {
headers
.get("grpc-status")
.is_some_and(|status| status == "16")
}
fn observe_stream_rejection(
response: http::Response<Body>,
credential: Arc<dyn Credential>,
rejected_token: Option<String>,
) -> http::Response<Body> {
let (parts, body) = response.into_parts();
if is_unauthenticated(&parts.headers)
&& let Some(token) = &rejected_token
{
credential.discard_rejected_token(token);
}
http::Response::from_parts(
parts,
Body::new(ObservedBody {
body,
credential,
rejected_token,
}),
)
}
struct ObservedBody {
body: Body,
credential: Arc<dyn Credential>,
rejected_token: Option<String>,
}
impl HttpBody for ObservedBody {
type Data = Bytes;
type Error = Status;
fn poll_frame(
mut self: Pin<&mut Self>,
context: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
let frame = std::task::ready!(Pin::new(&mut self.body).poll_frame(context));
if let Some(Ok(frame)) = &frame
&& frame.trailers_ref().is_some_and(is_unauthenticated)
&& let Some(token) = &self.rejected_token
{
self.credential.discard_rejected_token(token);
}
Poll::Ready(frame)
}
}
fn buffered_body(data: Bytes, trailers: Option<http::HeaderMap>) -> Body {
Body::new(BufferedBody {
data: Some(data),
trailers,
})
}
struct BufferedBody {
data: Option<Bytes>,
trailers: Option<http::HeaderMap>,
}
impl HttpBody for BufferedBody {
type Data = Bytes;
type Error = Status;
fn poll_frame(
mut self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
if let Some(data) = self.data.take() {
return Poll::Ready(Some(Ok(Frame::data(data))));
}
Poll::Ready(
self.trailers
.take()
.map(|trailers| Ok(Frame::trailers(trailers))),
)
}
}
#[derive(Clone)]
pub(crate) struct CredentialInterceptor {
credential: Option<Arc<dyn Credential>>,
}
impl CredentialInterceptor {
pub(crate) fn new(credential: Option<Arc<dyn Credential>>) -> Self {
Self { credential }
}
}
impl Interceptor for CredentialInterceptor {
fn call(&mut self, mut request: Request<()>) -> Result<Request<()>, Status> {
if let Some(credential) = &self.credential {
match credential.token() {
Ok(token) => {
let value = MetadataValue::try_from(format!("Bearer {token}"))
.map_err(|_| Status::unauthenticated("invalid bearer token metadata"))?;
request.metadata_mut().insert("authorization", value);
}
Err(_) if credential.token_after_rejection().is_some() => {}
Err(error) => return Err(Status::unauthenticated(error.to_string())),
}
}
Ok(request)
}
}
#[cfg(test)]
mod tests {
use std::convert::Infallible;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, Poll};
use dex_protocol::dex::worker_task_service_client::WorkerTaskServiceClient;
use dex_protocol::dex::worker_task_service_server::{
WorkerTaskService, WorkerTaskServiceServer,
};
use dex_protocol::dex::{
HealthInfo, PollWorkerTaskRequest, PollWorkerTaskResponse, ServerInfo,
};
use tokio::net::TcpListener;
use tokio::sync::{mpsc, oneshot};
use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream};
use tonic::codegen::{BoxFuture, Service, http};
use tonic::server::{NamedService, UnaryService};
use tonic::transport::{Endpoint, Server};
use tonic::{Response, Streaming};
use super::*;
use crate::{
BlobCache, BlobCacheConfig, Client, ClientOptions, PullWorkerOptions, Registry, Worker,
WorkerOptions,
};
#[derive(Clone)]
struct FakeFlowService {
headers: Arc<Mutex<Vec<Option<String>>>>,
}
impl NamedService for FakeFlowService {
const NAME: &'static str = "dex.FlowService";
}
impl<B> Service<http::Request<B>> for FakeFlowService
where
B: tonic::codegen::Body + Send + 'static,
B::Error: Into<tonic::codegen::StdError> + Send + 'static,
{
type Response = http::Response<tonic::body::Body>;
type Error = Infallible;
type Future = BoxFuture<Self::Response, Self::Error>;
fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, request: http::Request<B>) -> Self::Future {
let is_server_info = request.uri().path() == "/dex.FlowService/GetServerInfo";
let header = request
.headers()
.get("authorization")
.map(|value| value.to_str().expect("ASCII bearer header").to_string());
self.headers.lock().expect("headers lock").push(header);
Box::pin(async move {
if is_server_info {
let mut grpc = tonic::server::Grpc::new(tonic_prost::ProstCodec::default());
Ok(grpc.unary(GetServerInfo, request).await)
} else if request.uri().path() == "/dex.FlowService/SyncAttributeIndexes" {
Ok(Status::unimplemented("fake Attribute index synchronization").into_http())
} else {
let mut grpc = tonic::server::Grpc::new(tonic_prost::ProstCodec::default());
Ok(grpc.unary(HealthCheck, request).await)
}
})
}
}
struct HealthCheck;
impl UnaryService<()> for HealthCheck {
type Response = HealthInfo;
type Future = BoxFuture<Response<Self::Response>, Status>;
fn call(&mut self, _: Request<()>) -> Self::Future {
Box::pin(async { Ok(Response::new(HealthInfo::default())) })
}
}
struct GetServerInfo;
impl UnaryService<()> for GetServerInfo {
type Response = ServerInfo;
type Future = BoxFuture<Response<Self::Response>, Status>;
fn call(&mut self, _: Request<()>) -> Self::Future {
Box::pin(async {
Ok(Response::new(ServerInfo {
server_version: "fake".to_string(),
minimum_supported_protocol_version: 4,
current_protocol_version: 4,
..Default::default()
}))
})
}
}
#[derive(Clone)]
struct FakeTaskService {
headers: Arc<Mutex<Vec<Option<String>>>>,
}
#[tonic::async_trait]
impl WorkerTaskService for FakeTaskService {
type PollWorkerTaskStream = ReceiverStream<Result<PollWorkerTaskResponse, Status>>;
async fn poll_worker_task(
&self,
request: Request<Streaming<PollWorkerTaskRequest>>,
) -> Result<Response<Self::PollWorkerTaskStream>, Status> {
let header = request
.metadata()
.get("authorization")
.map(|value| value.to_str().expect("ASCII bearer header").to_string());
self.headers.lock().expect("headers lock").push(header);
let (_sender, receiver) = mpsc::channel(1);
Ok(Response::new(ReceiverStream::new(receiver)))
}
}
struct MutableCredential {
token: Mutex<Result<String, CredentialError>>,
calls: AtomicUsize,
}
impl Credential for MutableCredential {
fn token(&self) -> Result<String, CredentialError> {
self.calls.fetch_add(1, Ordering::SeqCst);
self.token.lock().expect("token lock").clone()
}
}
#[test]
fn file_credential_rotates_and_fails_closed() {
let directory = tempfile::tempdir().expect("token directory");
let path = directory.path().join("token");
assert!(FileCredential::new(&path).is_err());
fs::write(&path, " \n").expect("write empty token");
assert!(FileCredential::new(&path).is_err());
fs::write(&path, " first\n").expect("write initial token");
let credential = FileCredential::new(&path).expect("initial token");
assert_eq!(credential.token().expect("initial token"), "first");
fs::write(&path, " second-token \n").expect("rotate token");
assert_eq!(credential.token().expect("cached token"), "first");
credential.state.lock().expect("file state lock").checked_at -= Duration::from_secs(1);
assert_eq!(credential.token().expect("rotated token"), "second-token");
fs::write(&path, " \n").expect("empty token");
credential.state.lock().expect("file state lock").checked_at -= Duration::from_secs(1);
assert!(credential.token().is_err());
fs::remove_file(&path).expect("remove token");
credential.state.lock().expect("file state lock").checked_at -= Duration::from_secs(1);
assert!(credential.token().is_err());
fs::write(&path, "restored").expect("restore token");
credential.state.lock().expect("file state lock").checked_at -= Duration::from_secs(1);
assert_eq!(credential.token().expect("restored token"), "restored");
}
#[test]
fn file_credential_debug_redacts_token() {
let directory = tempfile::tempdir().expect("token directory");
let path = directory.path().join("token");
fs::write(&path, "private-token-value").expect("write token");
let credential = Arc::new(FileCredential::new(&path).expect("file credential"));
let debug = format!("{credential:?}");
assert!(!debug.contains("private-token-value"));
assert!(debug.contains("<redacted>"));
}
#[test]
fn flow_service_unary_uses_current_credential_and_default_has_no_header() {
let runtime = tokio::runtime::Runtime::new().expect("server runtime");
let listener = runtime
.block_on(TcpListener::bind("127.0.0.1:0"))
.expect("bind fake FlowService");
let address = listener.local_addr().expect("fake server address");
let headers = Arc::new(Mutex::new(Vec::new()));
let (shutdown, stopped) = oneshot::channel::<()>();
runtime.spawn(
Server::builder()
.add_service(FakeFlowService {
headers: Arc::clone(&headers),
})
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
stopped.await.expect("FlowService shutdown signal");
}),
);
let credential = Arc::new(MutableCredential {
token: Mutex::new(Ok("first".to_string())),
calls: AtomicUsize::new(0),
});
let directory = tempfile::tempdir().expect("blob cache directory");
let cache = Arc::new(
BlobCache::open(
BlobCacheConfig::new(directory.path(), 1024 * 1024, 0).expect("cache config"),
)
.expect("open cache"),
);
let without_credential = Client::try_new(
Registry::new(),
Arc::clone(&cache),
ClientOptions::new().server_address(address.to_string()),
)
.expect("default Client");
without_credential.health_check().expect("default call");
let with_credential = Client::try_new(
Registry::new(),
Arc::clone(&cache),
ClientOptions::new()
.server_address(address.to_string())
.credential(credential.clone()),
)
.expect("authenticated Client");
with_credential.health_check().expect("first call");
*credential.token.lock().expect("token lock") = Ok("second".to_string());
with_credential.health_check().expect("rotated call");
*credential.token.lock().expect("token lock") =
Err(CredentialError::new("token unavailable"));
let error = with_credential.health_check().expect_err("local error");
assert_eq!(
error.service_error().expect("service error").code(),
tonic::Code::Unauthenticated
);
*credential.token.lock().expect("token lock") = Ok("worker".to_string());
let worker = Worker::try_new(
Registry::new(),
cache,
WorkerOptions::new()
.server_address(address.to_string())
.credential(credential.clone())
.pull(PullWorkerOptions::new()),
)
.expect("authenticated Worker");
worker
.start()
.expect_err("fake Attribute index synchronization is unsupported");
assert_eq!(credential.calls.load(Ordering::SeqCst), 6);
assert_eq!(
*headers.lock().expect("headers lock"),
vec![
None,
None,
Some("Bearer first".to_string()),
Some("Bearer first".to_string()),
Some("Bearer second".to_string()),
Some("Bearer worker".to_string()),
Some("Bearer worker".to_string())
]
);
shutdown.send(()).expect("stop fake FlowService");
}
#[tokio::test(flavor = "multi_thread")]
async fn pull_stream_open_reevaluates_credential() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind fake WorkerTaskService");
let address = listener.local_addr().expect("fake server address");
let headers = Arc::new(Mutex::new(Vec::new()));
let (shutdown, stopped) = oneshot::channel::<()>();
let server = tokio::spawn(
Server::builder()
.add_service(WorkerTaskServiceServer::new(FakeTaskService {
headers: Arc::clone(&headers),
}))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
stopped.await.expect("WorkerTaskService shutdown signal");
}),
);
let credential = Arc::new(MutableCredential {
token: Mutex::new(Ok("first".to_string())),
calls: AtomicUsize::new(0),
});
let channel = Endpoint::from_shared(format!("http://{address}"))
.expect("endpoint")
.connect()
.await
.expect("connect fake WorkerTaskService");
let mut default_client = WorkerTaskServiceClient::new(CredentialChannel::new(
channel.clone(),
CredentialInterceptor::new(None),
));
default_client
.poll_worker_task(tokio_stream::empty::<PollWorkerTaskRequest>())
.await
.expect("open default task stream");
let mut client = WorkerTaskServiceClient::new(CredentialChannel::new(
channel,
CredentialInterceptor::new(Some(credential.clone())),
));
for token in ["first", "second"] {
*credential.token.lock().expect("token lock") = Ok(token.to_string());
client
.poll_worker_task(tokio_stream::empty::<PollWorkerTaskRequest>())
.await
.expect("open task stream");
}
*credential.token.lock().expect("token lock") =
Err(CredentialError::new("token unavailable"));
let error = client
.poll_worker_task(tokio_stream::empty::<PollWorkerTaskRequest>())
.await
.expect_err("local task stream error");
assert_eq!(error.code(), tonic::Code::Unauthenticated);
assert_eq!(credential.calls.load(Ordering::SeqCst), 3);
assert_eq!(
*headers.lock().expect("headers lock"),
vec![
None,
Some("Bearer first".to_string()),
Some("Bearer second".to_string())
]
);
shutdown.send(()).expect("stop fake WorkerTaskService");
server
.await
.expect("join fake WorkerTaskService")
.expect("serve task service");
}
}