atm0s_sdn_network/features/
vpn.rs

1use std::marker::PhantomData;
2
3#[cfg(feature = "vpn")]
4use crate::base::TransportMsg;
5#[cfg(feature = "vpn")]
6use atm0s_sdn_identity::{NodeId, NodeIdType};
7#[cfg(feature = "vpn")]
8use atm0s_sdn_router::{RouteAction, RouteRule, RouterTable};
9use derivative::Derivative;
10use sans_io_runtime::{collections::DynamicDeque, TaskSwitcherChild};
11
12use crate::base::{Buffer, Feature, FeatureContext, FeatureInput, FeatureOutput, FeatureWorker, FeatureWorkerContext, FeatureWorkerInput, FeatureWorkerOutput};
13
14pub const FEATURE_ID: u8 = 3;
15pub const FEATURE_NAME: &str = "vpn";
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub enum Control {}
19
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum Event {}
22
23#[derive(Debug, Clone)]
24pub struct ToWorker;
25
26#[derive(Debug, Clone)]
27pub struct ToController;
28
29pub type Output<UserData> = FeatureOutput<UserData, Event, ToWorker>;
30pub type WorkerOutput<UserData> = FeatureWorkerOutput<UserData, Control, Event, ToController>;
31
32#[derive(Debug, Derivative)]
33#[derivative(Default(bound = ""))]
34pub struct VpnFeature<UserData> {
35    _tmp: PhantomData<UserData>,
36    shutdown: bool,
37}
38
39impl<UserData> Feature<UserData, Control, Event, ToController, ToWorker> for VpnFeature<UserData> {
40    fn on_shared_input(&mut self, _ctx: &FeatureContext, _now: u64, _input: crate::base::FeatureSharedInput) {}
41
42    fn on_input(&mut self, _ctx: &FeatureContext, _now_ms: u64, _input: FeatureInput<'_, UserData, Control, ToController>) {}
43
44    fn on_shutdown(&mut self, _ctx: &FeatureContext, _now: u64) {
45        self.shutdown = true;
46    }
47}
48
49impl<UserData> TaskSwitcherChild<Output<UserData>> for VpnFeature<UserData> {
50    type Time = u64;
51
52    fn is_empty(&self) -> bool {
53        self.shutdown
54    }
55
56    fn empty_event(&self) -> Output<UserData> {
57        Output::OnResourceEmpty
58    }
59
60    fn pop_output(&mut self, _now: u64) -> Option<Output<UserData>> {
61        None
62    }
63}
64
65#[derive(Derivative)]
66#[derivative(Default(bound = ""))]
67pub struct VpnFeatureWorker<UserData> {
68    queue: DynamicDeque<WorkerOutput<UserData>, 16>,
69    shutdown: bool,
70}
71
72impl<UserData> VpnFeatureWorker<UserData> {
73    #[cfg(feature = "vpn")]
74    fn process_tun(&mut self, ctx: &FeatureWorkerContext, mut pkt: Buffer) {
75        #[cfg(any(target_os = "macos", target_os = "ios"))]
76        let to_ip = &pkt[20..24];
77        #[cfg(any(target_os = "linux", target_os = "android"))]
78        let to_ip = &pkt[16..20];
79        let dest = NodeId::build(ctx.node_id.geo1(), ctx.node_id.geo2(), ctx.node_id.group(), to_ip[3]);
80        if dest == ctx.node_id {
81            //This is for current node, just echo back
82            rewrite_tun_pkt(&mut pkt);
83            self.queue.push_back(FeatureWorkerOutput::TunPkt(pkt));
84        } else if let RouteAction::Next(remote) = ctx.router.path_to_node(dest) {
85            //TODO decrease TTL
86            //TODO how to avoid copy data here
87            self.queue
88                .push_back(FeatureWorkerOutput::RawDirect2(remote, TransportMsg::build(FEATURE_ID, 0, RouteRule::ToNode(dest), &pkt).take()));
89        }
90    }
91
92    fn process_udp(&mut self, _ctx: &FeatureWorkerContext, pkt: Buffer) {
93        #[cfg(feature = "vpn")]
94        {
95            self.queue.push_back(FeatureWorkerOutput::TunPkt(pkt));
96        }
97    }
98}
99
100impl<UserData> FeatureWorker<UserData, Control, Event, ToController, ToWorker> for VpnFeatureWorker<UserData> {
101    fn on_input(&mut self, ctx: &mut FeatureWorkerContext, _now: u64, input: FeatureWorkerInput<UserData, Control, ToWorker>) {
102        match input {
103            #[cfg(feature = "vpn")]
104            FeatureWorkerInput::TunPkt(pkt) => self.process_tun(ctx, pkt),
105            FeatureWorkerInput::Network(_conn, _header, pkt) => self.process_udp(ctx, pkt),
106            _ => {}
107        }
108    }
109
110    fn on_shutdown(&mut self, _ctx: &mut FeatureWorkerContext, _now: u64) {
111        log::info!("[VpnFeatureWorker] Shutdown");
112        self.shutdown = true;
113    }
114}
115
116impl<UserData> TaskSwitcherChild<WorkerOutput<UserData>> for VpnFeatureWorker<UserData> {
117    type Time = u64;
118
119    fn is_empty(&self) -> bool {
120        self.shutdown && self.queue.is_empty()
121    }
122
123    fn empty_event(&self) -> WorkerOutput<UserData> {
124        WorkerOutput::OnResourceEmpty
125    }
126
127    fn pop_output(&mut self, _now: u64) -> Option<WorkerOutput<UserData>> {
128        self.queue.pop_front()
129    }
130}
131
132#[cfg(feature = "vpn")]
133fn rewrite_tun_pkt(payload: &mut [u8]) {
134    #[cfg(any(target_os = "macos", target_os = "ios"))]
135    {
136        payload[2] = 0;
137        payload[3] = 2;
138    }
139    #[cfg(any(target_os = "linux", target_os = "android"))]
140    {
141        payload[2] = 8;
142        payload[3] = 0;
143    }
144}