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 Reopen(generated::ReopenRequest),
47 Pause(generated::PauseRequest),
49 Resume(generated::ResumeRequest),
51 Rename(generated::RenameRequest),
56 MintNamespace(generated::MintNamespaceRequest),
59}
60
61#[derive(Clone, Debug)]
63pub enum ForwardReply {
64 Start(generated::StartWorkflowResponse),
66 Signal(generated::SignalResponse),
68 Query(generated::QueryResponse),
70 Cancel(generated::CancelResponse),
72 Reopen(generated::ReopenResponse),
74 Pause(generated::PauseResponse),
76 Resume(generated::ResumeResponse),
78 Rename(generated::RenameResponse),
80 MintNamespace(generated::MintNamespaceResponse),
82}
83
84#[async_trait]
91pub trait RequestForwarder: Send + Sync {
92 async fn forward(
95 &self,
96 target: SocketAddr,
97 metadata: tonic::metadata::MetadataMap,
98 request: ForwardRequest,
99 ) -> Result<ForwardReply, Status>;
100}
101
102#[derive(Clone, Default)]
105pub struct GrpcRequestForwarder;
106
107impl GrpcRequestForwarder {
108 #[must_use]
110 pub const fn new() -> Self {
111 Self
112 }
113}
114
115#[must_use]
117pub fn current_hops(metadata: &tonic::metadata::MetadataMap) -> u32 {
118 metadata
119 .get(FORWARD_HOPS_METADATA)
120 .and_then(|value| value.to_str().ok())
121 .and_then(|value| value.parse().ok())
122 .unwrap_or(0)
123}
124
125fn stamp_next_hop(metadata: &mut tonic::metadata::MetadataMap) -> Result<(), Status> {
127 let next = current_hops(metadata)
128 .checked_add(1)
129 .ok_or_else(|| Status::internal("forward hop counter overflow"))?;
130 let value = tonic::metadata::MetadataValue::try_from(next.to_string())
131 .map_err(|_| Status::internal("invalid forward hop metadata value"))?;
132 metadata.insert(FORWARD_HOPS_METADATA, value);
133 Ok(())
134}
135
136async fn connect(target: SocketAddr) -> Result<WorkflowServiceClient<Channel>, Status> {
137 let uri = format!("http://{target}");
138 let endpoint = Endpoint::try_from(uri)
139 .map_err(|error| Status::unavailable(format!("invalid forward target: {error}")))?;
140 let channel = endpoint
141 .connect()
142 .await
143 .map_err(|error| Status::unavailable(format!("forward dial failed: {error}")))?;
144 Ok(WorkflowServiceClient::new(channel))
145}
146
147#[async_trait]
148impl RequestForwarder for GrpcRequestForwarder {
149 async fn forward(
150 &self,
151 target: SocketAddr,
152 mut metadata: tonic::metadata::MetadataMap,
153 request: ForwardRequest,
154 ) -> Result<ForwardReply, Status> {
155 stamp_next_hop(&mut metadata)?;
156 let mut client = connect(target).await?;
157 match request {
158 ForwardRequest::Start(message) => {
159 let mut outbound = Request::new(message);
160 *outbound.metadata_mut() = metadata;
161 client
162 .start_workflow(outbound)
163 .await
164 .map(|response| ForwardReply::Start(response.into_inner()))
165 }
166 ForwardRequest::Signal(message) => {
167 let mut outbound = Request::new(message);
168 *outbound.metadata_mut() = metadata;
169 client
170 .signal(outbound)
171 .await
172 .map(|response| ForwardReply::Signal(response.into_inner()))
173 }
174 ForwardRequest::Query(message) => {
175 let mut outbound = Request::new(message);
176 *outbound.metadata_mut() = metadata;
177 client
178 .query(outbound)
179 .await
180 .map(|response| ForwardReply::Query(response.into_inner()))
181 }
182 ForwardRequest::Cancel(message) => {
183 let mut outbound = Request::new(message);
184 *outbound.metadata_mut() = metadata;
185 client
186 .cancel(outbound)
187 .await
188 .map(|response| ForwardReply::Cancel(response.into_inner()))
189 }
190 ForwardRequest::Reopen(message) => {
191 let mut outbound = Request::new(message);
192 *outbound.metadata_mut() = metadata;
193 client
194 .reopen(outbound)
195 .await
196 .map(|response| ForwardReply::Reopen(response.into_inner()))
197 }
198 ForwardRequest::Pause(message) => {
199 let mut outbound = Request::new(message);
200 *outbound.metadata_mut() = metadata;
201 client
202 .pause(outbound)
203 .await
204 .map(|response| ForwardReply::Pause(response.into_inner()))
205 }
206 ForwardRequest::Resume(message) => {
207 let mut outbound = Request::new(message);
208 *outbound.metadata_mut() = metadata;
209 client
210 .resume(outbound)
211 .await
212 .map(|response| ForwardReply::Resume(response.into_inner()))
213 }
214 ForwardRequest::Rename(message) => {
215 let mut outbound = Request::new(message);
216 *outbound.metadata_mut() = metadata;
217 client
218 .rename(outbound)
219 .await
220 .map(|response| ForwardReply::Rename(response.into_inner()))
221 }
222 ForwardRequest::MintNamespace(message) => {
223 let mut outbound = Request::new(message);
224 *outbound.metadata_mut() = metadata;
225 client
226 .mint_namespace(outbound)
227 .await
228 .map(|response| ForwardReply::MintNamespace(response.into_inner()))
229 }
230 }
231 }
232}
233
234#[cfg(test)]
235mod tests {
236 use super::{FORWARD_HOPS_METADATA, MAX_FORWARD_HOPS, current_hops, stamp_next_hop};
237
238 #[test]
239 fn current_hops_defaults_to_zero_when_absent() {
240 let metadata = tonic::metadata::MetadataMap::new();
241 assert_eq!(current_hops(&metadata), 0);
242 }
243
244 #[test]
245 fn stamp_next_hop_increments_from_zero() -> Result<(), tonic::Status> {
246 let mut metadata = tonic::metadata::MetadataMap::new();
247 stamp_next_hop(&mut metadata)?;
248 assert_eq!(current_hops(&metadata), 1);
249 stamp_next_hop(&mut metadata)?;
250 assert_eq!(current_hops(&metadata), 2);
251 Ok(())
252 }
253
254 #[test]
255 fn malformed_hop_value_reads_as_zero() -> Result<(), tonic::Status> {
256 let mut metadata = tonic::metadata::MetadataMap::new();
257 metadata.insert(
258 FORWARD_HOPS_METADATA,
259 tonic::metadata::MetadataValue::try_from("not-a-number")
260 .map_err(|_| tonic::Status::internal("bad fixture"))?,
261 );
262 assert_eq!(current_hops(&metadata), 0);
263 Ok(())
264 }
265
266 #[test]
267 fn hop_cap_is_two() {
268 assert_eq!(MAX_FORWARD_HOPS, 2);
269 }
270}