feat: TUI supports inspecting websocket (#37)

sigoden · Oct 4, 2024 · 4eed92b · 4eed92b
1 parent 74799fb
commit 4eed92b
Show file tree

Hide file tree

Showing 6 changed files with 231 additions and 118 deletions.
diff --git a/assets/index.html b/assets/index.html
@@ -313,6 +313,7 @@
         }
 
         @media (max-width: 1000px) {
+
             .col-status,
             .col-type,
             .col-size,
@@ -421,6 +422,7 @@
          * @property {?number} size
          * @property {?number} time
          * @property {string} mime
+         * @property {?number} websocket_id
          */
 
         /**
@@ -987,7 +989,7 @@
             let date = formatDate(new Date(data.create));
             let value = data.body.value;
             if (data.body.encode == "base64") {
-                value = `data:;base64,${value}`;
+                value = base64ToText(data.body.value);
             }
             return `<div class="ws-msg">
                 <small>${arrow}<span class="ws-date">${date}</span></small>
@@ -1066,24 +1068,38 @@
             return value.startsWith('x-') ? value.substring(2) : value;
         }
 
+        /**
+         * @param {string} base64String
+         * @param {string} contentType
+         */
         function base64ToBlob(base64String, contentType) {
-            const byteCharacters = atob(base64String);
-            const byteArrays = [];
+            const byteArrays= base64ToUint8Array(base64String);
+            return new Blob([byteArrays], { type: contentType });
+        }
+
+        /**
+         * @param {string} base64String
+         */
+        function base64ToText(base64String) {
+            const byteArrays= base64ToUint8Array(base64String);
+            const decoder = new TextDecoder('utf-8');
+            return decoder.decode(byteArrays);
+        }
 
-            for (let offset = 0; offset < byteCharacters.length; offset += 512) {
-                const slice = byteCharacters.slice(offset, offset + 512);
+        /**
+         * @param {string} base64String
+         */
+        function base64ToUint8Array(base64String) {
+            const binaryString = atob(base64String);
 
-                const byteNumbers = new Array(slice.length);
-                for (let i = 0; i < slice.length; i++) {
-                    byteNumbers[i] = slice.charCodeAt(i);
-                }
+            const length = binaryString.length;
+            const bytes = new Uint8Array(length);
 
-                const byteArray = new Uint8Array(byteNumbers);
-                byteArrays.push(byteArray);
+            for (let i = 0; i < length; i++) {
+                bytes[i] = binaryString.charCodeAt(i);
             }
 
-            return new Blob(byteArrays, { type: contentType });
-
+            return bytes;
         }
 
         function formatSize(len) {

diff --git a/src/server.rs b/src/server.rs
@@ -442,6 +442,7 @@ impl Server {
 
         let Some((messages, receiver)) = self.state.subscribe_websocket(id).await else {
             *res.status_mut() = StatusCode::NOT_FOUND;
+            set_res_body(res, "Not found websocket");
             return Ok(());
         };
 

diff --git a/src/state.rs b/src/state.rs
@@ -22,7 +22,7 @@ pub struct State {
 
 impl State {
     pub fn new(print_mode: PrintMode) -> Self {
-        let (traffics_notifier, _) = broadcast::channel(16);
+        let (traffics_notifier, _) = broadcast::channel(128);
         let (websockets_notifier, _) = broadcast::channel(64);
         Self {
             print_mode,
@@ -183,14 +183,14 @@ impl State {
         let _ = self.websockets_notifier.send((id, message));
     }
 
-    pub async fn subscribe_websocket(&self, id: usize) -> Option<SubscribeWebSocket> {
+    pub async fn subscribe_websocket(&self, id: usize) -> Option<SubscribedWebSocket> {
         let websockets = self.websockets.lock().await;
         let messages = websockets.get(&id)?;
         Some((messages.to_vec(), self.websockets_notifier.subscribe()))
     }
 }
 
-pub type SubscribeWebSocket = (
+pub type SubscribedWebSocket = (
     Vec<WebsocketMessage>,
     broadcast::Receiver<(usize, WebsocketMessage)>,
 );
@@ -206,7 +206,7 @@ pub enum WebsocketMessage {
 #[derive(Debug, Clone, Serialize)]
 pub struct WebsocketData {
     #[serde(serialize_with = "crate::utils::serialize_datetime")]
-    create: OffsetDateTime,
-    server_to_client: bool,
-    body: Body,
+    pub create: OffsetDateTime,
+    pub server_to_client: bool,
+    pub body: Body,
 }
diff --git a/src/traffic.rs b/src/traffic.rs
@@ -1,7 +1,6 @@
 use crate::utils::*;
 
 use anyhow::{bail, Result};
-use base64::{engine::general_purpose::STANDARD, Engine as _};
 use http::{HeaderMap, StatusCode, Version};
 use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
@@ -232,6 +231,7 @@ impl Traffic {
             size: self.res_body_size,
             time: self.time(),
             mime: extract_mime(&self.res_headers).to_string(),
+            websocket_id: self.websocket_id,
         }
     }
 
@@ -331,6 +331,8 @@ pub struct TrafficHead {
     pub size: Option<u64>,
     pub time: Option<u64>,
     pub mime: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub websocket_id: Option<usize>,
 }
 
 impl TrafficHead {
@@ -443,7 +445,7 @@ impl Body {
             Ok(text) => Self::text(text),
             Err(_) => Body {
                 encode: "base64".to_string(),
-                value: STANDARD.encode(data),
+                value: base64_encode(data),
                 size: size as _,
             },
         }