Fix grpc sub entity ids (#1273)

dojoengine · Dec 13, 2023 · 6bbeecc · 6bbeecc
1 parent 772e4b6
commit 6bbeecc
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/crates/torii/grpc/src/server/mod.rs b/crates/torii/grpc/src/server/mod.rs
@@ -9,7 +9,6 @@ use std::sync::Arc;
 
 use dojo_types::schema::Ty;
 use futures::Stream;
-use hex::encode;
 use proto::world::{
     MetadataRequest, MetadataResponse, RetrieveEntitiesRequest, RetrieveEntitiesResponse,
     SubscribeModelsRequest, SubscribeModelsResponse,
@@ -282,7 +281,7 @@ impl DojoWorld {
 
     async fn subscribe_entities(
         &self,
-        ids: Vec<String>,
+        ids: Vec<FieldElement>,
     ) -> Result<Receiver<Result<proto::world::SubscribeEntityResponse, tonic::Status>>, Error> {
         self.entity_manager.add_subscriber(ids).await
     }
@@ -379,7 +378,13 @@ impl proto::world::world_server::World for DojoWorld {
         request: Request<SubscribeEntitiesRequest>,
     ) -> ServiceResult<Self::SubscribeEntitiesStream> {
         let SubscribeEntitiesRequest { ids } = request.into_inner();
-        let ids = ids.iter().map(encode).collect();
+        let ids = ids
+            .iter()
+            .map(|id| {
+                FieldElement::from_byte_slice_be(id)
+                    .map_err(|e| Status::invalid_argument(e.to_string()))
+            })
+            .collect::<Result<Vec<_>, _>>()?;
         let rx = self.subscribe_entities(ids).await.map_err(|e| Status::internal(e.to_string()))?;
 
         Ok(Response::new(Box::pin(ReceiverStream::new(rx)) as Self::SubscribeEntitiesStream))

diff --git a/crates/torii/grpc/src/server/subscriptions/entity.rs b/crates/torii/grpc/src/server/subscriptions/entity.rs
@@ -23,7 +23,7 @@ use crate::proto;
 
 pub struct EntitiesSubscriber {
     /// Entity ids that the subscriber is interested in
-    ids: HashSet<String>,
+    ids: HashSet<FieldElement>,
     /// The channel to send the response back to the subscriber.
     sender: Sender<Result<proto::world::SubscribeEntityResponse, tonic::Status>>,
 }
@@ -36,7 +36,7 @@ pub struct EntityManager {
 impl EntityManager {
     pub async fn add_subscriber(
         &self,
-        ids: Vec<String>,
+        ids: Vec<FieldElement>,
     ) -> Result<Receiver<Result<proto::world::SubscribeEntityResponse, tonic::Status>>, Error> {
         let id = rand::thread_rng().gen::<usize>();
         let (sender, receiver) = channel(1);
@@ -85,8 +85,9 @@ impl Service {
         let mut closed_stream = Vec::new();
 
         for (idx, sub) in subs.subscribers.read().await.iter() {
+            let felt_id = FieldElement::from_str(id).map_err(ParseError::FromStr)?;
             // publish all updates if ids is empty or only ids that are subscribed to
-            if sub.ids.is_empty() || sub.ids.contains(id) {
+            if sub.ids.is_empty() || sub.ids.contains(&felt_id) {
                 let models_query = r#"
                     SELECT group_concat(entity_model.model_id) as model_names
                     FROM entities
@@ -114,10 +115,7 @@ impl Service {
 
                 let resp = proto::world::SubscribeEntityResponse {
                     entity: Some(proto::types::Entity {
-                        id: FieldElement::from_str(id)
-                            .map_err(ParseError::FromStr)?
-                            .to_bytes_be()
-                            .to_vec(),
+                        id: felt_id.to_bytes_be().to_vec(),
                         models,
                     }),
                 };