1use ozzy_proto::{Envelope, Opcode};
4use ozzy_proto::{LinkSessionId, NodeId, RequestId};
5
6use super::{
7 COMMON_BYTES, ControlEncoding, Incoming, PrepareBatch, ReadFields, Reader, WireError,
8 WireLimits, WriteFields, Writer, validate_prefix, validate_scope,
9};
10use crate::OpNumber;
11use crate::flow::{Channel, ReceiveEpoch, Report};
12
13pub use crate::flow::Probe as FlowProbe;
14
15const PROBE_BYTES: usize = COMMON_BYTES + 40 + 8;
16const STATE_BYTES: usize = COMMON_BYTES + 16 + 8 + 40 + 40 + 8 + 8 + 4;
17
18fn validate_probe(probe: FlowProbe) -> Result<(), WireError> {
19 validate_scope(probe.scope)?;
20 validate_prefix(probe.tail)?;
21 if probe.available < probe.tail.op || probe.available.0 == u64::MAX {
22 return Err(WireError::Prefix);
23 }
24 if probe.request_id.as_bytes() == &[0; 16] {
25 return Err(WireError::Correlation);
26 }
27 Ok(())
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub struct FlowState {
33 pub handle: u32,
35 pub request_id: Option<RequestId>,
37 pub report: Report,
39 pub repair_limit: Option<OpNumber>,
42}
43
44impl FlowState {
45 fn validate(self) -> Result<(), WireError> {
46 self.report.validate_shape()?;
47 if self.handle == 0 {
48 return Err(WireError::Correlation);
49 }
50 if self.repair_limit == Some(OpNumber(u64::MAX)) {
51 return Err(WireError::Prefix);
52 }
53 validate_prefix(self.report.base)?;
54 validate_prefix(self.report.received)?;
55 if self.request_id.is_some_and(|id| id.as_bytes() == &[0; 16]) {
56 return Err(WireError::Correlation);
57 }
58 Ok(())
59 }
60
61 pub fn validate_response(self, probe: FlowProbe) -> Result<(), WireError> {
67 self.validate()?;
68 validate_probe(probe)?;
69 if self.request_id != Some(probe.request_id) || self.report.channel.scope != probe.scope {
70 return Err(WireError::Correlation);
71 }
72 Ok(())
73 }
74}
75
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub enum FlowMessage<'a> {
79 Probe(FlowProbe),
81 State(FlowState),
83 Prepare {
85 epoch: ReceiveEpoch,
87 batch: PrepareBatch<'a>,
89 },
90}
91
92pub fn encode_flow_probe(
94 sender: NodeId,
95 session: LinkSessionId,
96 probe: FlowProbe,
97 output: &mut [u8],
98 limits: WireLimits,
99) -> Result<ControlEncoding, WireError> {
100 validate_probe(probe)?;
101 let header = Envelope {
102 opcode: Opcode::ReplicaOpen,
103 response: false,
104 request_id: Some(probe.request_id),
105 sender,
106 session: Some(session),
107 }
108 .encode_header(PROBE_BYTES, 0, limits.envelope)?;
109 if output.len() < PROBE_BYTES {
110 return Err(WireError::Capacity);
111 }
112 let mut writer = Writer::with_len(output, PROBE_BYTES)?;
113 writer.scope(probe.scope, sender)?;
114 writer.prefix(probe.tail)?;
115 writer.u64(probe.available.0)?;
116 Ok(ControlEncoding {
117 header,
118 metadata_bytes: PROBE_BYTES,
119 })
120}
121
122pub fn encode_flow_state(
125 sender: NodeId,
126 session: LinkSessionId,
127 state: FlowState,
128 output: &mut [u8],
129 limits: WireLimits,
130) -> Result<ControlEncoding, WireError> {
131 state.validate()?;
132 let header = Envelope {
133 opcode: Opcode::ReplicaState,
134 response: state.request_id.is_some(),
135 request_id: state.request_id,
136 sender,
137 session: Some(session),
138 }
139 .encode_header(STATE_BYTES, 0, limits.envelope)?;
140 if output.len() < STATE_BYTES {
141 return Err(WireError::Capacity);
142 }
143 let mut writer = Writer::with_len(output, STATE_BYTES)?;
144 let report = state.report;
145 writer.scope(report.channel.scope, sender)?;
146 writer.bytes(&report.channel.epoch.get().to_be_bytes())?;
147 writer.u64(report.revision)?;
148 writer.prefix(report.base)?;
149 writer.prefix(report.received)?;
150 writer.u64(report.received_bytes)?;
151 writer.u64(state.repair_limit.map_or(u64::MAX, |op| op.0))?;
152 writer.bytes(&state.handle.to_be_bytes())?;
153 Ok(ControlEncoding {
154 header,
155 metadata_bytes: STATE_BYTES,
156 })
157}
158
159pub(super) fn decode_probe(
160 Incoming {
161 scope,
162 envelope,
163 mut reader,
164 payload,
165 }: Incoming<'_>,
166) -> Result<FlowProbe, WireError> {
167 if !payload.is_empty() {
168 return Err(WireError::Payload);
169 }
170 if envelope.response {
171 return Err(WireError::Correlation);
172 }
173 let probe = FlowProbe {
174 scope,
175 request_id: envelope.request_id.ok_or(WireError::Correlation)?,
176 tail: reader.prefix()?,
177 available: OpNumber(reader.u64()?),
178 };
179 reader.finish()?;
180 validate_probe(probe)?;
181 Ok(probe)
182}
183
184pub(super) fn decode_state(
185 Incoming {
186 scope,
187 envelope,
188 mut reader,
189 payload,
190 }: Incoming<'_>,
191) -> Result<FlowState, WireError> {
192 if !payload.is_empty() {
193 return Err(WireError::Payload);
194 }
195 if envelope.response != envelope.request_id.is_some() {
196 return Err(WireError::Correlation);
197 }
198 let epoch = ReceiveEpoch::new(u128::from_be_bytes(reader.array()?))?;
199 let state = FlowState {
200 handle: 1,
201 request_id: envelope.request_id,
202 report: Report {
203 channel: Channel { scope, epoch },
204 revision: reader.u64()?,
205 base: reader.prefix()?,
206 received: reader.prefix()?,
207 received_bytes: reader.u64()?,
208 },
209 repair_limit: match reader.u64()? {
210 u64::MAX => None,
211 value => Some(OpNumber(value)),
212 },
213 };
214 let state = FlowState {
215 handle: reader.u32()?,
216 ..state
217 };
218 reader.finish()?;
219 state.validate()?;
220 Ok(state)
221}
222
223pub fn flow_state_route(
226 packet: ozzy_proto::Packet<'_>,
227 limits: ozzy_proto::EnvelopeLimits,
228) -> Result<FlowState, WireError> {
229 packet
230 .envelope
231 .validate_frames(packet.metadata.len(), packet.payload.len(), limits)?;
232 if packet.envelope.opcode != Opcode::ReplicaState {
233 return Err(WireError::UnsupportedCommand);
234 }
235 let mut reader = Reader::new(packet.metadata);
236 decode_state(Incoming {
237 scope: reader.routing_scope(packet.envelope.sender)?,
238 envelope: packet.envelope,
239 reader,
240 payload: packet.payload,
241 })
242}