cc_lb_plugin_conformance/
dispatch.rs1#![cfg(feature = "dispatch")]
2
3use std::collections::BTreeSet;
4
5use cc_lb_plugin_wire::augmented_metadata::AugmentedMetadata;
6use cc_lb_plugin_wire::limits::{
7 HANDSHAKE_FUEL, HANDSHAKE_WALL_MS, SKIP_HANDSHAKE_IF_FRESH_TTL_SECS,
8};
9pub use cc_lb_plugin_wire::wire_function::FallbackPolicy;
10use cc_lb_plugin_wire::wire_function::WireFunction;
11use cc_lb_runtime_protocol::{BuildPluginError, build_plugin};
12use thiserror::Error;
13
14use crate::handshake::HandshakeError;
15
16pub fn run<F: WireFunction>(
22 wasm: &[u8],
23 request: F::Request,
24) -> Result<DispatchOutcome<F::Response>, RunError> {
25 let plugin = build_plugin(wasm, HANDSHAKE_WALL_MS, HANDSHAKE_FUEL)
26 .map_err(RunError::from_build_error)?;
27 let metadata = metadata_from_handshake(wasm, &BTreeSet::new()).map_err(RunError::Handshake)?;
28 let mut session = PluginSession { plugin, metadata };
29
30 Ok(session.dispatch::<F>(request))
31}
32
33#[non_exhaustive]
41pub struct PluginSession {
42 plugin: extism::Plugin,
43 metadata: AugmentedMetadata,
44}
45
46impl PluginSession {
47 pub fn new(wasm: &[u8]) -> Result<Self, HandshakeError> {
49 Self::new_with_caps(wasm, &BTreeSet::new())
50 }
51
52 pub fn new_with_caps(
54 wasm: &[u8],
55 host_capabilities: &BTreeSet<String>,
56 ) -> Result<Self, HandshakeError> {
57 let plugin = build_plugin(wasm, HANDSHAKE_WALL_MS, HANDSHAKE_FUEL)
58 .map_err(HandshakeError::from_build_error)?;
59 let metadata = metadata_from_handshake(wasm, host_capabilities)?;
60
61 Ok(Self { plugin, metadata })
62 }
63
64 pub fn dispatch<F: WireFunction>(
69 &mut self,
70 request: F::Request,
71 ) -> DispatchOutcome<F::Response> {
72 DispatchOutcome::from_protocol(cc_lb_runtime_protocol::dispatch::dispatch_wire_call::<F>(
73 &mut self.plugin,
74 &self.metadata,
75 request,
76 ))
77 }
78}
79
80#[non_exhaustive]
82#[derive(Clone, Debug, PartialEq, Eq)]
83pub enum DispatchOutcome<R> {
84 Ok(R),
86 Fallback(FallbackPolicy),
88}
89
90impl<R> DispatchOutcome<R> {
91 pub(crate) fn from_protocol(
92 value: cc_lb_runtime_protocol::dispatch::DispatchOutcome<R>,
93 ) -> Self {
94 match value {
95 cc_lb_runtime_protocol::dispatch::DispatchOutcome::Ok(response) => Self::Ok(response),
96 cc_lb_runtime_protocol::dispatch::DispatchOutcome::Fallback(policy) => {
97 Self::Fallback(policy)
98 }
99 _ => unreachable!(),
100 }
101 }
102}
103
104#[non_exhaustive]
106#[derive(Debug, Error)]
107pub enum RunError {
108 #[error("plugin instantiation failed: {reason}")]
110 Build { reason: String },
111 #[error("handshake failed: {0}")]
113 Handshake(HandshakeError),
114}
115
116impl RunError {
117 pub(crate) fn from_build_error(value: BuildPluginError) -> Self {
118 match value {
119 BuildPluginError::Instantiate { reason } => Self::Build { reason },
120 _ => unreachable!(),
121 }
122 }
123}
124
125impl HandshakeError {
126 pub(crate) fn from_build_error(value: BuildPluginError) -> Self {
127 match value {
128 BuildPluginError::Instantiate { reason } => Self::Instantiate { reason },
129 _ => unreachable!(),
130 }
131 }
132}
133
134fn metadata_from_handshake(
135 wasm: &[u8],
136 host_capabilities: &BTreeSet<String>,
137) -> Result<AugmentedMetadata, HandshakeError> {
138 let offer = cc_lb_runtime_protocol::handshake::build_offer(host_capabilities);
139 let accept = cc_lb_runtime_protocol::handshake::execute_handshake(wasm, &offer)
140 .map_err(HandshakeError::from_protocol)?;
141 let identity = cc_lb_runtime_protocol::identity::read_identity(wasm).map_err(|source| {
142 HandshakeError::InvalidIdentity {
143 field: "custom_section",
144 reason: source.to_string(),
145 }
146 })?;
147
148 Ok(AugmentedMetadata {
149 identity,
150 negotiated_functions: accept.chosen_versions,
151 negotiated_capabilities: accept.required_capabilities,
152 handshake_completed_at: 1,
153 self_check_passed: true,
154 self_check_completed_at: 1,
155 expires_at: 1 + SKIP_HANDSHAKE_IF_FRESH_TTL_SECS as i64,
156 })
157}