aboutsummaryrefslogtreecommitdiff
path: root/examples/std/src/bin/tcp_accept.rs
blob: 97ce77f42681d55a67229bff5b48303d3c13422e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
#![feature(type_alias_impl_trait)]

use core::fmt::Write as _;
use std::default::Default;

use clap::Parser;
use embassy_executor::{Executor, Spawner};
use embassy_net::tcp::TcpSocket;
use embassy_net::{Config, Ipv4Address, Ipv4Cidr, Stack, StackResources};
use embassy_time::{Duration, Timer};
use embedded_io::asynch::Write as _;
use heapless::Vec;
use log::*;
use rand_core::{OsRng, RngCore};
use static_cell::StaticCell;

#[path = "../tuntap.rs"]
mod tuntap;

use crate::tuntap::TunTapDevice;

macro_rules! singleton {
    ($val:expr) => {{
        type T = impl Sized;
        static STATIC_CELL: StaticCell<T> = StaticCell::new();
        STATIC_CELL.init_with(move || $val)
    }};
}

#[derive(Parser)]
#[clap(version = "1.0")]
struct Opts {
    /// TAP device name
    #[clap(long, default_value = "tap0")]
    tap: String,
    /// use a static IP instead of DHCP
    #[clap(long)]
    static_ip: bool,
}

#[embassy_executor::task]
async fn net_task(stack: &'static Stack<TunTapDevice>) -> ! {
    stack.run().await
}

#[derive(Default)]
struct StrWrite(pub heapless::Vec<u8, 30>);

impl core::fmt::Write for StrWrite {
    fn write_str(&mut self, s: &str) -> Result<(), core::fmt::Error> {
        self.0.extend_from_slice(s.as_bytes()).unwrap();
        Ok(())
    }
}

#[embassy_executor::task]
async fn main_task(spawner: Spawner) {
    let opts: Opts = Opts::parse();

    // Init network device
    let device = TunTapDevice::new(&opts.tap).unwrap();

    // Choose between dhcp or static ip
    let config = if opts.static_ip {
        Config::Static(embassy_net::StaticConfig {
            address: Ipv4Cidr::new(Ipv4Address::new(192, 168, 69, 2), 24),
            dns_servers: Vec::new(),
            gateway: Some(Ipv4Address::new(192, 168, 69, 1)),
        })
    } else {
        Config::Dhcp(Default::default())
    };

    // Generate random seed
    let mut seed = [0; 8];
    OsRng.fill_bytes(&mut seed);
    let seed = u64::from_le_bytes(seed);

    // Init network stack
    let stack = &*singleton!(Stack::new(device, config, singleton!(StackResources::<3>::new()), seed));

    // Launch network task
    spawner.spawn(net_task(stack)).unwrap();

    // Then we can use it!
    let mut rx_buffer = [0; 4096];
    let mut tx_buffer = [0; 4096];

    loop {
        let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
        socket.set_timeout(Some(Duration::from_secs(10)));
        info!("Listening on TCP:9999...");
        if let Err(_) = socket.accept(9999).await {
            warn!("accept error");
            continue;
        }

        info!("Accepted a connection");

        // Write some quick output
        for i in 1..=5 {
            let mut w = StrWrite::default();
            write!(w, "{}!  ", i).unwrap();
            let r = socket.write_all(&w.0).await;
            if let Err(e) = r {
                warn!("write error: {:?}", e);
                return;
            }

            Timer::after(Duration::from_millis(500)).await;
        }
        info!("Closing the connection");
        socket.abort();
        info!("Flushing the RST out...");
        socket.flush().await;
        info!("Finished with the socket");
    }
}

static EXECUTOR: StaticCell<Executor> = StaticCell::new();

fn main() {
    env_logger::builder()
        .filter_level(log::LevelFilter::Debug)
        .filter_module("async_io", log::LevelFilter::Info)
        .format_timestamp_nanos()
        .init();

    let executor = EXECUTOR.init(Executor::new());
    executor.run(|spawner| {
        spawner.spawn(main_task(spawner)).unwrap();
    });
}