Tell clients their peer id on connection in Hello message

Antonio Scandurra and Nathan Sobo created

Co-Authored-By: Nathan Sobo <nathan@zed.dev>

Change summary

crates/client/src/client.rs |  82 ++++++++++++-
crates/collab/src/rpc.rs    |   2 
crates/rpc/proto/zed.proto  | 231 +++++++++++++++++++-------------------
crates/rpc/src/proto.rs     |   1 
crates/rpc/src/rpc.rs       |   2 
5 files changed, 194 insertions(+), 124 deletions(-)

Detailed changes

crates/client/src/client.rs 🔗

@@ -143,11 +143,16 @@ pub enum Status {
     Authenticating,
     Connecting,
     ConnectionError,
-    Connected { connection_id: ConnectionId },
+    Connected {
+        peer_id: PeerId,
+        connection_id: ConnectionId,
+    },
     ConnectionLost,
     Reauthenticating,
     Reconnecting,
-    ReconnectionError { next_reconnection: Instant },
+    ReconnectionError {
+        next_reconnection: Instant,
+    },
 }
 
 impl Status {
@@ -663,6 +668,7 @@ impl Client {
             self.set_status(Status::Reconnecting, cx);
         }
 
+        let mut timeout = cx.background().timer(CONNECTION_TIMEOUT).fuse();
         futures::select_biased! {
             connection = self.establish_connection(&credentials, cx).fuse() => {
                 match connection {
@@ -671,8 +677,14 @@ impl Client {
                         if !read_from_keychain && IMPERSONATE_LOGIN.is_none() {
                             write_credentials_to_keychain(&credentials, cx).log_err();
                         }
-                        self.set_connection(conn, cx);
-                        Ok(())
+
+                        futures::select_biased! {
+                            result = self.set_connection(conn, cx).fuse() => result,
+                            _ = timeout => {
+                                self.set_status(Status::ConnectionError, cx);
+                                Err(anyhow!("timed out waiting on hello message from server"))
+                            }
+                        }
                     }
                     Err(EstablishConnectionError::Unauthorized) => {
                         self.state.write().credentials.take();
@@ -695,21 +707,65 @@ impl Client {
                     }
                 }
             }
-            _ = cx.background().timer(CONNECTION_TIMEOUT).fuse() => {
+            _ = &mut timeout => {
                 self.set_status(Status::ConnectionError, cx);
                 Err(anyhow!("timed out trying to establish connection"))
             }
         }
     }
 
-    fn set_connection(self: &Arc<Self>, conn: Connection, cx: &AsyncAppContext) {
+    async fn set_connection(
+        self: &Arc<Self>,
+        conn: Connection,
+        cx: &AsyncAppContext,
+    ) -> Result<()> {
         let executor = cx.background();
         log::info!("add connection to peer");
         let (connection_id, handle_io, mut incoming) = self
             .peer
             .add_connection(conn, move |duration| executor.timer(duration));
-        log::info!("set status to connected {}", connection_id);
-        self.set_status(Status::Connected { connection_id }, cx);
+        let handle_io = cx.background().spawn(handle_io);
+
+        let peer_id = async {
+            log::info!("waiting for server hello");
+            let message = incoming
+                .next()
+                .await
+                .ok_or_else(|| anyhow!("no hello message received"))?;
+            log::info!("got server hello");
+            let hello_message_type_name = message.payload_type_name().to_string();
+            let hello = message
+                .into_any()
+                .downcast::<TypedEnvelope<proto::Hello>>()
+                .map_err(|_| {
+                    anyhow!(
+                        "invalid hello message received: {:?}",
+                        hello_message_type_name
+                    )
+                })?;
+            Ok(PeerId(hello.payload.peer_id))
+        };
+
+        let peer_id = match peer_id.await {
+            Ok(peer_id) => peer_id,
+            Err(error) => {
+                self.peer.disconnect(connection_id);
+                return Err(error);
+            }
+        };
+
+        log::info!(
+            "set status to connected (connection id: {}, peer id: {})",
+            connection_id,
+            peer_id
+        );
+        self.set_status(
+            Status::Connected {
+                peer_id,
+                connection_id,
+            },
+            cx,
+        );
         cx.foreground()
             .spawn({
                 let cx = cx.clone();
@@ -807,14 +863,18 @@ impl Client {
             })
             .detach();
 
-        let handle_io = cx.background().spawn(handle_io);
         let this = self.clone();
         let cx = cx.clone();
         cx.foreground()
             .spawn(async move {
                 match handle_io.await {
                     Ok(()) => {
-                        if *this.status().borrow() == (Status::Connected { connection_id }) {
+                        if *this.status().borrow()
+                            == (Status::Connected {
+                                connection_id,
+                                peer_id,
+                            })
+                        {
                             this.set_status(Status::SignedOut, &cx);
                         }
                     }
@@ -825,6 +885,8 @@ impl Client {
                 }
             })
             .detach();
+
+        Ok(())
     }
 
     fn authenticate(self: &Arc<Self>, cx: &AsyncAppContext) -> Task<Result<Credentials>> {

crates/collab/src/rpc.rs 🔗

@@ -369,6 +369,8 @@ impl Server {
                 });
 
             tracing::info!(%user_id, %login, %connection_id, %address, "connection opened");
+            this.peer.send(connection_id, proto::Hello { peer_id: connection_id.0 })?;
+            tracing::info!(%user_id, %login, %connection_id, %address, "sent hello message");
 
             if let Some(send_connection_id) = send_connection_id.as_mut() {
                 let _ = send_connection_id.send(connection_id).await;

crates/rpc/proto/zed.proto 🔗

@@ -6,125 +6,130 @@ message Envelope {
     optional uint32 responding_to = 2;
     optional uint32 original_sender_id = 3;
     oneof payload {
-        Ack ack = 4;
-        Error error = 5;
-        Ping ping = 6;
-        Test test = 7;
+        Hello hello = 4;
+        Ack ack = 5;
+        Error error = 6;
+        Ping ping = 7;
+        Test test = 8;
         
-        CreateRoom create_room = 8;
-        CreateRoomResponse create_room_response = 9;
-        JoinRoom join_room = 10;
-        JoinRoomResponse join_room_response = 11;
-        LeaveRoom leave_room = 12;
-        Call call = 13;
-        IncomingCall incoming_call = 14;
-        CallCanceled call_canceled = 15;
-        CancelCall cancel_call = 16;
-        DeclineCall decline_call = 17;
-        UpdateParticipantLocation update_participant_location = 18;
-        RoomUpdated room_updated = 19;
-
-        ShareProject share_project = 20;
-        ShareProjectResponse share_project_response = 21;
-        UnshareProject unshare_project = 22;
-        JoinProject join_project = 23;
-        JoinProjectResponse join_project_response = 24;
-        LeaveProject leave_project = 25;
-        AddProjectCollaborator add_project_collaborator = 26;
-        RemoveProjectCollaborator remove_project_collaborator = 27;
-
-        GetDefinition get_definition = 28;
-        GetDefinitionResponse get_definition_response = 29;
-        GetTypeDefinition get_type_definition = 30;
-        GetTypeDefinitionResponse get_type_definition_response = 31;
-        GetReferences get_references = 32;
-        GetReferencesResponse get_references_response = 33;
-        GetDocumentHighlights get_document_highlights = 34;
-        GetDocumentHighlightsResponse get_document_highlights_response = 35;
-        GetProjectSymbols get_project_symbols = 36;
-        GetProjectSymbolsResponse get_project_symbols_response = 37;
-        OpenBufferForSymbol open_buffer_for_symbol = 38;
-        OpenBufferForSymbolResponse open_buffer_for_symbol_response = 39;
-
-        UpdateProject update_project = 40;
-        RegisterProjectActivity register_project_activity = 41;
-        UpdateWorktree update_worktree = 42;
-        UpdateWorktreeExtensions update_worktree_extensions = 43;
-
-        CreateProjectEntry create_project_entry = 44;
-        RenameProjectEntry rename_project_entry = 45;
-        CopyProjectEntry copy_project_entry = 46;
-        DeleteProjectEntry delete_project_entry = 47;
-        ProjectEntryResponse project_entry_response = 48;
-
-        UpdateDiagnosticSummary update_diagnostic_summary = 49;
-        StartLanguageServer start_language_server = 50;
-        UpdateLanguageServer update_language_server = 51;
-
-        OpenBufferById open_buffer_by_id = 52;
-        OpenBufferByPath open_buffer_by_path = 53;
-        OpenBufferResponse open_buffer_response = 54;
-        CreateBufferForPeer create_buffer_for_peer = 55;
-        UpdateBuffer update_buffer = 56;
-        UpdateBufferFile update_buffer_file = 57;
-        SaveBuffer save_buffer = 58;
-        BufferSaved buffer_saved = 59;
-        BufferReloaded buffer_reloaded = 60;
-        ReloadBuffers reload_buffers = 61;
-        ReloadBuffersResponse reload_buffers_response = 62;
-        FormatBuffers format_buffers = 63;
-        FormatBuffersResponse format_buffers_response = 64;
-        GetCompletions get_completions = 65;
-        GetCompletionsResponse get_completions_response = 66;
-        ApplyCompletionAdditionalEdits apply_completion_additional_edits = 67;
-        ApplyCompletionAdditionalEditsResponse apply_completion_additional_edits_response = 68;
-        GetCodeActions get_code_actions = 69;
-        GetCodeActionsResponse get_code_actions_response = 70;
-        GetHover get_hover = 71;
-        GetHoverResponse get_hover_response = 72;
-        ApplyCodeAction apply_code_action = 73;
-        ApplyCodeActionResponse apply_code_action_response = 74;
-        PrepareRename prepare_rename = 75;
-        PrepareRenameResponse prepare_rename_response = 76;
-        PerformRename perform_rename = 77;
-        PerformRenameResponse perform_rename_response = 78;
-        SearchProject search_project = 79;
-        SearchProjectResponse search_project_response = 80;
-
-        GetChannels get_channels = 81;
-        GetChannelsResponse get_channels_response = 82;
-        JoinChannel join_channel = 83;
-        JoinChannelResponse join_channel_response = 84;
-        LeaveChannel leave_channel = 85;
-        SendChannelMessage send_channel_message = 86;
-        SendChannelMessageResponse send_channel_message_response = 87;
-        ChannelMessageSent channel_message_sent = 88;
-        GetChannelMessages get_channel_messages = 89;
-        GetChannelMessagesResponse get_channel_messages_response = 90;
-
-        UpdateContacts update_contacts = 91;
-        UpdateInviteInfo update_invite_info = 92;
-        ShowContacts show_contacts = 93;
-
-        GetUsers get_users = 94;
-        FuzzySearchUsers fuzzy_search_users = 95;
-        UsersResponse users_response = 96;
-        RequestContact request_contact = 97;
-        RespondToContactRequest respond_to_contact_request = 98;
-        RemoveContact remove_contact = 99;
-
-        Follow follow = 100;
-        FollowResponse follow_response = 101;
-        UpdateFollowers update_followers = 102;
-        Unfollow unfollow = 103;
-        GetPrivateUserInfo get_private_user_info = 104;
-        GetPrivateUserInfoResponse get_private_user_info_response = 105;
-        UpdateDiffBase update_diff_base = 106;
+        CreateRoom create_room = 9;
+        CreateRoomResponse create_room_response = 10;
+        JoinRoom join_room = 11;
+        JoinRoomResponse join_room_response = 12;
+        LeaveRoom leave_room = 13;
+        Call call = 14;
+        IncomingCall incoming_call = 15;
+        CallCanceled call_canceled = 16;
+        CancelCall cancel_call = 17;
+        DeclineCall decline_call = 18;
+        UpdateParticipantLocation update_participant_location = 19;
+        RoomUpdated room_updated = 20;
+
+        ShareProject share_project = 21;
+        ShareProjectResponse share_project_response = 22;
+        UnshareProject unshare_project = 23;
+        JoinProject join_project = 24;
+        JoinProjectResponse join_project_response = 25;
+        LeaveProject leave_project = 26;
+        AddProjectCollaborator add_project_collaborator = 27;
+        RemoveProjectCollaborator remove_project_collaborator = 28;
+
+        GetDefinition get_definition = 29;
+        GetDefinitionResponse get_definition_response = 30;
+        GetTypeDefinition get_type_definition = 31;
+        GetTypeDefinitionResponse get_type_definition_response = 32;
+        GetReferences get_references = 33;
+        GetReferencesResponse get_references_response = 34;
+        GetDocumentHighlights get_document_highlights = 35;
+        GetDocumentHighlightsResponse get_document_highlights_response = 36;
+        GetProjectSymbols get_project_symbols = 37;
+        GetProjectSymbolsResponse get_project_symbols_response = 38;
+        OpenBufferForSymbol open_buffer_for_symbol = 39;
+        OpenBufferForSymbolResponse open_buffer_for_symbol_response = 40;
+
+        UpdateProject update_project = 41;
+        RegisterProjectActivity register_project_activity = 42;
+        UpdateWorktree update_worktree = 43;
+        UpdateWorktreeExtensions update_worktree_extensions = 44;
+
+        CreateProjectEntry create_project_entry = 45;
+        RenameProjectEntry rename_project_entry = 46;
+        CopyProjectEntry copy_project_entry = 47;
+        DeleteProjectEntry delete_project_entry = 48;
+        ProjectEntryResponse project_entry_response = 49;
+
+        UpdateDiagnosticSummary update_diagnostic_summary = 50;
+        StartLanguageServer start_language_server = 51;
+        UpdateLanguageServer update_language_server = 52;
+
+        OpenBufferById open_buffer_by_id = 53;
+        OpenBufferByPath open_buffer_by_path = 54;
+        OpenBufferResponse open_buffer_response = 55;
+        CreateBufferForPeer create_buffer_for_peer = 56;
+        UpdateBuffer update_buffer = 57;
+        UpdateBufferFile update_buffer_file = 58;
+        SaveBuffer save_buffer = 59;
+        BufferSaved buffer_saved = 60;
+        BufferReloaded buffer_reloaded = 61;
+        ReloadBuffers reload_buffers = 62;
+        ReloadBuffersResponse reload_buffers_response = 63;
+        FormatBuffers format_buffers = 64;
+        FormatBuffersResponse format_buffers_response = 65;
+        GetCompletions get_completions = 66;
+        GetCompletionsResponse get_completions_response = 67;
+        ApplyCompletionAdditionalEdits apply_completion_additional_edits = 68;
+        ApplyCompletionAdditionalEditsResponse apply_completion_additional_edits_response = 69;
+        GetCodeActions get_code_actions = 70;
+        GetCodeActionsResponse get_code_actions_response = 71;
+        GetHover get_hover = 72;
+        GetHoverResponse get_hover_response = 73;
+        ApplyCodeAction apply_code_action = 74;
+        ApplyCodeActionResponse apply_code_action_response = 75;
+        PrepareRename prepare_rename = 76;
+        PrepareRenameResponse prepare_rename_response = 77;
+        PerformRename perform_rename = 78;
+        PerformRenameResponse perform_rename_response = 79;
+        SearchProject search_project = 80;
+        SearchProjectResponse search_project_response = 81;
+
+        GetChannels get_channels = 82;
+        GetChannelsResponse get_channels_response = 83;
+        JoinChannel join_channel = 84;
+        JoinChannelResponse join_channel_response = 85;
+        LeaveChannel leave_channel = 86;
+        SendChannelMessage send_channel_message = 87;
+        SendChannelMessageResponse send_channel_message_response = 88;
+        ChannelMessageSent channel_message_sent = 89;
+        GetChannelMessages get_channel_messages = 90;
+        GetChannelMessagesResponse get_channel_messages_response = 91;
+
+        UpdateContacts update_contacts = 92;
+        UpdateInviteInfo update_invite_info = 93;
+        ShowContacts show_contacts = 94;
+
+        GetUsers get_users = 95;
+        FuzzySearchUsers fuzzy_search_users = 96;
+        UsersResponse users_response = 97;
+        RequestContact request_contact = 98;
+        RespondToContactRequest respond_to_contact_request = 99;
+        RemoveContact remove_contact = 100;
+
+        Follow follow = 101;
+        FollowResponse follow_response = 102;
+        UpdateFollowers update_followers = 103;
+        Unfollow unfollow = 104;
+        GetPrivateUserInfo get_private_user_info = 105;
+        GetPrivateUserInfoResponse get_private_user_info_response = 106;
+        UpdateDiffBase update_diff_base = 107;
     }
 }
 
 // Messages
 
+message Hello {
+    uint32 peer_id = 1;
+}
+
 message Ping {}
 
 message Ack {}

crates/rpc/src/proto.rs 🔗

@@ -121,6 +121,7 @@ messages!(
     (GetProjectSymbols, Background),
     (GetProjectSymbolsResponse, Background),
     (GetUsers, Foreground),
+    (Hello, Foreground),
     (IncomingCall, Foreground),
     (UsersResponse, Foreground),
     (JoinChannel, Foreground),

crates/rpc/src/rpc.rs 🔗

@@ -6,4 +6,4 @@ pub use conn::Connection;
 pub use peer::*;
 mod macros;
 
-pub const PROTOCOL_VERSION: u32 = 37;
+pub const PROTOCOL_VERSION: u32 = 38;