1
use std::{net::SocketAddr, path::Path};
2

            
3
use echo_protocol::EchoChannel;
4
use jetstream::prelude::*;
5
use jetstream_macros::service;
6
use jetstream_rpc::Framed;
7
use okstd::prelude::*;
8
use s2n_quic::{client::Connect, provider::tls, Client, Server};
9

            
10
#[service]
11
pub trait Echo {
12
    async fn ping(&mut self) -> Result<(), Error>;
13
}
14

            
15
struct EchoImpl {}
16

            
17
impl Echo for EchoImpl {
18
    async fn ping(&mut self) -> Result<(), Error> {
19
        eprintln!("Ping received");
20
        eprintln!("Pong sent");
21
        Ok(())
22
    }
23
}
24

            
25
pub static CA_CERT_PEM: &str =
26
    concat!(env!("CARGO_MANIFEST_DIR"), "/certs/ca-cert.pem");
27
pub static CLIENT_CERT_PEM: &str =
28
    concat!(env!("CARGO_MANIFEST_DIR"), "/certs/client-cert.pem");
29
pub static CLIENT_KEY_PEM: &str =
30
    concat!(env!("CARGO_MANIFEST_DIR"), "/certs/client-key.pem");
31
pub static SERVER_CERT_PEM: &str =
32
    concat!(env!("CARGO_MANIFEST_DIR"), "/certs/server-cert.pem");
33
pub static SERVER_KEY_PEM: &str =
34
    concat!(env!("CARGO_MANIFEST_DIR"), "/certs/server-key.pem");
35

            
36
async fn server() -> Result<(), Box<dyn std::error::Error>> {
37
    let tls = tls::default::Server::builder()
38
        .with_trusted_certificate(Path::new(CA_CERT_PEM))?
39
        .with_certificate(
40
            Path::new(SERVER_CERT_PEM),
41
            Path::new(SERVER_KEY_PEM),
42
        )?
43
        .with_client_authentication()?
44
        .build()?;
45

            
46
    let mut server = Server::builder()
47
        .with_tls(tls)?
48
        .with_io("127.0.0.1:4433")?
49
        .start()?;
50

            
51
    while let Some(mut connection) = server.accept().await {
52
        // spawn a new task for the connection
53
        tokio::spawn(async move {
54
            eprintln!(
55
                "Connection accepted from {:?}",
56
                connection.remote_addr()
57
            );
58

            
59
            while let Ok(Some(stream)) =
60
                connection.accept_bidirectional_stream().await
61
            {
62
                // spawn a new task for the stream
63
                tokio::spawn(async move {
64
                    eprintln!(
65
                        "Stream opened from {:?}",
66
                        &stream.connection().remote_addr()
67
                    );
68
                    let echo = EchoImpl {};
69
                    let servercodec: jetstream::prelude::server::service::ServerCodec<
70
                        echo_protocol::EchoService<EchoImpl>,
71
                    > = Default::default();
72
                    let framed = Framed::new(stream, servercodec);
73
                    let mut serv = echo_protocol::EchoService { inner: echo };
74
                    server::service::run(&mut serv, framed).await.unwrap();
75
                });
76
            }
77
        });
78
    }
79

            
80
    Ok(())
81
}
82

            
83
async fn client() -> Result<(), Box<dyn std::error::Error>> {
84
    let tls = tls::default::Client::builder()
85
        .with_certificate(Path::new(CA_CERT_PEM))?
86
        .with_client_identity(
87
            Path::new(CLIENT_CERT_PEM),
88
            Path::new(CLIENT_KEY_PEM),
89
        )?
90
        .build()?;
91

            
92
    let client = Client::builder()
93
        .with_tls(tls)?
94
        .with_io("0.0.0.0:0")?
95
        .start()?;
96

            
97
    let addr: SocketAddr = "127.0.0.1:4433".parse()?;
98
    let connect = Connect::new(addr).with_server_name("localhost");
99
    let mut connection = client.connect(connect).await?;
100

            
101
    // ensure the connection doesn't time out with inactivity
102
    connection.keep_alive(true)?;
103

            
104
    // open a new stream and split the receiving and sending sides
105
    let stream = connection.open_bidirectional_stream().await?;
106
    let client_codec: jetstream_client::ClientCodec<EchoChannel> =
107
        Default::default();
108
    let mut framed = Framed::new(stream, client_codec);
109
    let mut chan = EchoChannel {
110
        inner: Box::new(&mut framed),
111
    };
112
    eprintln!("Ping sent");
113
    chan.ping().await?;
114
    eprintln!("Pong received");
115
    Ok(())
116
}
117

            
118
#[okstd::main]
119
async fn main() {
120
    tokio::select! {
121
      _ = server() => {},
122
      _ = client() => {},
123
    }
124
}