From fddb77759f0dbaca2d9fdab534f19f2cc2266840 Mon Sep 17 00:00:00 2001
From: quininer <quininer@live.com>
Date: Fri, 23 Mar 2018 19:03:30 +0800
Subject: [PATCH] feat: handle CloseNotify alert

---
 src/futures_impl.rs | 24 +++++++++++++++++++-----
 src/lib.rs          | 27 +++++++++++++++++++++------
 2 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/src/futures_impl.rs b/src/futures_impl.rs
index 22c637c..86bc9da 100644
--- a/src/futures_impl.rs
+++ b/src/futures_impl.rs
@@ -100,13 +100,27 @@ impl<S, C> AsyncRead for TlsStream<S, C>
         C: Session
 {
     fn poll_read(&mut self, ctx: &mut Context, buf: &mut [u8]) -> Poll<usize, Error> {
-        let (io, session) = self.get_mut();
-        let mut taskio = TaskStream { io, task: ctx };
-        let mut stream = Stream::new(session, &mut taskio);
+        if self.eof {
+            return Ok(Async::Ready(0));
+        }
 
-        match io::Read::read(&mut stream, buf) {
+        // TODO nll
+        let result = {
+            let (io, session) = self.get_mut();
+            let mut taskio = TaskStream { io, task: ctx };
+            let mut stream = Stream::new(session, &mut taskio);
+            io::Read::read(&mut stream, buf)
+        };
+
+        match result {
+            Ok(0) => { self.eof = true; Ok(Async::Ready(0)) },
             Ok(n) => Ok(Async::Ready(n)),
-            Err(ref e) if e.kind() == io::ErrorKind::ConnectionAborted => Ok(Async::Ready(0)),
+            Err(ref e) if e.kind() == io::ErrorKind::ConnectionAborted => {
+                self.eof = true;
+                self.is_shutdown = true;
+                self.session.send_close_notify();
+                Ok(Async::Ready(0))
+            },
             Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Ok(Async::Pending),
             Err(e) => Err(e)
         }
diff --git a/src/lib.rs b/src/lib.rs
index 3337e0d..18b3ae2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -56,7 +56,7 @@ pub fn connect_async_with_session<S>(stream: S, session: ClientSession)
     where S: io::Read + io::Write
 {
     ConnectAsync(MidHandshake {
-        inner: Some(TlsStream { session, io: stream, is_shutdown: false })
+        inner: Some(TlsStream { session, io: stream, is_shutdown: false, eof: false })
     })
 }
 
@@ -77,7 +77,7 @@ pub fn accept_async_with_session<S>(stream: S, session: ServerSession)
     where S: io::Read + io::Write
 {
     AcceptAsync(MidHandshake {
-        inner: Some(TlsStream { session, io: stream, is_shutdown: false })
+        inner: Some(TlsStream { session, io: stream, is_shutdown: false, eof: false })
     })
 }
 
@@ -92,6 +92,7 @@ struct MidHandshake<S, C> {
 #[derive(Debug)]
 pub struct TlsStream<S, C> {
     is_shutdown: bool,
+    eof: bool,
     io: S,
     session: C
 }
@@ -112,12 +113,26 @@ impl<S, C> io::Read for TlsStream<S, C>
     where S: io::Read + io::Write, C: Session
 {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        let (io, session) = self.get_mut();
-        let mut stream = Stream::new(session, io);
+        if self.eof {
+            return Ok(0);
+        }
 
-        match stream.read(buf) {
+        // TODO nll
+        let result = {
+            let (io, session) = self.get_mut();
+            let mut stream = Stream::new(session, io);
+            stream.read(buf)
+        };
+
+        match result {
+            Ok(0) => { self.eof = true; Ok(0) },
             Ok(n) => Ok(n),
-            Err(ref e) if e.kind() == io::ErrorKind::ConnectionAborted => Ok(0),
+            Err(ref e) if e.kind() == io::ErrorKind::ConnectionAborted => {
+                self.eof = true;
+                self.is_shutdown = true;
+                self.session.send_close_notify();
+                Ok(0)
+            },
             Err(e) => Err(e)
         }
     }