fixed flag loadig isue for stream stats page (#738)

* fixed flag loadig isue for stream stats page
* Stream display country flag fix

backend/src/api/model/active_user_manager.rs

@@ -3526,6 +3526,7 @@ mod tests {
         assert!(connection_data.sessions.iter().all(|session| session.token != token));
     }
 
+    #[allow(clippy::too_many_lines)]
     #[tokio::test]
     async fn terminate_session_promotes_soft_stream_after_releasing_capacity() {
         let config = Config::default();

frontend/src/app/components/home.rs

@@ -38,6 +38,7 @@ pub fn Home() -> Html {
     let system_info = use_state(|| None::<Rc<SystemInfo>>);
     let view_visible = use_state(|| if setup_mode { Some(ViewType::Config) } else { None });
     let theme = use_state(Theme::get_current_theme);
+    let force_update = use_force_update();
 
     let handle_theme_select = {
         let set_theme = theme.clone();
@@ -155,15 +156,20 @@ pub fn Home() -> Html {
 
     {
         let flags_service = services.flags.clone();
+        let flags_loaded = force_update.clone();
         use_effect_with(geoip_enabled, move |geoip_enabled| {
             let cancelled = Rc::new(Cell::new(false));
             if *geoip_enabled {
                 let flags_service = flags_service.clone();
+                let flags_loaded = flags_loaded.clone();
                 let cancelled = cancelled.clone();
                 spawn_local(async move {
                     while !cancelled.get() && !flags_service.is_loaded() {
                         match flags_service.ensure_loaded_from_assets().await {
-                            Ok(FlagsLoadState::Loaded) => break,
+                            Ok(FlagsLoadState::Loaded) => {
+                                flags_loaded.force_update();
+                                break;
+                            }
                             Ok(FlagsLoadState::InProgress) => {
                                 TimeoutFuture::new(250).await;
                             }
@@ -173,7 +179,12 @@ pub fn Home() -> Html {
                             }
                         }
                     }
+                    if flags_service.is_loaded() && !cancelled.get() {
+                        flags_loaded.force_update();
+                    }
                 });
+            } else {
+                flags_loaded.force_update();
             }
             move || cancelled.set(true)
         });

frontend/src/services/websocket_service.rs

@@ -10,6 +10,7 @@ use shared::{
 };
 use std::{
     cell::{Cell, RefCell},
+    collections::VecDeque,
     rc::Rc,
 };
 use wasm_bindgen::{closure::Closure, JsCast};
@@ -40,6 +41,7 @@ pub struct WebSocketService {
     connected: Rc<Cell<bool>>,
     attempt_counter: Rc<Cell<u16>>,
     ws: Rc<RefCell<Option<WebSocket>>>,
+    pending_messages: Rc<RefCell<VecDeque<Vec<u8>>>>,
     status_service: Rc<StatusService>,
     event_service: Rc<EventService>,
     ws_path: String,
@@ -58,6 +60,7 @@ impl WebSocketService {
             connected: Rc::new(Cell::new(false)),
             attempt_counter: Rc::new(Cell::new(0)),
             ws: Rc::new(RefCell::new(None)),
+            pending_messages: Rc::new(RefCell::new(VecDeque::new())),
             status_service,
             event_service,
             ws_path: concat_path_leading_slash(&base_href, "ws"),
@@ -74,6 +77,7 @@ impl WebSocketService {
             connected: self.connected.clone(),
             attempt_counter: self.attempt_counter.clone(),
             ws: self.ws.clone(),
+            pending_messages: self.pending_messages.clone(),
             status_service: self.status_service.clone(),
             event_service: self.event_service.clone(),
             ws_path: self.ws_path.clone(),
@@ -85,9 +89,19 @@ impl WebSocketService {
     }
 
     pub fn connect_ws_with_backoff(&self) {
-        if self.connected.get() {
+        let has_active_socket = self
+            .ws
+            .borrow()
+            .as_ref()
+            .map(|ws| {
+                let state = ws.ready_state();
+                state == WebSocket::CONNECTING || state == WebSocket::OPEN
+            })
+            .unwrap_or(false);
+        if has_active_socket {
             return;
         }
+        *self.ws.borrow_mut() = None;
         match WebSocket::new(&self.ws_path) {
             Err(err) => error!("Failed to open websocket connection: {err:?}"),
             Ok(socket) => {
@@ -101,12 +115,20 @@ impl WebSocketService {
                     let ws_onmessage_clone = ws_clone.clone();
                     let event_service = self.event_service.clone();
                     let attempt_counter = self.attempt_counter.clone();
+                    let connected = self.connected.clone();
+                    let pending_messages = self.pending_messages.clone();
                     let onmessage_callback =
                         Closure::<dyn FnMut(MessageEvent)>::wrap(Box::new(move |event: MessageEvent| {
                             trace!("WebSocket received message: {event:?}");
-                            if let Some(response) = handle_socket_protocol_msg(event, &event_service, &attempt_counter)
-                            {
-                                Self::try_send_message(ws_onmessage_clone.borrow().as_ref(), response);
+                            if let Some(response) = handle_socket_protocol_msg(
+                                event,
+                                &event_service,
+                                &attempt_counter,
+                                &connected,
+                                &pending_messages,
+                                &ws_onmessage_clone,
+                            ) {
+                                Self::try_send_message(ws_onmessage_clone.borrow().as_ref(), &response);
                             }
                         }));
                     socket.set_onmessage(Some(onmessage_callback.as_ref().unchecked_ref()));
@@ -120,14 +142,13 @@ impl WebSocketService {
                     let ws_open_clone = ws_clone.clone();
                     let connected_clone = self.connected.clone();
                     let onopen_callback = Closure::<dyn FnMut(Event)>::wrap(Box::new(move |_event: Event| {
-                        // on open is called on a connect attempt, it does not mean it is connected!
+                        // on open starts the protocol handshake; application messages wait until authorization.
                         trace!("WebSocket connection opened.");
-                        if Self::try_send_message(
+                        connected_clone.set(false);
+                        Self::try_send_message(
                             ws_open_clone.borrow().as_ref(),
-                            ProtocolMessage::Version(PROTOCOL_VERSION),
-                        ) {
-                            connected_clone.set(true);
-                        }
+                            &ProtocolMessage::Version(PROTOCOL_VERSION),
+                        );
                     }));
                     socket.set_onopen(Some(onopen_callback.as_ref().unchecked_ref()));
                     ws_onopen_ref.borrow_mut().replace(onopen_callback);
@@ -224,25 +245,47 @@ impl WebSocketService {
         );
     }
 
-    fn try_send_message(ws_opt: Option<&WebSocket>, msg: ProtocolMessage) -> bool {
+    fn try_send_message(ws_opt: Option<&WebSocket>, msg: &ProtocolMessage) -> bool {
+        match msg.to_bytes() {
+            Ok(bytes) => Self::try_send_bytes(ws_opt, bytes.as_ref()),
+            Err(err) => {
+                error!("Failed to create WebSocket protocol version message: {err}");
+                false
+            }
+        }
+    }
+
+    fn try_send_bytes(ws_opt: Option<&WebSocket>, bytes: &[u8]) -> bool {
         if let Some(ws) = ws_opt {
-            match msg.to_bytes() {
-                Ok(bytes) => {
-                    if let Err(err) = ws.send_with_u8_array(bytes.as_ref()) {
-                        error!("Failed to send a websocket message: {err:?}");
-                    } else {
-                        return true;
-                    }
-                }
-                Err(err) => {
-                    error!("Failed to create WebSocket protocol version message: {err}")
-                }
+            if ws.ready_state() != WebSocket::OPEN {
+                return false;
+            }
+            if let Err(err) = ws.send_with_u8_array(bytes) {
+                error!("Failed to send a websocket message: {err:?}");
+            } else {
+                return true;
             }
         }
         false
     }
 
-    pub fn send_message(&self, msg: ProtocolMessage) -> bool { Self::try_send_message(self.ws.borrow().as_ref(), msg) }
+    pub fn send_message(&self, msg: ProtocolMessage) -> bool {
+        if self.connected.get() && Self::try_send_message(self.ws.borrow().as_ref(), &msg) {
+            return true;
+        }
+
+        match msg.to_bytes() {
+            Ok(bytes) => {
+                self.pending_messages.borrow_mut().push_back(bytes.to_vec());
+                trace!("Queued websocket message until connection is ready.");
+                true
+            }
+            Err(err) => {
+                error!("Failed to create WebSocket message: {err}");
+                false
+            }
+        }
+    }
 
     pub async fn get_server_status(&self) {
         if self.connected.get() {
@@ -269,6 +312,9 @@ fn handle_socket_protocol_msg(
     event: MessageEvent,
     event_service: &Rc<EventService>,
     attempt_counter: &Rc<Cell<u16>>,
+    connected: &Rc<Cell<bool>>,
+    pending_messages: &Rc<RefCell<VecDeque<Vec<u8>>>>,
+    ws: &Rc<RefCell<Option<WebSocket>>>,
 ) -> Option<ProtocolMessage> {
     if let Ok(buf) = event.data().dyn_into::<ArrayBuffer>() {
         let array = Uint8Array::new(&buf);
@@ -277,12 +323,16 @@ fn handle_socket_protocol_msg(
             Ok(message) => {
                 match message {
                     ProtocolMessage::Unauthorized => {
+                        connected.set(false);
+                        pending_messages.borrow_mut().clear();
                         event_service.broadcast(EventMessage::Unauthorized);
                     }
                     ProtocolMessage::Error(err) => {
                         error!("{err}");
                     }
                     ProtocolMessage::Authorized => {
+                        connected.set(true);
+                        flush_pending_messages(pending_messages, ws);
                         event_service.broadcast(EventMessage::WebSocketStatus(true));
                     }
                     ProtocolMessage::ActiveUserResponse(event) => {
@@ -332,6 +382,8 @@ fn handle_socket_protocol_msg(
                         if let Some(token) = get_token() {
                             return Some(ProtocolMessage::Auth(token));
                         } else {
+                            connected.set(true);
+                            flush_pending_messages(pending_messages, ws);
                             event_service.broadcast(EventMessage::WebSocketStatus(true));
                         }
                     }
@@ -355,3 +407,26 @@ fn handle_socket_protocol_msg(
     }
     None
 }
+
+fn flush_pending_messages(pending_messages: &Rc<RefCell<VecDeque<Vec<u8>>>>, ws: &Rc<RefCell<Option<WebSocket>>>) {
+    let mut pending = pending_messages.borrow_mut();
+    if pending.is_empty() {
+        return;
+    }
+
+    let ws_borrow = ws.borrow();
+    let Some(ws) = ws_borrow.as_ref() else {
+        return;
+    };
+
+    let mut remaining = VecDeque::new();
+    while let Some(bytes) = pending.pop_front() {
+        if !WebSocketService::try_send_bytes(Some(ws), bytes.as_slice()) {
+            remaining.push_back(bytes);
+            remaining.append(&mut pending);
+            break;
+        }
+    }
+
+    *pending = remaining;
+}