examples/tls-echo-server/src/main.rs

   1 use async_std::{io, task};
   2 use connect::{tls::TlsListener, ConnectDatagram, SinkExt, StreamExt};
   3 use log::*;
   4 use rustls::{Certificate, NoClientAuth, PrivateKey, ServerConfig};
   5 use rustls_pemfile::{certs, rsa_private_keys};
   6 use std::env;
   7 use std::fs::File;
   8 use std::io::BufReader;
   9
  10 #[async_std::main]
  11 async fn main() -> anyhow::Result<()> {
  12     env_logger::init();
  13
  14     // Get ip address from cmd line args
  15     let (ip_address, cert_path, key_path) = parse_args();
  16
  17     let certs = certs(&mut BufReader::new(File::open(cert_path)?))
  18         .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid cert"))?;
  19     assert!(certs.len() > 0);
  20
  21     let rustls_certs: Vec<Certificate> = certs.into_iter().map(|v| Certificate(v)).collect();
  22     assert!(rustls_certs.len() > 0);
  23
  24     let mut keys = rsa_private_keys(&mut BufReader::new(File::open(key_path)?))
  25         .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid key"))?;
  26
  27     let mut config = ServerConfig::new(NoClientAuth::new());
  28     config
  29         .set_single_cert(rustls_certs, PrivateKey(keys.remove(0)))
  30         .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?;
  31
  32     // create a server
  33     let mut server = TlsListener::bind(ip_address, config.into()).await?;
  34
  35     // handle server connections
  36     // wait for a connection to come in and be accepted
  37     while let Some(mut conn) = server.next().await {
  38         info!("Handling connection from {}", conn.peer_addr());
  39
  40         task::spawn(async move {
  41             while let Some(envelope) = conn.reader().next().await {
  42                 // handle message based on intended recipient
  43                 if envelope.tag() == 65535 {
  44                     // if recipient is 65535, we do custom processing
  45                     let data = envelope.data().to_vec();
  46                     let msg =
  47                         String::from_utf8(data).expect("could not build String from payload bytes");
  48                     info!("Received a message \"{}\" from {}", msg, conn.peer_addr());
  49
  50                     let reply = ConnectDatagram::with_tag(envelope.tag(), msg.into_bytes())
  51                         .expect("could not construct new datagram from built String");
  52
  53                     conn.writer()
  54                         .send(reply)
  55                         .await
  56                         .expect("Could not send message back to source connection");
  57                     info!("Sent message back to original sender");
  58                 } else {
  59                     // if recipient is anything else, we just send it back
  60                     warn!(
  61                         "Received a message for unknown recipient {} from {}",
  62                         envelope.tag(),
  63                         conn.peer_addr()
  64                     );
  65
  66                     conn.writer()
  67                         .send(envelope)
  68                         .await
  69                         .expect("Could not send message back to source connection");
  70                     info!("Sent message back to original sender");
  71                 }
  72             }
  73         });
  74     }
  75
  76     Ok(())
  77 }
  78
  79 fn parse_args() -> (String, String, String) {
  80     let args: Vec<String> = env::args().collect();
  81
  82     let ip_address = match args.get(1) {
  83         Some(addr) => addr,
  84         None => {
  85             error!("Need to pass IP address to connect to as first command line argument");
  86             panic!();
  87         }
  88     };
  89
  90     let cert_path = match args.get(2) {
  91         Some(d) => d,
  92         None => {
  93             error!("Need to pass path to cert file as second command line argument");
  94             panic!();
  95         }
  96     };
  97
  98     let key_path = match args.get(3) {
  99         Some(d) => d,
 100         None => {
 101             error!("Need to pass path to key file as third command line argument");
 102             panic!();
 103         }
 104     };
 105
 106     (
 107         ip_address.to_string(),
 108         cert_path.to_string(),
 109         key_path.to_string(),
 110     )
 111 }