Merge branch 'main' of https://github.com/swimos/swim-rust into fluvio

Author: SirCipher

example_apps/transit/src/ui.rs

@@ -39,8 +39,8 @@ impl UiConfig {
 pub fn ui_server_router(port: u16) -> Router {
     Router::new()
         .route("/index.html", get(index_html))
-        .route("/dist/main/swimos-transit.js", get(transit_js))
-        .route("/dist/main/swimos-transit.js.map", get(transit_js_map))
+        .route("/dist/main/swim-transit.js", get(transit_js))
+        .route("/dist/main/swim-transit.js.map", get(transit_js_map))
         .with_state(UiConfig { port })
 }
 
@@ -49,7 +49,7 @@ const INDEX: &str = "ui/index.html";
 async fn index_html(State(UiConfig { port }): State<UiConfig>) -> impl IntoResponse {
     if let Ok(file) = File::open(INDEX).await {
         let lines = LinesStream::new(BufReader::new(file).lines()).map_ok(move |mut line| {
-            if line == "const portParam = baseUri.port();" {
+            if line.contains("const portParam = baseUri.port();") {
                 format!("const portParam = {};\n", port)
             } else {
                 line.push('\n');
@@ -78,15 +78,15 @@ async fn index_html(State(UiConfig { port }): State<UiConfig>) -> impl IntoRespo
 
 async fn transit_js() -> impl IntoResponse {
     let headers = [(header::CONTENT_TYPE, JS.clone())];
-    (headers, load_file("dist/main/swimos-transit.js").await)
+    (headers, load_file("dist/main/swim-transit.js").await)
 }
 
 static HTML: HeaderValue = HeaderValue::from_static("text/html; charset=utf-8");
 static JS: HeaderValue = HeaderValue::from_static("application/json");
 
 async fn transit_js_map() -> impl IntoResponse {
     let headers = [(header::CONTENT_TYPE, JS.clone())];
-    (headers, load_file("dist/main/swimos-transit.js.map").await)
+    (headers, load_file("dist/main/swim-transit.js.map").await)
 }
 
 async fn load_file(path: &str) -> impl IntoResponse {

server/swimos_server_app/src/server/http/mod.rs

@@ -38,6 +38,7 @@ use ratchet::{
     WebSocketStream,
 };
 use ratchet_core::server::UpgradeRequestParts;
+use std::env::{var, VarError};
 use std::{
     marker::PhantomData,
     net::SocketAddr,
@@ -61,6 +62,7 @@ use tokio::{
     sync::mpsc,
     time::{sleep, Sleep},
 };
+use tracing::{debug, warn};
 
 mod resolver;
 #[cfg(test)]
@@ -375,6 +377,37 @@ where
     }
 }
 
+/// Returns whether to validate the upgrade's requested subprotocol. Users may disable this check
+/// by setting WS_NO_SUBPROTOCOL_CHECK=true if they are running a legacy version of Swim which would be
+/// incompatible with this version. If the environment variable has not been set or it is invalid
+/// then the subprotocol will still be verified and a log entry will be emitted.
+///
+/// Defaults to `true`.
+fn validate_subprotocol() -> bool {
+    match var("WS_NO_SUBPROTOCOL_CHECK") {
+        Ok(s) => match s.as_str() {
+            "true" => {
+                debug!("Skipping subprotocol check due to WS_NO_SUBPROTOCOL_CHECK=true");
+                false
+            }
+            "false" => true,
+            s => {
+                warn!("WS_NO_SUBPROTOCOL_CHECK set to an invalid value. Ignoring. Should be 'true' or 'false': {}", s);
+                true
+            }
+        },
+        Err(VarError::NotPresent) => true,
+        Err(VarError::NotUnicode(value)) => {
+            let value = value.to_string_lossy();
+            warn!(
+                "WS_NO_SUBPROTOCOL_CHECK set to non-Unicode value. Ignoring: {}",
+                value
+            );
+            true
+        }
+    }
+}
+
 /// Perform the websocket negotiation and assign the upgrade future to the target parameter.
 fn perform_upgrade<Ext, Sock, B>(
     config: WebSocketConfig,
@@ -392,7 +425,7 @@ where
 {
     match result {
         Ok(parts) => {
-            if parts.subprotocol.is_none() {
+            if parts.subprotocol.is_none() && validate_subprotocol() {
                 // We can only speak warp0 so fail the upgrade.
                 let response = swimos_http::fail_upgrade("Failed to negotiate warp0 subprotocol");
                 (response, None)

swimos_utilities/swimos_byte_channel/src/coop/mod.rs

@@ -19,7 +19,7 @@ use std::{
     task::{Context, Poll},
 };
 
-use futures::Future;
+use futures::{ready, Future};
 use pin_project::pin_project;
 
 #[cfg(test)]
@@ -71,7 +71,7 @@ fn set_budget(n: usize) {
     })
 }
 
-/// Wraps a futures and ensures that the byte channel budget is reset each time it is polled.
+/// Wraps a futures and ensures that the task budget is reset each time it is polled.
 #[pin_project]
 #[derive(Debug, Clone, Copy)]
 pub struct RunWithBudget<F> {
@@ -95,22 +95,27 @@ impl<F> RunWithBudget<F> {
     }
 }
 
-/// Extension trait to allow futures to be run with a byte channel budget.
+/// Extension trait to allow futures to be run with a task budget.
 pub trait BudgetedFutureExt: Sized + Future {
-    /// Run this future with the default byte channel budget.
+    /// Run this future with the default task budget.
     fn budgeted(self) -> RunWithBudget<Self> {
         RunWithBudget::new(self)
     }
 
-    /// Run this future with the specified byte channel budget.
+    /// Run this future with the specified task budget.
     fn with_budget(self, budget: NonZeroUsize) -> RunWithBudget<Self> {
         RunWithBudget::with_budget(budget, self)
     }
 
-    /// Run this future wit a specified byte channel budget or the default if not is specified.
+    /// Run this future wit a specified task budget or the default if not is specified.
     fn with_budget_or_default(self, budget: Option<NonZeroUsize>) -> RunWithBudget<Self> {
         RunWithBudget::with_budget(budget.unwrap_or(DEFAULT_START_BUDGET), self)
     }
+
+    /// Run this future, consuming budget if it does work.
+    fn consuming(self) -> BudgetConsumer<Self> {
+        BudgetConsumer::new(self)
+    }
 }
 
 impl<F: Future> BudgetedFutureExt for F {}
@@ -133,3 +138,26 @@ impl<F: Future> Future for RunWithBudget<F> {
         projected.fut.poll(cx)
     }
 }
+
+/// Wraps a future to consume the task budget when it performs work.
+#[pin_project]
+pub struct BudgetConsumer<F> {
+    #[pin]
+    fut: F,
+}
+
+impl<F> BudgetConsumer<F> {
+    pub fn new(fut: F) -> Self {
+        BudgetConsumer { fut }
+    }
+}
+
+impl<F: Future> Future for BudgetConsumer<F> {
+    type Output = F::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        ready!(consume_budget(cx));
+        let projected = self.project();
+        track_progress(projected.fut.poll(cx))
+    }
+}

swimos_utilities/swimos_byte_channel/src/coop/tests.rs

@@ -12,9 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use futures::task::{waker, ArcWake};
+use futures::{
+    future::{pending, ready},
+    task::{waker, ArcWake},
+};
 use std::{
     cell::Cell,
+    future::Future,
     num::NonZeroUsize,
     sync::{atomic::AtomicBool, Arc},
     task::{Context, Poll},
@@ -25,6 +29,8 @@ use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 
 use crate::{byte_channel, RunWithBudget};
 
+use super::BudgetedFutureExt;
+
 struct TestWaker(AtomicBool);
 
 impl TestWaker {
@@ -165,3 +171,48 @@ async fn with_budget_sets_budget() {
 
     fut.await;
 }
+
+#[test]
+fn consume_budget_consumes() {
+    let w = Arc::new(TestWaker::default());
+    let waker = waker(w.clone());
+    let mut cx = Context::from_waker(&waker);
+    super::set_budget(2);
+
+    let fut = pin!(ready(0).consuming());
+
+    assert_eq!(fut.poll(&mut cx), Poll::Ready(0));
+    assert_eq!(super::TASK_BUDGET.with(Cell::get), Some(1));
+
+    assert!(!w.triggered());
+}
+
+#[test]
+fn consume_budget_pending_no_consume() {
+    let w = Arc::new(TestWaker::default());
+    let waker = waker(w.clone());
+    let mut cx = Context::from_waker(&waker);
+    super::set_budget(2);
+
+    let fut = pin!(pending::<i32>().consuming());
+
+    assert_eq!(fut.poll(&mut cx), Poll::Pending);
+    assert_eq!(super::TASK_BUDGET.with(Cell::get), Some(2));
+    assert!(!w.triggered());
+}
+
+#[test]
+fn consume_budget_yields_on_exhaustion() {
+    let w = Arc::new(TestWaker::default());
+    let waker = waker(w.clone());
+    let mut cx = Context::from_waker(&waker);
+    super::set_budget(1);
+
+    let mut fut = pin!(ready(0).consuming());
+
+    assert_eq!(fut.as_mut().poll(&mut cx), Poll::Pending);
+    assert_eq!(super::TASK_BUDGET.with(Cell::get), None);
+    assert!(w.triggered());
+
+    assert_eq!(fut.poll(&mut cx), Poll::Ready(0));
+}

swimos_utilities/swimos_byte_channel/src/lib.rs

@@ -38,4 +38,4 @@ mod coop;
 pub use channel::{are_connected, byte_channel, ByteReader, ByteWriter};
 
 #[cfg(feature = "coop")]
-pub use coop::{BudgetedFutureExt, RunWithBudget};
+pub use coop::{BudgetConsumer, BudgetedFutureExt, RunWithBudget};