Use Dropshot 'channel' for serial console endpoint

lif · lif · commit f414d0d1e2c6 · 2022-09-08T03:53:29.000Z
This moves the websocket-related boilerplate into the code generated by Dropshot's `#[channel]` annotation macro. (rel: oxidecomputer/dropshot#403 )
diff --git a/bin/propolis-server/src/lib/server.rs b/bin/propolis-server/src/lib/server.rs
@@ -2,14 +2,15 @@
 
 use anyhow::Result;
 use dropshot::{
-    endpoint, ApiDescription, HttpError, HttpResponseCreated, HttpResponseOk,
-    HttpResponseUpdatedNoContent, Path, RequestContext, TypedBody,
+    channel, endpoint, ApiDescription, HttpError, HttpResponseCreated,
+    HttpResponseOk, HttpResponseUpdatedNoContent, Path, RequestContext,
+    TypedBody, WebsocketConnection,
 };
 use futures::future::Fuse;
 use futures::stream::{SplitSink, SplitStream};
 use futures::{FutureExt, SinkExt, StreamExt};
-use hyper::upgrade::{self, Upgraded};
-use hyper::{header, Body, Response, StatusCode};
+use hyper::upgrade::Upgraded;
+use hyper::{Body, Response};
 use propolis::hw::qemu::ramfb::RamFb;
 use rfb::server::VncServer;
 use slog::{error, info, o, Logger};
@@ -21,7 +22,7 @@ use thiserror::Error;
 use tokio::sync::{mpsc, oneshot, watch, MappedMutexGuard, Mutex, MutexGuard};
 use tokio::task::JoinHandle;
 use tokio_tungstenite::tungstenite::protocol::{Role, WebSocketConfig};
-use tokio_tungstenite::tungstenite::{self, handshake, Message};
+use tokio_tungstenite::tungstenite::{self, Message};
 use tokio_tungstenite::WebSocketStream;
 
 use propolis::dispatch::AsyncCtx;
@@ -757,74 +758,20 @@ async fn instance_serial_task(
     Ok(())
 }
 
-#[endpoint {
-    method = GET,
+#[channel {
+    protocol = WEBSOCKETS,
     path = "/instance/serial",
 }]
 async fn instance_serial(
     rqctx: Arc<RequestContext<DropshotEndpointContext>>,
-) -> Result<Response<Body>, HttpError> {
+    websock: WebsocketConnection,
+) -> dropshot::WebsocketChannelResult {
     let mut inst = rqctx.context().instance().await?;
 
     let serial = inst.serial.clone();
-    let request = &mut *rqctx.request.lock().await;
-
-    if !request
-        .headers()
-        .get(header::CONNECTION)
-        .and_then(|hv| hv.to_str().ok())
-        .map(|hv| {
-            hv.split(|c| c == ',' || c == ' ')
-                .any(|vs| vs.eq_ignore_ascii_case("upgrade"))
-        })
-        .unwrap_or(false)
-    {
-        return Err(HttpError::for_bad_request(
-            None,
-            "expected connection upgrade".to_string(),
-        ));
-    }
-    if !request
-        .headers()
-        .get(header::UPGRADE)
-        .and_then(|v| v.to_str().ok())
-        .map(|v| {
-            v.split(|c| c == ',' || c == ' ')
-                .any(|v| v.eq_ignore_ascii_case("websocket"))
-        })
-        .unwrap_or(false)
-    {
-        return Err(HttpError::for_bad_request(
-            None,
-            "unexpected protocol for upgrade".to_string(),
-        ));
-    }
-    if request
-        .headers()
-        .get(header::SEC_WEBSOCKET_VERSION)
-        .map(|v| v.as_bytes())
-        != Some(b"13")
-    {
-        return Err(HttpError::for_bad_request(
-            None,
-            "missing or invalid websocket version".to_string(),
-        ));
-    }
-    let accept_key = request
-        .headers()
-        .get(header::SEC_WEBSOCKET_KEY)
-        .map(|hv| hv.as_bytes())
-        .map(|key| handshake::derive_accept_key(key))
-        .ok_or_else(|| {
-            HttpError::for_bad_request(
-                None,
-                "missing websocket key".to_string(),
-            )
-        })?;
 
     let actx = inst.instance.async_ctx();
-    let ws_log = rqctx.log.new(o!());
-    let err_log = ws_log.clone();
+    let err_log = rqctx.log.new(o!());
 
     // Create or get active serial task handle and channels
     let serial_task = inst.serial_task.get_or_insert_with(move || {
@@ -836,49 +783,33 @@ async fn instance_serial(
                 websocks_recv,
                 close_recv,
                 serial,
-                ws_log.clone(),
+                err_log.clone(),
                 &actx,
             )
             .await
             {
-                error!(ws_log, "Failed to spawn instance serial task: {}", e);
+                error!(err_log, "Failed to spawn instance serial task: {}", e);
             }
         });
 
         SerialTask { task, close_ch: Some(close_ch), websocks_ch }
     });
 
-    let upgrade_fut = upgrade::on(request);
     let config =
         WebSocketConfig { max_send_queue: Some(4096), ..Default::default() };
     let websocks_send = serial_task.websocks_ch.clone();
-    tokio::spawn(async move {
-        let upgraded = match upgrade_fut.await {
-            Ok(u) => u,
-            Err(e) => {
-                error!(err_log, "Serial socket upgrade failed: {}", e);
-                return;
-            }
-        };
-
-        let ws_stream = WebSocketStream::from_raw_socket(
-            upgraded,
-            Role::Server,
-            Some(config),
-        )
-        .await;
 
-        if let Err(e) = websocks_send.send(ws_stream).await {
-            error!(err_log, "Serial socket hand-off failed: {}", e);
-        }
-    });
+    let ws_stream = WebSocketStream::from_raw_socket(
+        websock.into_inner(),
+        Role::Server,
+        Some(config),
+    )
+    .await;
 
-    Ok(Response::builder()
-        .status(StatusCode::SWITCHING_PROTOCOLS)
-        .header(header::CONNECTION, "Upgrade")
-        .header(header::UPGRADE, "websocket")
-        .header(header::SEC_WEBSOCKET_ACCEPT, accept_key)
-        .body(Body::empty())?)
+    websocks_send
+        .send(ws_stream.into())
+        .await
+        .map_err(|e| format!("Serial socket hand-off failed: {}", e).into())
 }
 
 // This endpoint is meant to only be called during a migration from the destination
