1use std::{
2 collections::{HashMap, VecDeque},
3 fmt::Debug,
4 ops::Deref,
5};
6
7use atm0s_sdn_identity::NodeId;
8use atm0s_sdn_router::RouteRule;
9use sans_io_runtime::{collections::DynamicDeque, return_if_none, TaskSwitcherChild};
10
11use crate::base::{
12 Buffer, Feature, FeatureContext, FeatureControlActor, FeatureInput, FeatureOutput, FeatureSharedInput, FeatureWorker, FeatureWorkerContext, FeatureWorkerInput, FeatureWorkerOutput,
13 NetOutgoingMeta, Ttl,
14};
15
16pub const FEATURE_ID: u8 = 7;
17pub const FEATURE_NAME: &str = "socket";
18
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub enum Control {
21 Bind(u16),
22 Connect(u16, NodeId, u16),
23 SendTo(u16, NodeId, u16, Buffer, u8),
24 Send(u16, Buffer, u8),
25 Unbind(u16),
26}
27
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub enum Event {
30 RecvFrom(u16, NodeId, u16, Buffer, u8),
31}
32
33#[derive(Debug, Clone)]
34pub enum ToWorker<UserData> {
35 BindSocket(u16, FeatureControlActor<UserData>),
36 ConnectSocket(u16, NodeId, u16),
37 UnbindSocket(u16),
38}
39
40#[derive(Debug, Clone)]
41pub struct ToController;
42
43struct Socket<UserData> {
44 target: Option<(NodeId, u16)>,
45 actor: FeatureControlActor<UserData>,
46}
47
48pub type Output<UserData> = FeatureOutput<UserData, Event, ToWorker<UserData>>;
49pub type WorkerOutput<UserData> = FeatureWorkerOutput<UserData, Control, Event, ToController>;
50
51pub struct SocketFeature<UserData> {
52 sockets: HashMap<u16, Socket<UserData>>,
53 queue: VecDeque<Output<UserData>>,
54 shutdown: bool,
55}
56
57impl<UserData> SocketFeature<UserData> {
58 fn send_to(&mut self, src: u16, dest_node: NodeId, dest_port: u16, mut data: Buffer, meta: u8) {
59 embed_meta(src, dest_port, &mut data);
60 let meta: NetOutgoingMeta = NetOutgoingMeta::new(true, Default::default(), meta, false);
61 self.queue.push_back(FeatureOutput::SendRoute(RouteRule::ToNode(dest_node), meta, data));
62 }
63}
64
65impl<UserData> Default for SocketFeature<UserData> {
66 fn default() -> Self {
67 Self {
68 sockets: HashMap::new(),
69 queue: VecDeque::new(),
70 shutdown: false,
71 }
72 }
73}
74
75impl<UserData: Copy + Debug + Eq> Feature<UserData, Control, Event, ToController, ToWorker<UserData>> for SocketFeature<UserData> {
76 fn on_shared_input(&mut self, _ctx: &FeatureContext, _now: u64, _input: FeatureSharedInput) {}
77
78 fn on_input(&mut self, ctx: &FeatureContext, _now_ms: u64, input: FeatureInput<'_, UserData, Control, ToController>) {
79 match input {
80 FeatureInput::Control(actor, control) => match control {
81 Control::Bind(port) => {
82 if self.sockets.contains_key(&port) {
83 log::warn!("[SocketFeature] Bind failed, port already in use: {}", port);
84 return;
85 }
86 self.sockets.insert(port, Socket { target: None, actor });
87 self.queue.push_back(FeatureOutput::ToWorker(true, ToWorker::BindSocket(port, actor)));
88 }
89 Control::Connect(port, dest_node, dest_port) => {
90 if let Some(socket) = self.sockets.get_mut(&port) {
91 if socket.actor == actor {
92 socket.target = Some((dest_node, dest_port));
93 self.queue.push_back(FeatureOutput::ToWorker(true, ToWorker::ConnectSocket(port, dest_node, dest_port)));
94 } else {
95 log::warn!("[SocketFeature] Connect failed, actor mismatch: {:?} != {:?}", socket.actor, actor);
96 }
97 } else {
98 log::warn!("[SocketFeature] Connect failed, port not found: {}", port);
99 }
100 }
101 Control::SendTo(port, dest_node, dest_port, data, meta) => {
102 if let Some(socket) = self.sockets.get(&port) {
103 if socket.actor == actor {
104 if dest_node == ctx.node_id {
105 if self.sockets.contains_key(&dest_port) {
106 self.queue.push_back(FeatureOutput::Event(actor, Event::RecvFrom(dest_port, ctx.node_id, port, data, meta)));
107 } else {
108 log::warn!("[SocketFeature] SendTo failed, port not found: {}", dest_port);
109 }
110 } else {
111 self.send_to(port, dest_node, dest_port, data, meta);
112 }
113 } else {
114 log::warn!("[SocketFeature] SendTo failed, actor mismatch: {:?} != {:?}", socket.actor, actor);
115 }
116 } else {
117 log::warn!("[SocketFeature] SendTo failed, port not found: {}", port);
118 }
119 }
120 Control::Send(port, data, meta) => {
121 if let Some(socket) = self.sockets.get(&port) {
122 if let Some((dest_node, dest_port)) = socket.target {
123 if dest_node == ctx.node_id {
124 if self.sockets.contains_key(&dest_port) {
125 self.queue.push_back(FeatureOutput::Event(actor, Event::RecvFrom(dest_port, ctx.node_id, port, data, meta)));
126 } else {
127 log::warn!("[SocketFeature] SendTo failed, port not found: {}", dest_port);
128 }
129 } else {
130 self.send_to(port, dest_node, dest_port, data, meta);
131 }
132 } else {
133 log::warn!("[SocketFeature] Send failed, target not found: {}", port);
134 }
135 } else {
136 log::warn!("[SocketFeature] Send failed, port not found: {}", port);
137 }
138 }
139 Control::Unbind(port) => {
140 if let Some(socket) = self.sockets.get(&port) {
141 if socket.actor == actor {
142 self.sockets.remove(&port);
143 self.queue.push_back(FeatureOutput::ToWorker(true, ToWorker::UnbindSocket(port)));
144 } else {
145 log::warn!("[SocketFeature] Unbind failed, actor mismatch: {:?} != {:?}", socket.actor, actor);
146 }
147 } else {
148 log::warn!("[SocketFeature] Unbind failed, port not found: {}", port);
149 }
150 }
151 },
152 FeatureInput::Net(_, meta, mut buf) | FeatureInput::Local(meta, mut buf) => {
153 let from_node = if let Some(source) = meta.source {
154 source
155 } else {
156 log::warn!("[SocketFeature] Recv failed, source not set");
157 return;
158 };
159 let (pkt_src, pkt_dest) = if let Some(res) = extract_meta(&mut buf) {
160 res
161 } else {
162 log::warn!("[SocketFeature] Recv failed, invalid data");
163 return;
164 };
165 if let Some(socket) = self.sockets.get(&pkt_dest) {
166 if let Some((dest_node, dest_port)) = socket.target {
167 if dest_node != from_node {
168 log::warn!("[SocketFeature] Recv failed, node mismatch: {} != {}", dest_node, from_node);
169 return;
170 }
171 if dest_port != pkt_dest {
172 log::warn!("[SocketFeature] Recv failed, port mismatch: {} != {}", dest_port, pkt_dest);
173 return;
174 }
175 self.queue.push_back(FeatureOutput::Event(socket.actor, Event::RecvFrom(pkt_dest, from_node, pkt_src, buf, meta.meta)));
176 } else {
177 self.queue.push_back(FeatureOutput::Event(socket.actor, Event::RecvFrom(pkt_dest, from_node, pkt_src, buf, meta.meta)));
178 }
179 } else {
180 log::warn!("[SocketFeature] Recv failed, port not found: {}", pkt_dest);
181 }
182 }
183 _ => {}
184 }
185 }
186
187 fn on_shutdown(&mut self, _ctx: &FeatureContext, _now: u64) {
188 self.shutdown = true;
189 }
190}
191
192impl<UserData> TaskSwitcherChild<Output<UserData>> for SocketFeature<UserData> {
193 type Time = u64;
194
195 fn is_empty(&self) -> bool {
196 self.shutdown && self.queue.is_empty()
197 }
198
199 fn empty_event(&self) -> Output<UserData> {
200 Output::OnResourceEmpty
201 }
202
203 fn pop_output(&mut self, _now: u64) -> Option<Output<UserData>> {
204 self.queue.pop_front()
205 }
206}
207
208pub struct SocketFeatureWorker<UserData> {
209 sockets: HashMap<u16, Socket<UserData>>,
210 queue: DynamicDeque<WorkerOutput<UserData>, 16>,
211 shutdown: bool,
212}
213
214impl<UserData: Copy> SocketFeatureWorker<UserData> {
215 fn process_incoming(&mut self, from_node: NodeId, mut buf: Buffer, meta: u8) {
216 let (pkt_src, pkt_dest) = return_if_none!(extract_meta(&mut buf));
217 let socket = return_if_none!(self.sockets.get(&pkt_dest));
218 if let Some((dest_node, dest_port)) = socket.target {
219 if dest_node != from_node {
220 log::warn!("[SocketFeature] Recv failed, node mismatch: {} != {}", dest_node, from_node);
221 return;
222 }
223 if dest_port != pkt_dest {
224 log::warn!("[SocketFeature] Recv failed, port mismatch: {} != {}", dest_port, pkt_dest);
225 return;
226 }
227 self.queue.push_back(FeatureWorkerOutput::Event(socket.actor, Event::RecvFrom(pkt_dest, from_node, pkt_src, buf, meta)));
228 } else {
229 self.queue.push_back(FeatureWorkerOutput::Event(socket.actor, Event::RecvFrom(pkt_dest, from_node, pkt_src, buf, meta)))
230 }
231 }
232}
233
234impl<UserData> Default for SocketFeatureWorker<UserData> {
235 fn default() -> Self {
236 Self {
237 sockets: HashMap::new(),
238 queue: Default::default(),
239 shutdown: false,
240 }
241 }
242}
243
244impl<UserData: Clone + Copy + Eq> FeatureWorker<UserData, Control, Event, ToController, ToWorker<UserData>> for SocketFeatureWorker<UserData> {
245 fn on_input(&mut self, _ctx: &mut FeatureWorkerContext, _now: u64, input: FeatureWorkerInput<UserData, Control, ToWorker<UserData>>) {
246 match input {
247 FeatureWorkerInput::Network(_conn, meta, buf) => {
248 let from_node = return_if_none!(meta.source);
249 self.process_incoming(from_node, buf, meta.meta);
250 }
251 FeatureWorkerInput::FromController(_, control) => match control {
252 ToWorker::BindSocket(port, actor) => {
253 log::info!("[SocketFeatureWorker] BindSocket: {port}");
254 self.sockets.insert(port, Socket { target: None, actor });
255 }
256 ToWorker::ConnectSocket(port, dest_node, dest_port) => {
257 log::info!("[SocketFeatureWorker] ConnectSocket: {port} => {dest_node}:{dest_port}");
258 if let Some(socket) = self.sockets.get_mut(&port) {
259 socket.target = Some((dest_node, dest_port));
260 }
261 }
262 ToWorker::UnbindSocket(port) => {
263 log::info!("[SocketFeatureWorker] UnbindSocket: {port}");
264 self.sockets.remove(&port);
265 }
266 },
267 FeatureWorkerInput::Control(actor, control) => {
268 let (port, (dest_node, dest_port), mut data, meta) = match control {
269 Control::Send(port, data, meta) => {
270 let socket = return_if_none!(self.sockets.get(&port));
271 if actor == socket.actor {
272 let target = return_if_none!(socket.target);
273 (port, target, data, meta)
274 } else {
275 return;
276 }
277 }
278 Control::SendTo(port, dest_node, dest_port, data, meta) => {
279 let socket = return_if_none!(self.sockets.get(&port));
280 if actor == socket.actor {
281 (port, (dest_node, dest_port), data, meta)
282 } else {
283 return;
284 }
285 }
286 _ => {
287 self.queue.push_back(FeatureWorkerOutput::ForwardControlToController(actor, control));
288 return;
289 }
290 };
291
292 embed_meta(port, dest_port, &mut data);
293 let outgoing_meta = NetOutgoingMeta::new(true, Ttl::default(), meta, false);
294 self.queue.push_back(FeatureWorkerOutput::SendRoute(RouteRule::ToNode(dest_node), outgoing_meta, data));
295 }
296 FeatureWorkerInput::Local(meta, buf) => {
297 let from_node = return_if_none!(meta.source);
298 self.process_incoming(from_node, buf, meta.meta);
299 }
300 _ => {}
301 }
302 }
303
304 fn on_shutdown(&mut self, _ctx: &mut FeatureWorkerContext, _now: u64) {
305 log::info!("[SocketFeatureWorker] Shutdown");
306 self.shutdown = true;
307 }
308}
309
310impl<UserData> TaskSwitcherChild<WorkerOutput<UserData>> for SocketFeatureWorker<UserData> {
311 type Time = u64;
312
313 fn is_empty(&self) -> bool {
314 self.shutdown && self.queue.is_empty()
315 }
316
317 fn empty_event(&self) -> WorkerOutput<UserData> {
318 WorkerOutput::OnResourceEmpty
319 }
320
321 fn pop_output(&mut self, _now: u64) -> Option<WorkerOutput<UserData>> {
322 self.queue.pop_front()
323 }
324}
325
326fn embed_meta(src: u16, dest: u16, data: &mut Buffer) {
327 data.ensure_front(4);
328 data.push_front(&dest.to_be_bytes());
329 data.push_front(&src.to_be_bytes());
330}
331
332fn extract_meta(buf: &mut Buffer) -> Option<(u16, u16)> {
333 let src_buf2 = buf.pop_front(2)?;
334 let src_buf = src_buf2.deref();
335 let src = u16::from_be_bytes([src_buf[0], src_buf[1]]);
336
337 let dest_buf2 = buf.pop_front(2)?;
338 let dest_buf = dest_buf2.deref();
339 let dest = u16::from_be_bytes([dest_buf[0], dest_buf[1]]);
340
341 Some((src, dest))
342}