From e34b5c08ba280f2a8318b2ed607309e41cb9f4d6 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Fri, 3 Aug 2018 19:24:53 -0700
Subject: [PATCH] allow to pass extra information from acceptor to application
 level

---
 src/server/h1.rs      |  4 ++++
 src/server/h2.rs      | 11 +++++++++--
 src/server/message.rs |  8 ++++++++
 src/server/mod.rs     |  7 +++++++
 4 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/server/h1.rs b/src/server/h1.rs
index 085cea005..2c07f0cf4 100644
--- a/src/server/h1.rs
+++ b/src/server/h1.rs
@@ -368,6 +368,10 @@ where
                         self.payload = Some(PayloadType::new(&msg.inner.headers, ps));
                     }
 
+                    // stream extensions
+                    msg.inner_mut().stream_extensions =
+                        self.stream.get_mut().extensions();
+
                     // set remote addr
                     msg.inner_mut().addr = self.addr;
 
diff --git a/src/server/h2.rs b/src/server/h2.rs
index cb5367c5e..9f0725022 100644
--- a/src/server/h2.rs
+++ b/src/server/h2.rs
@@ -14,6 +14,7 @@ use tokio_io::{AsyncRead, AsyncWrite};
 use tokio_timer::Delay;
 
 use error::{Error, PayloadError};
+use extensions::Extensions;
 use http::{StatusCode, Version};
 use payload::{Payload, PayloadStatus, PayloadWriter};
 use uri::Url;
@@ -22,7 +23,7 @@ use super::error::ServerError;
 use super::h2writer::H2Writer;
 use super::input::PayloadType;
 use super::settings::WorkerSettings;
-use super::{HttpHandler, HttpHandlerTask, Writer};
+use super::{HttpHandler, HttpHandlerTask, IoStream, Writer};
 
 bitflags! {
     struct Flags: u8 {
@@ -42,6 +43,7 @@ where
     state: State<IoWrapper<T>>,
     tasks: VecDeque<Entry<H>>,
     keepalive_timer: Option<Delay>,
+    extensions: Option<Rc<Extensions>>,
 }
 
 enum State<T: AsyncRead + AsyncWrite> {
@@ -52,12 +54,13 @@ enum State<T: AsyncRead + AsyncWrite> {
 
 impl<T, H> Http2<T, H>
 where
-    T: AsyncRead + AsyncWrite + 'static,
+    T: IoStream + 'static,
     H: HttpHandler + 'static,
 {
     pub fn new(
         settings: Rc<WorkerSettings<H>>, io: T, addr: Option<SocketAddr>, buf: Bytes,
     ) -> Self {
+        let extensions = io.extensions();
         Http2 {
             flags: Flags::empty(),
             tasks: VecDeque::new(),
@@ -68,6 +71,7 @@ where
             keepalive_timer: None,
             addr,
             settings,
+            extensions,
         }
     }
 
@@ -206,6 +210,7 @@ where
                                 resp,
                                 self.addr,
                                 &self.settings,
+                                self.extensions.clone(),
                             ));
                         }
                         Ok(Async::NotReady) => {
@@ -324,6 +329,7 @@ impl<H: HttpHandler + 'static> Entry<H> {
     fn new(
         parts: Parts, recv: RecvStream, resp: SendResponse<Bytes>,
         addr: Option<SocketAddr>, settings: &Rc<WorkerSettings<H>>,
+        extensions: Option<Rc<Extensions>>,
     ) -> Entry<H>
     where
         H: HttpHandler + 'static,
@@ -338,6 +344,7 @@ impl<H: HttpHandler + 'static> Entry<H> {
             inner.method = parts.method;
             inner.version = parts.version;
             inner.headers = parts.headers;
+            inner.stream_extensions = extensions;
             *inner.payload.borrow_mut() = Some(payload);
             inner.addr = addr;
         }
diff --git a/src/server/message.rs b/src/server/message.rs
index 395d7b7c3..43f7e1425 100644
--- a/src/server/message.rs
+++ b/src/server/message.rs
@@ -35,6 +35,7 @@ pub(crate) struct InnerRequest {
     pub(crate) info: RefCell<ConnectionInfo>,
     pub(crate) payload: RefCell<Option<Payload>>,
     pub(crate) settings: ServerSettings,
+    pub(crate) stream_extensions: Option<Rc<Extensions>>,
     pool: &'static RequestPool,
 }
 
@@ -82,6 +83,7 @@ impl Request {
                 info: RefCell::new(ConnectionInfo::default()),
                 payload: RefCell::new(None),
                 extensions: RefCell::new(Extensions::new()),
+                stream_extensions: None,
             }),
         }
     }
@@ -189,6 +191,12 @@ impl Request {
         }
     }
 
+    /// Io stream extensions
+    #[inline]
+    pub fn stream_extensions(&self) -> Option<&Extensions> {
+        self.inner().stream_extensions.as_ref().map(|e| e.as_ref())
+    }
+
     /// Server settings
     #[inline]
     pub fn server_settings(&self) -> &ServerSettings {
diff --git a/src/server/mod.rs b/src/server/mod.rs
index 55de25db4..baf004926 100644
--- a/src/server/mod.rs
+++ b/src/server/mod.rs
@@ -1,5 +1,6 @@
 //! Http server
 use std::net::Shutdown;
+use std::rc::Rc;
 use std::{io, net, time};
 
 use bytes::{BufMut, BytesMut};
@@ -36,6 +37,7 @@ pub use self::helpers::write_content_length;
 use actix::Message;
 use body::Binary;
 use error::Error;
+use extensions::Extensions;
 use header::ContentEncoding;
 use httpresponse::HttpResponse;
 
@@ -287,6 +289,11 @@ pub trait IoStream: AsyncRead + AsyncWrite + 'static {
             }
         }
     }
+
+    /// Extra io stream extensions
+    fn extensions(&self) -> Option<Rc<Extensions>> {
+        None
+    }
 }
 
 impl IoStream for TcpStream {