1 // A Rustls stub for TryTLS
2 //
3 // Author: Joachim Viide
4 // See: https://github.com/HowNetWorks/trytls-rustls-stub
5 //
6
7 use webpki_roots;
8
9 use rustls::{ClientConfig, ClientConnection, Error, OwnedTrustAnchor, RootCertStore};
10 use std::convert::TryInto;
11 use std::env;
12 use std::error::Error as StdError;
13 use std::fs::File;
14 use std::io::{BufReader, Read, Write};
OpenStream(streamID uint32, options OpenStreamOptions)15 use std::net::TcpStream;
16 use std::process;
17 use std::sync::Arc;
18
19 enum Verdict {
CloseStream(streamID uint32)20 Accept,
21 Reject(Error),
22 }
23
24 fn parse_args(args: &[String]) -> Result<(String, u16, ClientConfig), Box<dyn StdError>> {
25 let mut root_store = RootCertStore::empty();
26 match args.len() {
27 3 => {
28 root_store.add_server_trust_anchors(
29 webpki_roots::TLS_SERVER_ROOTS
30 .0
31 .iter()
32 .map(|ta| {
33 OwnedTrustAnchor::from_subject_spki_name_constraints(
34 ta.subject,
35 ta.spki,
36 ta.name_constraints,
37 )
38 }),
39 );
40 }
41 4 => {
42 let f = File::open(&args[3])?;
43 root_store
44 .add_parsable_certificates(&rustls_pemfile::certs(&mut BufReader::new(f)).unwrap());
45 }
46 _ => {
47 return Err(From::from("Incorrect number of arguments"));
48 }
49 };
50 let config = rustls::ClientConfig::builder()
51 .with_safe_defaults()
52 .with_root_certificates(root_store)
53 .with_no_client_auth();
54
55 let port = args[2].parse()?;
56 Ok((args[1].clone(), port, config))
57 }
58
59 fn communicate(
60 host: String,
61 port: u16,
62 config: ClientConfig,
63 ) -> Result<Verdict, Box<dyn StdError>> {
64 let server_name = host.as_str().try_into().unwrap();
65 let rc_config = Arc::new(config);
66 let mut client = ClientConnection::new(rc_config, server_name).unwrap();
67 let mut stream = TcpStream::connect((&*host, port))?;
68
69 client
70 .writer()
71 .write_all(b"GET / HTTP/1.0\r\nConnection: close\r\nContent-Length: 0\r\n\r\n")?;
72 loop {
73 while client.wants_write() {
74 client.write_tls(&mut stream)?;
75 }
76
77 if client.wants_read() {
78 if client.read_tls(&mut stream)? == 0 {
79 return Err(From::from("Connection closed"));
80 }
81
82 if let Err(err) = client.process_new_packets() {
83 return match err {
84 Error::InvalidCertificateData(_)
85 | Error::InvalidCertificateSignature
86 | Error::InvalidCertificateSignatureType
87 | Error::InvalidCertificateEncoding
88 | Error::AlertReceived(_) => Ok(Verdict::Reject(err)),
89 _ => Err(From::from(format!("{:?}", err))),
90 };
91 }
92
93 if client.reader().read(&mut [0])? > 0 {
94 return Ok(Verdict::Accept);
95 }
96 }
97 }
98 }
99
100 fn main() {
101 let args: Vec<String> = env::args().collect();
102 let (host, port, config) = parse_args(&args).unwrap_or_else(|err| {
103 println!("Argument error: {}", err);
104 process::exit(2);
105 });
106
107 match communicate(host, port, config) {
108 Ok(Verdict::Accept) => {
109 println!("ACCEPT");
110 process::exit(0);
111 }
112 Ok(Verdict::Reject(reason)) => {
113 println!("{:?}", reason);
114 println!("REJECT");
115 process::exit(0);
116 }
117 Err(err) => {
118 println!("{}", err);
119 process::exit(1);
120 }
121 }
122 }
123