Add backwards compatibility, use Option<> where appropriate

proxy-wasm · Nov 20, 2020 · 05d21f3 · 05d21f3
1 parent 9e2e8e1
commit 05d21f3
Show file tree

Hide file tree

Showing 8 changed files with 73 additions and 60 deletions.
diff --git a/examples/http_auth_random.rs b/examples/http_auth_random.rs
@@ -20,21 +20,7 @@ use std::time::Duration;
 #[no_mangle]
 pub fn _start() {
     proxy_wasm::set_log_level(LogLevel::Trace);
-    proxy_wasm::set_root_context(|_| -> Box<dyn RootContext> { Box::new(HttpAuthRandomRoot) });
-}
-
-struct HttpAuthRandomRoot;
-
-impl Context for HttpAuthRandomRoot {}
-
-impl RootContext for HttpAuthRandomRoot {
-    fn get_type(&self) -> ContextType {
-        ContextType::HttpContext
-    }
-
-    fn create_http_context(&self, _context_id: u32) -> Box<dyn HttpContext> {
-        Box::new(HttpAuthRandom)
-    }
+    proxy_wasm::set_http_context(|_, _| -> Box<dyn HttpContext> { Box::new(HttpAuthRandom) });
 }
 
 struct HttpAuthRandom;

diff --git a/examples/http_body.rs b/examples/http_body.rs
@@ -26,12 +26,12 @@ struct HttpBodyRoot;
 impl Context for HttpBodyRoot {}
 
 impl RootContext for HttpBodyRoot {
-    fn get_type(&self) -> ContextType {
-        ContextType::HttpContext
+    fn get_type(&self) -> Option<ContextType> {
+        Some(ContextType::HttpContext)
     }
 
-    fn create_http_context(&self, _context_id: u32) -> Box<dyn HttpContext> {
-        Box::new(HttpBody)
+    fn create_http_context(&self, _context_id: u32) -> Option<Box<dyn HttpContext>> {
+        Some(Box::new(HttpBody))
     }
 }
 

diff --git a/examples/http_config.rs b/examples/http_config.rs
@@ -49,25 +49,21 @@ impl Context for HttpConfigHeaderRootContext {}
 
 impl RootContext for HttpConfigHeaderRootContext {
 
-    fn on_vm_start(&mut self, _vm_configuration_size: usize) -> bool {
-        true
-    }
-
     fn on_configure(&mut self, _plugin_configuration_size: usize) -> bool {
         if let Some(config_bytes) = self.get_configuration() {
             self.header_content = str::from_utf8(config_bytes.as_ref()).unwrap().to_owned()
         }
         true
     }
 
-    fn create_http_context(&self, _context_id: u32) -> Box<dyn HttpContext> {
-        Box::new(HttpConfigHeader{
+    fn create_http_context(&self, _context_id: u32) -> Option<Box<dyn HttpContext>> {
+        Some(Box::new(HttpConfigHeader{
             header_content: self.header_content.clone(),
-        })
+        }))
     }
 
-    fn get_type(&self) -> ContextType {
-        ContextType::HttpContext
+    fn get_type(&self) -> Option<ContextType> {
+        Some(ContextType::HttpContext)
     }
 
 }
diff --git a/examples/http_headers.rs b/examples/http_headers.rs
@@ -27,14 +27,14 @@ struct HttpHeadersRoot;
 impl Context for HttpHeadersRoot {}
 
 impl RootContext for HttpHeadersRoot {
-    fn get_type(&self) -> ContextType {
-        ContextType::HttpContext
+    fn get_type(&self) -> Option<ContextType> {
+        Some(ContextType::HttpContext)
     }
 
-    fn create_http_context(&self, _context_id: u32) -> Box<dyn HttpContext> {
-        Box::new(HttpHeaders {
+    fn create_http_context(&self, _context_id: u32) -> Option<Box<dyn HttpContext>> {
+        Some(Box::new(HttpHeaders {
             context_id: _context_id,
-        })
+        }))
     }
 }
 

diff --git a/src/dispatcher.rs b/src/dispatcher.rs
@@ -26,6 +26,14 @@ 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));
 }
@@ -38,7 +46,9 @@ 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>>,
@@ -49,7 +59,9 @@ 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()),
@@ -60,6 +72,14 @@ 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 register_callout(&self, token_id: u32) {
         if self
             .callouts
@@ -88,7 +108,13 @@ impl Dispatcher {
 
     fn create_stream_context(&self, context_id: u32, root_context_id: u32) {
         let new_context = match self.roots.borrow().get(&root_context_id) {
-            Some(root_context) => root_context.create_stream_context(context_id),
+            Some(root_context) => match self.new_stream.get() {
+                Some(f) => f(context_id, root_context_id),
+                None => match root_context.create_stream_context(context_id) {
+                    Some(stream_context) => stream_context,
+                    None => panic!("create_stream_context returned None"),
+                },
+            },
             None => panic!("invalid root_context_id"),
         };
         if self
@@ -103,7 +129,13 @@ impl Dispatcher {
 
     fn create_http_context(&self, context_id: u32, root_context_id: u32) {
         let new_context = match self.roots.borrow().get(&root_context_id) {
-            Some(root_context) => root_context.create_http_context(context_id),
+            Some(root_context) => match self.new_http_stream.get() {
+                Some(f) => f(context_id, root_context_id),
+                None => match root_context.create_http_context(context_id) {
+                    Some(stream_context) => stream_context,
+                    None => panic!("create_http_context returned None"),
+                },
+            },
             None => panic!("invalid root_context_id"),
         };
         if self
@@ -119,16 +151,18 @@ 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);
         } 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 ContextType on root_context"),
+                Some(ContextType::HttpContext) => self.create_http_context(context_id, root_context_id),
+                Some(ContextType::StreamContext) => self.create_stream_context(context_id, root_context_id),
+                None => panic!("missing ContextType on root_context"),
             }
         } else {
-            panic!("invalid root_context_id");
+            panic!("invalid root_context_id and missing constructors");
         }
     }
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -28,5 +28,13 @@ 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() {}
diff --git a/src/traits.rs b/src/traits.rs
@@ -122,16 +122,16 @@ pub trait RootContext: Context {
 
     fn on_log(&mut self) {}
 
-    fn create_http_context(&self, _context_id: u32) -> Box<dyn HttpContext> {
-        Box::new(EmptyHttpContext)
+    fn create_http_context(&self, _context_id: u32) -> Option<Box<dyn HttpContext>> {
+        None
     }
 
-    fn create_stream_context(&self, _context_id: u32) -> Box<dyn StreamContext> {
-        Box::new(EmptyStreamContext)
+    fn create_stream_context(&self, _context_id: u32) -> Option<Box<dyn StreamContext>> {
+        None
     }
 
-    fn get_type(&self) -> ContextType {
-        ContextType::RootContext
+    fn get_type(&self) -> Option<ContextType> {
+        None
     }
 }
 
@@ -171,11 +171,6 @@ pub trait StreamContext: Context {
     fn on_log(&mut self) {}
 }
 
-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
@@ -320,8 +315,3 @@ pub trait HttpContext: Context {
 
     fn on_log(&mut self) {}
 }
-
-struct EmptyHttpContext;
-
-impl HttpContext for EmptyHttpContext {}
-impl Context for EmptyHttpContext {}
diff --git a/src/types.rs b/src/types.rs
@@ -50,9 +50,8 @@ pub enum Status {
 #[repr(u32)]
 #[derive(Debug)]
 pub enum ContextType {
-    RootContext = 0,
-    HttpContext = 1,
-    StreamContext = 2,
+    HttpContext = 0,
+    StreamContext = 1,
 }
 
 #[repr(u32)]