frontend websocket init

This commit is contained in:
Jovi Hsu 2021-11-08 13:58:09 +08:00
parent b2d56e28c4
commit da0819c214
9 changed files with 159 additions and 25 deletions

View File

@ -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 {

View File

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

View File

@ -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;

View File

@ -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",

View File

@ -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)

View File

@ -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(
params: web::Json<AuthInfo>,
data: web::Data<AppData>,
) -> Result<HttpResponse, Error> {
pub async fn auth(params: web::Json<AuthInfo>, req: HttpRequest) -> 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!({

View File

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

View File

@ -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);
}
}
}

View File

@ -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()}
</>
}
}
}