From b80dedc354956e4f20568eca89731d31cc5dde52 Mon Sep 17 00:00:00 2001
From: eatradish <sakiiily@aosc.io>
Date: Wed, 19 Jun 2024 05:48:08 +0800
Subject: [PATCH] feat(worker): filter machine port to recv message

---
 worker/src/build.rs | 7 ++++---
 worker/src/lib.rs   | 4 ++++
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/worker/src/build.rs b/worker/src/build.rs
index 49ce7bd..c4ae52f 100644
--- a/worker/src/build.rs
+++ b/worker/src/build.rs
@@ -5,6 +5,7 @@ use common::{JobOk, WorkerJobUpdateRequest, WorkerPollRequest, WorkerPollRespons
 use flume::{unbounded, Receiver, Sender};
 use futures_util::StreamExt;
 use log::{error, info, warn};
+use reqwest::Url;
 use std::{
     path::Path,
     process::{Output, Stdio},
@@ -371,7 +372,7 @@ async fn build_worker_inner(args: &Args) -> anyhow::Result<()> {
         logical_cores: num_cpus::get() as i32,
     };
 
-    let ws = args.websocket.clone();
+    let ws = Url::parse(&args.websocket)?.join(&args.port.to_string())?;
 
     let (tx, rx) = unbounded();
 
@@ -430,10 +431,10 @@ pub async fn build_worker(args: Args) -> ! {
     }
 }
 
-pub async fn websocket_connect(rx: Receiver<Message>, ws: String) -> ! {
+pub async fn websocket_connect(rx: Receiver<Message>, ws: Url) -> ! {
     loop {
         info!("Starting websocket connect");
-        match connect_async(&ws).await {
+        match connect_async(ws.as_str()).await {
             Ok((ws_stream, _)) => {
                 let (write, _) = ws_stream.split();
                 let rx = rx.clone().into_stream();
diff --git a/worker/src/lib.rs b/worker/src/lib.rs
index 00dc5f4..26d9402 100644
--- a/worker/src/lib.rs
+++ b/worker/src/lib.rs
@@ -53,6 +53,10 @@ pub struct Args {
     /// Websocket uri
     #[arg(short = 'w', long, env = "BUILDIT_WS")]
     pub websocket: String,
+
+    /// Worker machine relay port
+    #[arg(short = 'w', long, env = "BUILDIT_RELAY_PORT")]
+    pub port: u16,
 }
 
 pub fn get_memory_bytes() -> i64 {