rust_thrift_tls/
tls_socket.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::io::{self, ErrorKind, Read, Write};
19use std::net::{Shutdown, TcpStream};
20use std::sync::{Arc, Mutex};
21
22use thrift::transport::{ReadHalf, TIoChannel, WriteHalf};
23use thrift::{new_transport_error, TransportErrorKind};
24
25use rustls::StreamOwned as RusTLSStream;
26use rustls::{ClientSession, RootCertStore, ServerSession, Session};
27use webpki;
28
29use super::X509Credentials;
30
31pub type TLSStream<S> = Arc<Mutex<RusTLSStream<S, TcpStream>>>;
32
33/// Bidirectional TCP/IP channel.
34///
35pub struct TLSTTcpChannel<S>
36where
37    S: Session,
38{
39    stream: Option<TLSStream<S>>,
40    shutdown: Shutdown,
41}
42
43impl<S> TLSTTcpChannel<S>
44where
45    S: Session,
46{
47    /// Create an uninitialized `TLSTTcpChannel`.
48    ///
49    /// The returned instance must be opened using `TLSTTcpChannel::open(...)`
50    /// before it can be used.
51    pub fn new() -> TLSTTcpChannel<S> {
52        TLSTTcpChannel {
53            stream: None,
54            shutdown: Shutdown::Both,
55        }
56    }
57
58    /// Shut down this channel.
59    ///
60    /// Both send and receive halves are closed, and this instance can no
61    /// longer be used to communicate with another endpoint.
62    pub fn close(&mut self) -> thrift::Result<()> {
63        let shutdown_direction = self.shutdown;
64        self.if_set(|s| s.get_mut().shutdown(shutdown_direction))
65            .map_err(From::from)
66    }
67
68    fn if_set<F, T>(&mut self, mut stream_operation: F) -> io::Result<T>
69    where
70        F: FnMut(&mut RusTLSStream<S, TcpStream>) -> io::Result<T>,
71    {
72        if let Some(ref mut s) = self.stream {
73            stream_operation(&mut s.lock().unwrap())
74        } else {
75            Err(io::Error::new(
76                ErrorKind::NotConnected,
77                "tcp endpoint not connected",
78            ))
79        }
80    }
81}
82
83impl TLSTTcpChannel<ServerSession> {
84    /// Create a `TLSTTcpChannel` that wraps an existing `TLSStream`.
85    ///
86    /// The passed-in stream is assumed to have been opened before being wrapped
87    /// by the created `TLSTTcpChannel` instance.
88    pub fn with_stream(stream: TLSStream<ServerSession>) -> TLSTTcpChannel<ServerSession> {
89        TLSTTcpChannel {
90            stream: Some(stream),
91            shutdown: Shutdown::Both,
92        }
93    }
94}
95
96impl TLSTTcpChannel<ClientSession> {
97    /// Connect to `remote_address`, which should have the form `host:port`.
98    /// Client authentication can be enabled by passing a `rust_thrift_tls::X509Credentials`
99    /// By Default `webpki_roots::TLS_SERVER_ROOTS` is used to validate server certs
100    /// that can be overrode by passing a cusrom `rustls::RootCertStore`
101    pub fn open(
102        &mut self,
103        remote_address: &str,
104        key_pair: Option<X509Credentials>,
105        root_cert_store: Option<RootCertStore>,
106    ) -> thrift::Result<()> {
107        if self.stream.is_some() {
108            Err(new_transport_error(
109                TransportErrorKind::AlreadyOpen,
110                "TLS session connection previously opened",
111            ))
112        } else {
113            let tsap: Vec<&str> = remote_address.rsplit(':').collect();
114            if tsap.len() != 2 {
115                return Err(new_transport_error(
116                    TransportErrorKind::Unknown,
117                    format!("Invalid remote address: '{}'", remote_address),
118                ));
119            }
120
121            let dns_name = match webpki::DNSNameRef::try_from_ascii_str(tsap[1]) {
122                Ok(dns_nameref) => dns_nameref,
123                Err(e) => {
124                    return Err(new_transport_error(
125                        TransportErrorKind::Unknown,
126                        format!("Invalid DNS name: '{}'", e),
127                    ))
128                }
129            };
130            let config = super::make_tls_client_config(key_pair, root_cert_store);
131
132            let sess = ClientSession::new(&config, dns_name);
133            let sock = TcpStream::connect(remote_address).unwrap();
134            self.stream = Some(Arc::new(Mutex::new(RusTLSStream::new(sess, sock))));
135
136            Ok(())
137        }
138    }
139}
140
141impl<S> TIoChannel for TLSTTcpChannel<S>
142where
143    S: Session,
144{
145    fn split(self) -> thrift::Result<(ReadHalf<Self>, WriteHalf<Self>)>
146    where
147        Self: Sized,
148    {
149        if let Some(stream) = self.stream {
150            let read_half = ReadHalf::new(TLSTTcpChannel {
151                stream: Some(stream.clone()),
152                shutdown: Shutdown::Read,
153            });
154            let write_half = WriteHalf::new(TLSTTcpChannel {
155                stream: Some(stream),
156                shutdown: Shutdown::Write,
157            });
158            Ok((read_half, write_half))
159        } else {
160            Err(new_transport_error(
161                TransportErrorKind::Unknown,
162                "cannot clone underlying tcp stream",
163            ))
164        }
165    }
166}
167
168impl<S> Read for TLSTTcpChannel<S>
169where
170    S: Session,
171{
172    fn read(&mut self, b: &mut [u8]) -> io::Result<usize> {
173        self.if_set(|s| s.read(b))
174    }
175}
176
177impl<S> Write for TLSTTcpChannel<S>
178where
179    S: Session,
180{
181    fn write(&mut self, b: &[u8]) -> io::Result<usize> {
182        self.if_set(|s| s.write(b))
183    }
184
185    fn flush(&mut self) -> io::Result<()> {
186        self.if_set(|s| s.flush())
187    }
188}
rust_thrift_tls/tls_socket.rs

rust_thrift_tls/
tls_socket.rs
