Move context creation to fn of RootContext instead of callback

Signed-off-by: Daniel Grimm <dgrimm@redhat.com>

Author: dgn

examples/http_auth_random.rs

@@ -20,10 +20,23 @@ use std::time::Duration;
 #[no_mangle]
 pub fn _start() {
     proxy_wasm::set_log_level(LogLevel::Trace);
-    proxy_wasm::set_http_context(|_, _| -> Box<dyn HttpContext> { Box::new(HttpAuthRandom) });
+    proxy_wasm::set_root_context(|_| -> Box<dyn RootContext> { Box::new(HttpAuthRandomRoot) });
 }
 
 struct HttpAuthRandom;
+struct HttpAuthRandomRoot;
+
+impl RootContext for HttpAuthRandomRoot {
+    fn get_type(&self) -> ContextType {
+        ContextType::HttpContext
+    }
+
+    fn create_http_context(&self, _root_context_id: u32, _context_id: u32) -> Box<dyn HttpContext> {
+        Box::new(HttpAuthRandom)
+    }
+}
+
+impl Context for HttpAuthRandomRoot {}
 
 impl HttpContext for HttpAuthRandom {
     fn on_http_request_headers(&mut self, _: usize) -> Action {

examples/http_body.rs

@@ -18,10 +18,24 @@ use proxy_wasm::types::*;
 #[no_mangle]
 pub fn _start() {
     proxy_wasm::set_log_level(LogLevel::Trace);
-    proxy_wasm::set_http_context(|_, _| -> Box<dyn HttpContext> { Box::new(HttpBody) });
+    proxy_wasm::set_root_context(|_| -> Box<dyn RootContext> { Box::new(HttpBodyRoot) });
 }
 
 struct HttpBody;
+struct HttpBodyRoot;
+
+impl RootContext for HttpBodyRoot {
+    fn get_type(&self) -> ContextType {
+        ContextType::HttpContext
+    }
+
+    fn create_http_context(&self, _context_id: u32, _root_context_id: u32) -> Box<dyn HttpContext> {
+        Box::new(HttpBody)
+    }
+}
+
+impl Context for HttpBodyRoot {}
+
 
 impl Context for HttpBody {}
 

examples/http_headers.rs

@@ -19,11 +19,25 @@ use proxy_wasm::types::*;
 #[no_mangle]
 pub fn _start() {
     proxy_wasm::set_log_level(LogLevel::Trace);
-    proxy_wasm::set_http_context(|context_id, _| -> Box<dyn HttpContext> {
-        Box::new(HttpHeaders { context_id })
+    proxy_wasm::set_root_context(|_| -> Box<dyn RootContext> {
+        Box::new(HttpHeadersRoot)
     });
 }
 
+struct HttpHeadersRoot;
+
+impl Context for HttpHeadersRoot {}
+
+impl RootContext for HttpHeadersRoot {
+    fn get_type(&self) -> ContextType {
+        ContextType::HttpContext
+    }
+
+    fn create_http_context(&self, _context_id: u32, _root_context_id: u32) -> Box<dyn HttpContext> {
+        Box::new(HttpHeaders{context_id: _context_id})
+    }
+}
+
 struct HttpHeaders {
     context_id: u32,
 }

src/dispatcher.rs

@@ -26,14 +26,6 @@ pub(crate) fn set_root_context(callback: NewRootContext) {
     DISPATCHER.with(|dispatcher| dispatcher.set_root_context(callback));
 }
 
-pub(crate) fn set_stream_context(callback: NewStreamContext) {
-    DISPATCHER.with(|dispatcher| dispatcher.set_stream_context(callback));
-}
-
-pub(crate) fn set_http_context(callback: NewHttpContext) {
-    DISPATCHER.with(|dispatcher| dispatcher.set_http_context(callback));
-}
-
 pub(crate) fn register_callout(token_id: u32) {
     DISPATCHER.with(|dispatcher| dispatcher.register_callout(token_id));
 }
@@ -46,9 +38,7 @@ impl RootContext for NoopRoot {}
 struct Dispatcher {
     new_root: Cell<Option<NewRootContext>>,
     roots: RefCell<HashMap<u32, Box<dyn RootContext>>>,
-    new_stream: Cell<Option<NewStreamContext>>,
     streams: RefCell<HashMap<u32, Box<dyn StreamContext>>>,
-    new_http_stream: Cell<Option<NewHttpContext>>,
     http_streams: RefCell<HashMap<u32, Box<dyn HttpContext>>>,
     active_id: Cell<u32>,
     callouts: RefCell<HashMap<u32, u32>>,
@@ -59,9 +49,7 @@ impl Dispatcher {
         Dispatcher {
             new_root: Cell::new(None),
             roots: RefCell::new(HashMap::new()),
-            new_stream: Cell::new(None),
             streams: RefCell::new(HashMap::new()),
-            new_http_stream: Cell::new(None),
             http_streams: RefCell::new(HashMap::new()),
             active_id: Cell::new(0),
             callouts: RefCell::new(HashMap::new()),
@@ -72,14 +60,6 @@ impl Dispatcher {
         self.new_root.set(Some(callback));
     }
 
-    fn set_stream_context(&self, callback: NewStreamContext) {
-        self.new_stream.set(Some(callback));
-    }
-
-    fn set_http_context(&self, callback: NewHttpContext) {
-        self.new_http_stream.set(Some(callback));
-    }
-
     fn create_root_context(&self, context_id: u32) {
         let new_context = match self.new_root.get() {
             Some(f) => f(context_id),
@@ -96,12 +76,9 @@ impl Dispatcher {
     }
 
     fn create_stream_context(&self, context_id: u32, root_context_id: u32) {
-        if !self.roots.borrow().contains_key(&root_context_id) {
-            panic!("invalid root_context_id")
-        }
-        let new_context = match self.new_stream.get() {
-            Some(f) => f(context_id, root_context_id),
-            None => panic!("missing constructor"),
+        let new_context = match self.roots.borrow().get(&root_context_id) {
+            Some(root_context) => root_context.create_stream_context(context_id, root_context_id),
+            None => panic!("invalid root_context_id"),
         };
         if self
             .streams
@@ -114,12 +91,9 @@ impl Dispatcher {
     }
 
     fn create_http_context(&self, context_id: u32, root_context_id: u32) {
-        if !self.roots.borrow().contains_key(&root_context_id) {
-            panic!("invalid root_context_id")
-        }
-        let new_context = match self.new_http_stream.get() {
-            Some(f) => f(context_id, root_context_id),
-            None => panic!("missing constructor"),
+        let new_context = match self.roots.borrow().get(&root_context_id) {
+            Some(root_context) => root_context.create_http_context(context_id, root_context_id),
+            None => panic!("invalid root_context_id"),
         };
         if self
             .http_streams
@@ -144,11 +118,13 @@ impl Dispatcher {
 
     fn on_create_context(&self, context_id: u32, root_context_id: u32) {
         if root_context_id == 0 {
-            self.create_root_context(context_id)
-        } else if self.new_http_stream.get().is_some() {
-            self.create_http_context(context_id, root_context_id);
-        } else if self.new_stream.get().is_some() {
-            self.create_stream_context(context_id, root_context_id);
+            self.create_root_context(context_id);
+        } else if let Some(root_context) = self.roots.borrow().get(&root_context_id) {
+            match root_context.get_type() {
+                ContextType::HttpContext => self.create_http_context(context_id, root_context_id),
+                ContextType::StreamContext => self.create_stream_context(context_id, root_context_id),
+                ContextType::RootContext => panic!("missing constructors"),
+            }
         } else {
             panic!("missing constructors")
         }

src/lib.rs

@@ -28,13 +28,5 @@ pub fn set_root_context(callback: types::NewRootContext) {
     dispatcher::set_root_context(callback);
 }
 
-pub fn set_stream_context(callback: types::NewStreamContext) {
-    dispatcher::set_stream_context(callback);
-}
-
-pub fn set_http_context(callback: types::NewHttpContext) {
-    dispatcher::set_http_context(callback);
-}
-
 #[no_mangle]
 pub extern "C" fn proxy_abi_version_0_1_0() {}

src/traits.rs

@@ -121,6 +121,18 @@ pub trait RootContext: Context {
     fn on_queue_ready(&mut self, _queue_id: u32) {}
 
     fn on_log(&mut self) {}
+
+    fn create_http_context(&self, _context_id: u32, _root_context_id: u32) -> Box<dyn HttpContext> {
+        Box::new(EmptyHttpContext)
+    }
+
+    fn create_stream_context(&self, _context_id: u32, _root_context_id: u32) -> Box<dyn StreamContext> {
+        Box::new(EmptyStreamContext)
+    }
+
+    fn get_type(&self) -> ContextType {
+        ContextType::RootContext
+    }
 }
 
 pub trait StreamContext: Context {
@@ -159,6 +171,16 @@ pub trait StreamContext: Context {
     fn on_log(&mut self) {}
 }
 
+struct EmptyHttpContext;
+
+impl HttpContext for EmptyHttpContext {}
+impl Context for EmptyHttpContext {}
+
+struct EmptyStreamContext;
+
+impl StreamContext for EmptyStreamContext {}
+impl Context for EmptyStreamContext {}
+
 pub trait HttpContext: Context {
     fn on_http_request_headers(&mut self, _num_headers: usize) -> Action {
         Action::Continue

src/types.rs

@@ -85,3 +85,9 @@ pub enum MetricType {
 }
 
 pub type Bytes = Vec<u8>;
+
+pub enum ContextType {
+    RootContext = 0,
+    HttpContext = 1,
+    StreamContext = 2,
+}