frontend websocket init

backend/src/agent/agent.rs

@@ -27,7 +27,12 @@ impl Decoder for TcpCodec {
 
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
         info!("recv from server: {:?}", src);
-        Ok(Some(Bytes::from(src.to_vec())))
+        if 0 == src.len() {
+            return Ok(None);
+        }
+        let web_bytes = Bytes::from(src.to_vec());
+        src.clear();
+        Ok(Some(web_bytes))
     }
 }
 
@@ -181,10 +186,12 @@ impl Handler<AgentManagerMsg> for AgentManager {
     fn handle(&mut self, msg: AgentManagerMsg, _ctx: &mut Context<Self>) -> Self::Result {
         match msg {
             AgentManagerMsg::Add(addr) => {
+                info!("add agent: {:?}", addr.0);
                 self.agents.insert(addr.0, addr.1);
                 AgentManagerResult::NoReturn
             }
             AgentManagerMsg::Get(aid) => {
+                info!("get agent: {}", aid);
                 if let Some(addr) = self.agents.get(&aid) {
                     AgentManagerResult::Success(addr.clone())
                 } else {

backend/src/agent/mod.rs

@@ -1,4 +1,4 @@
-pub mod remote;
-pub mod ws;
-pub mod resolver;
 pub mod agent;
+pub mod remote;
+pub mod resolver;
+pub mod ws;

backend/src/agent/remote.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use actix_session::Session;
 use actix_web::*;
 use serde::{Deserialize, Serialize};
@@ -6,8 +8,6 @@ use serde_json::json;
 use log::info;
 use rand::Rng;
 
-use crate::AppData;
-
 use super::agent;
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -22,14 +22,14 @@ pub struct RemoteInfo {
 
 #[post("/target/validate")]
 pub async fn target_validate(
-    data: web::Data<AppData>,
+    req: HttpRequest,
     params: web::Json<RemoteInfo>,
 ) -> Result<HttpResponse, Error> {
     let remote = params.into_inner();
     info!("{:?}", remote);
-    // let resolved = data.resolver.send(ResolveMsg::Resolve(remote.host)).await;
+    let app_data = req.app_data::<Arc<crate::AppData>>().unwrap();
 
-    match data.resolver.lockup(remote.host).await {
+    match app_data.resolver.lockup(remote.host).await {
         Some(ipaddr) => {
             let json = json!({
                 "status": "success",
@@ -49,18 +49,19 @@ pub async fn target_validate(
 
 #[post("/target/ssh")]
 pub async fn target_ssh(
+    req: HttpRequest,
     session: Session,
-    data: web::Data<AppData>,
     params: web::Json<RemoteInfo>,
 ) -> Result<HttpResponse, Error> {
     let aid = rand::thread_rng().gen::<u32>();
+    let app_data = req.app_data::<Arc<crate::AppData>>().unwrap();
     let remote = params.into_inner();
     let agent = agent::Agent::new(aid, (remote.ip, remote.port)).await;
 
     match agent {
         Some(addr) => {
             // add to agent list
-            let _ = data
+            let _ = app_data
                 .agents
                 .send(agent::AgentManagerMsg::Add((aid, addr)))
                 .await;

backend/src/agent/ws.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use actix::{Actor, Addr, Message, StreamHandler};
 use actix::{AsyncContext, Handler};
 use actix_session::Session;
@@ -7,8 +9,6 @@ use actix_web::{web, Error, HttpRequest, HttpResponse};
 use actix_web_actors::ws;
 use log::*;
 
-use crate::AppData;
-
 use super::agent::*;
 
 #[derive(Message)]
@@ -63,12 +63,17 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WsSession {
 pub async fn ws_index(
     req: HttpRequest,
     session: Session,
-    data: web::Data<AppData>,
     stream: web::Payload,
 ) -> Result<HttpResponse, Error> {
     let aid = session.get::<u32>("aid").unwrap_or(Some(0)).unwrap();
+    let app_data = req.app_data::<Arc<crate::AppData>>().unwrap();
 
-    let resp = match data.agents.send(AgentManagerMsg::Get(aid)).await.unwrap() {
+    let resp = match app_data
+        .agents
+        .send(AgentManagerMsg::Get(aid))
+        .await
+        .unwrap()
+    {
         AgentManagerResult::Success(agent) => ws::start(WsSession { agent }, &req, stream),
         _ => Err(Error::from(actix_web::error::ErrorInternalServerError(
             "Agent not found",

backend/src/main.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use actix::Addr;
 use actix_files as fs;
 use actix_session::CookieSession;
@@ -61,9 +63,10 @@ async fn main() -> std::io::Result<()> {
 
     info!("Server starts at http://127.0.0.1:8080");
     let private_key = rand::thread_rng().gen::<[u8; 32]>();
+    let app_data = Arc::new(AppData::new());
     HttpServer::new(move || {
         App::new()
-            .app_data(AppData::new())
+            .app_data(app_data.clone())
             .wrap(CookieSession::signed(&private_key).secure(false))
             .wrap(middleware::Compress::new(ContentEncoding::Gzip))
             .service(index)

backend/src/user/auth.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use actix::{Actor, Addr, Context, Handler, Message, MessageResponse};
 use actix_web::*;
 use serde::{Deserialize, Serialize};
@@ -5,8 +7,6 @@ use serde_json::json;
 
 use log::info;
 
-use crate::AppData;
-
 #[derive(MessageResponse)]
 #[allow(dead_code)]
 enum AuthResult {
@@ -63,12 +63,13 @@ impl Handler<AuthMsg> for Authenticator {
 }
 
 #[post("/auth")]
-pub async fn auth(
+pub async fn auth(params: web::Json<AuthInfo>, req: HttpRequest) -> Result<HttpResponse, Error> {
-    params: web::Json<AuthInfo>,
-    data: web::Data<AppData>,
-) -> Result<HttpResponse, Error> {
     let auth_info = params.into_inner();
-    let res = data.authenticator.send(AuthMsg::DoAuth(auth_info)).await;
+    let app_data = req.app_data::<Arc<crate::AppData>>().unwrap();
+    let res = app_data
+        .authenticator
+        .send(AuthMsg::DoAuth(auth_info))
+        .await;
 
     match res {
         Ok(AuthResult::AuthSuccess) => Ok(HttpResponse::Ok().json(json!({

frontend/src/components/mod.rs

@@ -1,2 +1,3 @@
 pub mod auth;
 pub mod host;
+pub mod ws;

frontend/src/components/ws.rs

@@ -0,0 +1,106 @@
+use yew::prelude::*;
+use yew::services::websocket::{WebSocketService, WebSocketStatus, WebSocketTask};
+use yew::services::ConsoleService;
+use yew::{format::Binary, utils::host};
+
+pub struct WebsocketCtx {
+    ws: Option<WebSocketTask>,
+    link: ComponentLink<Self>,
+    error_msg: String,
+    onrecv: Callback<Vec<u8>>,
+}
+
+#[derive(Clone, PartialEq, Properties)]
+pub struct WebsocketProps {
+    #[prop_or_default]
+    pub onrecv: Callback<Vec<u8>>,
+}
+
+pub enum WebsocketMsg {
+    Connect,
+    Disconnected,
+    Ignore,
+    Send(Binary),
+    Recv(Binary),
+}
+
+impl Component for WebsocketCtx {
+    type Message = WebsocketMsg;
+    type Properties = WebsocketProps;
+
+    fn create(props: Self::Properties, link: ComponentLink<Self>) -> Self {
+        Self {
+            ws: None,
+            link: link,
+            error_msg: String::new(),
+            onrecv: props.onrecv,
+        }
+    }
+
+    fn update(&mut self, msg: Self::Message) -> ShouldRender {
+        match msg {
+            WebsocketMsg::Connect => {
+                ConsoleService::log("Connecting");
+                let cbout = self.link.callback(|data| WebsocketMsg::Recv(data));
+                let cbnot = self.link.callback(|input| {
+                    ConsoleService::log(&format!("Notification: {:?}", input));
+                    match input {
+                        WebSocketStatus::Closed | WebSocketStatus::Error => {
+                            WebsocketMsg::Disconnected
+                        }
+                        _ => WebsocketMsg::Ignore,
+                    }
+                });
+                if self.ws.is_none() {
+                    let task = WebSocketService::connect_binary(
+                        &format!("ws://{}/ws", host().unwrap()),
+                        cbout,
+                        cbnot,
+                    );
+                    self.ws = Some(task.unwrap());
+                }
+                true
+            }
+            WebsocketMsg::Disconnected => {
+                self.ws = None;
+                self.error_msg = "Disconnected".to_string();
+                true
+            }
+            WebsocketMsg::Ignore => false,
+            WebsocketMsg::Send(data) => {
+                if let Some(ref mut ws) = self.ws {
+                    ws.send_binary(data);
+                }
+                false
+            }
+            WebsocketMsg::Recv(Ok(s)) => {
+                // ConsoleService::log(&format!("recv {:?}", s));
+                self.onrecv.emit(s);
+                false
+            }
+            WebsocketMsg::Recv(Err(s)) => {
+                self.error_msg = format!("Error when reading from server: {}\n", &s.to_string());
+                self.link.send_message(WebsocketMsg::Disconnected);
+                true
+            }
+        }
+    }
+
+    fn change(&mut self, _prop: Self::Properties) -> ShouldRender {
+        false
+    }
+
+    fn view(&self) -> Html {
+        html! {
+            <>
+            {self.error_msg.clone()}
+            </>
+        }
+    }
+
+    fn rendered(&mut self, first_render: bool) {
+        if first_render && self.ws.is_none() {
+            self.link.send_message(WebsocketMsg::Connect);
+        }
+    }
+}

frontend/src/pages/page_ssh.rs

@@ -22,6 +22,7 @@ pub enum SshMsg {
     SshConnect((String, u16)),
     SshConnectResp(Result<Value, anyhow::Error>),
     SshConnected,
+    SshRecv(Vec<u8>),
 }
 
 impl Component for PageSsh {
@@ -88,6 +89,10 @@ impl Component for PageSsh {
                 self.connected = true;
                 true
             }
+            SshMsg::SshRecv(v) => {
+                self.error_msg = String::from_utf8(v).unwrap();
+                true
+            }
         }
     }
 
@@ -96,8 +101,8 @@ impl Component for PageSsh {
     }
 
     fn view(&self) -> Html {
-        let connect_ssh = self.link.callback(SshMsg::SshConnect);
         if !self.connected {
+            let connect_ssh = self.link.callback(SshMsg::SshConnect);
             html! {
                 <>
                     <components::host::Host onsubmit=connect_ssh/>
@@ -105,8 +110,13 @@ impl Component for PageSsh {
                 </>
             }
         } else {
+            let recv_msg = self.link.callback(|v| SshMsg::SshRecv(v));
             html! {
-                <></>
+                <>
+                    <components::ws::WebsocketCtx
+                        onrecv=recv_msg/>
+                    {self.error_msg.clone()}
+                </>
             }
         }
     }