use std::future::Future; use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll, Wake, Waker}; use url::Url; use crate::credssp::ServerError; use crate::network_client::{AsyncNetworkClient, NetworkClient, NetworkProtocol}; use crate::{AcceptSecurityContextResult, Error, InitializeSecurityContextResult}; pub struct Interrupt { value_to_yield: Option, yielded_value: YieldedValue, resumed_value: ResumedValue, ready_to_resume: bool, } impl Future for Interrupt where YieldTy: Unpin, { type Output = ResumeTy; fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll { let this = self.get_mut(); if this.ready_to_resume { let resumed_value = this.resumed_value.try_lock().unwrap().take().unwrap(); Poll::Ready(resumed_value) } else { let value_to_yield = this.value_to_yield.take().unwrap(); *this.yielded_value.try_lock().unwrap() = Some(value_to_yield); this.ready_to_resume = true; Poll::Pending } } } #[derive(Debug)] pub struct YieldPoint { yielded_value: YieldedValue, resumed_value: ResumedValue, } impl Clone for YieldPoint { fn clone(&self) -> Self { Self { yielded_value: self.yielded_value.clone(), resumed_value: self.resumed_value.clone(), } } } impl YieldPoint { pub fn suspend(&mut self, value: YieldTy) -> Interrupt { Interrupt { value_to_yield: Some(value), yielded_value: Arc::clone(&self.yielded_value), resumed_value: Arc::clone(&self.resumed_value), ready_to_resume: false, } } } type YieldedValue = Arc>>; type ResumedValue = Arc>>; type PinnedFuture<'a, T> = Pin + Send + 'a>>; pub enum GeneratorState { Suspended(YieldTy), Completed(OutTy), } pub struct Generator<'a, YieldTy, ResumeTy, OutTy> { yielded_value: YieldedValue, resumed_value: ResumedValue, generator: PinnedFuture<'a, OutTy>, } impl<'a, YieldTy, ResumeTy, OutTy> Generator<'a, YieldTy, ResumeTy, OutTy> where OutTy: Send + 'a, { pub fn new(producer: Producer) -> Self where Producer: FnOnce(YieldPoint) -> Task, Task: Future + Send + 'a, { let yielded_value = Arc::new(Mutex::new(None)); let resumed_value = Arc::new(Mutex::new(None)); let yield_point = YieldPoint { yielded_value: Arc::clone(&yielded_value), resumed_value: Arc::clone(&resumed_value), }; Self { yielded_value, resumed_value, generator: Box::pin(producer(yield_point)), } } pub fn start(&mut self) -> GeneratorState { self.step() } pub fn resume(&mut self, value: ResumeTy) -> GeneratorState { *self.resumed_value.try_lock().unwrap() = Some(value); self.step() } fn step(&mut self) -> GeneratorState { match execute_one_step(&mut self.generator) { None => { let value = self.yielded_value.try_lock().unwrap().take().unwrap(); GeneratorState::Suspended(value) } Some(value) => GeneratorState::Completed(value), } } } fn execute_one_step(task: &mut PinnedFuture<'_, OutTy>) -> Option { struct NoopWake; impl Wake for NoopWake { fn wake(self: Arc) { // do nothing } } let waker = Waker::from(Arc::new(NoopWake)); let mut context = Context::from_waker(&waker); match task.as_mut().poll(&mut context) { Poll::Pending => None, Poll::Ready(item) => Some(item), } } /// Utility types and methods impl<'a, YieldTy, ResumeTy, OutTy> Generator<'a, YieldTy, ResumeTy, Result> where OutTy: Send + 'a, { pub fn resolve_to_result(&mut self) -> Result { let state = self.start(); match state { GeneratorState::Suspended(_) => Err(Error::new( crate::ErrorKind::UnsupportedFunction, "cannot finish generator", )), GeneratorState::Completed(res) => res, } } pub fn unwrap(&mut self) -> OutTy { self.resolve_to_result().unwrap() } pub fn expect(&mut self, msg: &str) -> OutTy { self.resolve_to_result().expect(msg) } } impl<'a, YieldTy, ResumeTy, OutTy> Generator<'a, YieldTy, ResumeTy, Result> where OutTy: Send + 'a, { pub fn resolve_to_result(&mut self) -> Result { let state = self.start(); match state { GeneratorState::Suspended(_) => Err(ServerError { ts_request: None, error: Error::new(crate::ErrorKind::UnsupportedFunction, "cannot finish generator"), }), GeneratorState::Completed(res) => res, } } } #[derive(Debug, Clone)] pub struct NetworkRequest { pub protocol: NetworkProtocol, pub url: Url, pub data: Vec, // avoid life time problem, suspend requires 'static life time } impl std::fmt::Debug for Generator<'_, YieldTy, ResumeTy, OutTy> where YieldTy: std::fmt::Debug, ResumeTy: std::fmt::Debug, { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Generator") .field("yielded_value", &self.yielded_value) .field("resumed_value", &self.resumed_value) .finish() } } pub type GeneratorInitSecurityContext<'a> = Generator<'a, NetworkRequest, crate::Result>, crate::Result>; pub type GeneratorAcceptSecurityContext<'a> = Generator<'a, NetworkRequest, crate::Result>, crate::Result>; pub type GeneratorChangePassword<'a> = Generator<'a, NetworkRequest, crate::Result>, crate::Result<()>>; pub(crate) type YieldPointLocal = YieldPoint>>; impl<'a, YieldType, ResumeType, OutType, ErrorType> From> for Generator<'a, YieldType, ResumeType, Result> where OutType: Send + 'a, ErrorType: Send + 'a, { fn from(value: Result) -> Self { Generator::new(move |_| async move { value }) } } /// Utilities for working with network client impl<'a, OutTy> Generator<'a, NetworkRequest, crate::Result>, OutTy> where OutTy: 'a + Send, { #[cfg(feature = "network_client")] pub fn resolve_with_default_network_client(&mut self) -> OutTy { let network_client = crate::network_client::reqwest_network_client::ReqwestNetworkClient; self.resolve_with_client(&network_client) } pub fn resolve_with_client(&mut self, network_client: &dyn NetworkClient) -> OutTy { let mut state = self.start(); loop { match state { GeneratorState::Suspended(ref request) => { state = self.resume(NetworkClient::send(network_client, request)); } GeneratorState::Completed(res) => { return res; } } } } pub async fn resolve_with_async_client(&mut self, network_client: &mut dyn AsyncNetworkClient) -> OutTy { let mut state = self.start(); loop { match state { GeneratorState::Suspended(ref request) => { state = self.resume(network_client.send(request).await); } GeneratorState::Completed(client_state) => { return client_state; } } } } }