aion_server/routing/
forwarder.rs1use std::net::SocketAddr;
13
14use async_trait::async_trait;
15use tonic::transport::{Channel, Endpoint};
16use tonic::{Request, Status};
17
18use aion_proto::generated::{self, workflow_service_client::WorkflowServiceClient};
19
20pub const FORWARD_HOPS_METADATA: &str = "x-aion-forward-hops";
25
26pub const MAX_FORWARD_HOPS: u32 = 2;
30
31#[derive(Clone, Debug)]
36pub enum ForwardRequest {
37 Start(generated::StartWorkflowRequest),
39 Signal(generated::SignalRequest),
41 Query(generated::QueryRequest),
43 Cancel(generated::CancelRequest),
45 RetireWorkloop(generated::RetireWorkloopRequest),
50 Reopen(generated::ReopenRequest),
52 Pause(generated::PauseRequest),
54 Resume(generated::ResumeRequest),
56 Rename(generated::RenameRequest),
61 MintNamespace(generated::MintNamespaceRequest),
64}
65
66#[derive(Clone, Debug)]
68pub enum ForwardReply {
69 Start(generated::StartWorkflowResponse),
71 Signal(generated::SignalResponse),
73 Query(generated::QueryResponse),
75 Cancel(generated::CancelResponse),
77 RetireWorkloop(generated::RetireWorkloopResponse),
79 Reopen(generated::ReopenResponse),
81 Pause(generated::PauseResponse),
83 Resume(generated::ResumeResponse),
85 Rename(generated::RenameResponse),
87 MintNamespace(generated::MintNamespaceResponse),
89}
90
91#[async_trait]
98pub trait RequestForwarder: Send + Sync {
99 async fn forward(
102 &self,
103 target: SocketAddr,
104 metadata: tonic::metadata::MetadataMap,
105 request: ForwardRequest,
106 ) -> Result<ForwardReply, Status>;
107}
108
109#[derive(Clone, Default)]
112pub struct GrpcRequestForwarder;
113
114impl GrpcRequestForwarder {
115 #[must_use]
117 pub const fn new() -> Self {
118 Self
119 }
120}
121
122#[must_use]
124pub fn current_hops(metadata: &tonic::metadata::MetadataMap) -> u32 {
125 metadata
126 .get(FORWARD_HOPS_METADATA)
127 .and_then(|value| value.to_str().ok())
128 .and_then(|value| value.parse().ok())
129 .unwrap_or(0)
130}
131
132fn stamp_next_hop(metadata: &mut tonic::metadata::MetadataMap) -> Result<(), Status> {
134 let next = current_hops(metadata)
135 .checked_add(1)
136 .ok_or_else(|| Status::internal("forward hop counter overflow"))?;
137 let value = tonic::metadata::MetadataValue::try_from(next.to_string())
138 .map_err(|_| Status::internal("invalid forward hop metadata value"))?;
139 metadata.insert(FORWARD_HOPS_METADATA, value);
140 Ok(())
141}
142
143async fn connect(target: SocketAddr) -> Result<WorkflowServiceClient<Channel>, Status> {
144 let uri = format!("http://{target}");
145 let endpoint = Endpoint::try_from(uri)
146 .map_err(|error| Status::unavailable(format!("invalid forward target: {error}")))?;
147 let channel = endpoint
148 .connect()
149 .await
150 .map_err(|error| Status::unavailable(format!("forward dial failed: {error}")))?;
151 Ok(WorkflowServiceClient::new(channel))
152}
153
154#[async_trait]
155impl RequestForwarder for GrpcRequestForwarder {
156 async fn forward(
157 &self,
158 target: SocketAddr,
159 mut metadata: tonic::metadata::MetadataMap,
160 request: ForwardRequest,
161 ) -> Result<ForwardReply, Status> {
162 stamp_next_hop(&mut metadata)?;
163 let mut client = connect(target).await?;
164 match request {
165 ForwardRequest::Start(message) => {
166 let mut outbound = Request::new(message);
167 *outbound.metadata_mut() = metadata;
168 client
169 .start_workflow(outbound)
170 .await
171 .map(|response| ForwardReply::Start(response.into_inner()))
172 }
173 ForwardRequest::Signal(message) => {
174 let mut outbound = Request::new(message);
175 *outbound.metadata_mut() = metadata;
176 client
177 .signal(outbound)
178 .await
179 .map(|response| ForwardReply::Signal(response.into_inner()))
180 }
181 ForwardRequest::Query(message) => {
182 let mut outbound = Request::new(message);
183 *outbound.metadata_mut() = metadata;
184 client
185 .query(outbound)
186 .await
187 .map(|response| ForwardReply::Query(response.into_inner()))
188 }
189 ForwardRequest::Cancel(message) => {
190 let mut outbound = Request::new(message);
191 *outbound.metadata_mut() = metadata;
192 client
193 .cancel(outbound)
194 .await
195 .map(|response| ForwardReply::Cancel(response.into_inner()))
196 }
197 ForwardRequest::Reopen(message) => {
198 let mut outbound = Request::new(message);
199 *outbound.metadata_mut() = metadata;
200 client
201 .reopen(outbound)
202 .await
203 .map(|response| ForwardReply::Reopen(response.into_inner()))
204 }
205 ForwardRequest::Pause(message) => {
206 let mut outbound = Request::new(message);
207 *outbound.metadata_mut() = metadata;
208 client
209 .pause(outbound)
210 .await
211 .map(|response| ForwardReply::Pause(response.into_inner()))
212 }
213 ForwardRequest::Resume(message) => {
214 let mut outbound = Request::new(message);
215 *outbound.metadata_mut() = metadata;
216 client
217 .resume(outbound)
218 .await
219 .map(|response| ForwardReply::Resume(response.into_inner()))
220 }
221 ForwardRequest::Rename(message) => {
222 let mut outbound = Request::new(message);
223 *outbound.metadata_mut() = metadata;
224 client
225 .rename(outbound)
226 .await
227 .map(|response| ForwardReply::Rename(response.into_inner()))
228 }
229 ForwardRequest::RetireWorkloop(message) => {
230 let mut outbound = Request::new(message);
231 *outbound.metadata_mut() = metadata;
232 client
233 .retire_workloop(outbound)
234 .await
235 .map(|response| ForwardReply::RetireWorkloop(response.into_inner()))
236 }
237 ForwardRequest::MintNamespace(message) => {
238 let mut outbound = Request::new(message);
239 *outbound.metadata_mut() = metadata;
240 client
241 .mint_namespace(outbound)
242 .await
243 .map(|response| ForwardReply::MintNamespace(response.into_inner()))
244 }
245 }
246 }
247}
248
249#[cfg(test)]
250mod tests {
251 use super::{FORWARD_HOPS_METADATA, MAX_FORWARD_HOPS, current_hops, stamp_next_hop};
252
253 #[test]
254 fn current_hops_defaults_to_zero_when_absent() {
255 let metadata = tonic::metadata::MetadataMap::new();
256 assert_eq!(current_hops(&metadata), 0);
257 }
258
259 #[test]
260 fn stamp_next_hop_increments_from_zero() -> Result<(), tonic::Status> {
261 let mut metadata = tonic::metadata::MetadataMap::new();
262 stamp_next_hop(&mut metadata)?;
263 assert_eq!(current_hops(&metadata), 1);
264 stamp_next_hop(&mut metadata)?;
265 assert_eq!(current_hops(&metadata), 2);
266 Ok(())
267 }
268
269 #[test]
270 fn malformed_hop_value_reads_as_zero() -> Result<(), tonic::Status> {
271 let mut metadata = tonic::metadata::MetadataMap::new();
272 metadata.insert(
273 FORWARD_HOPS_METADATA,
274 tonic::metadata::MetadataValue::try_from("not-a-number")
275 .map_err(|_| tonic::Status::internal("bad fixture"))?,
276 );
277 assert_eq!(current_hops(&metadata), 0);
278 Ok(())
279 }
280
281 #[test]
282 fn hop_cap_is_two() {
283 assert_eq!(MAX_FORWARD_HOPS, 2);
284 }
285}