Add SSZ response support for validator_balances endpoints

beacon_node/http_api/src/beacon/states.rs

@@ -779,17 +779,32 @@ pub fn post_beacon_state_validator_balances<T: BeaconChainTypes>(
         .and(warp::path("validator_balances"))
         .and(warp::path::end())
         .and(warp_utils::json::json_no_body())
+        .and(warp::header::optional::<api_types::Accept>("accept"))
         .then(
             |state_id: StateId,
              task_spawner: TaskSpawner<T::EthSpec>,
              chain: Arc<BeaconChain<T>>,
-             query: ValidatorBalancesRequestBody| {
-                task_spawner.blocking_json_task(Priority::P1, move || {
-                    crate::validators::get_beacon_state_validator_balances(
+             query: ValidatorBalancesRequestBody,
+             accept_header: Option<api_types::Accept>| {
+                task_spawner.blocking_response_task(Priority::P1, move || {
+                    let result = crate::validators::get_beacon_state_validator_balances(
                         state_id,
                         chain,
                         Some(&query.ids),
-                    )
+                    )?;
+                    match accept_header {
+                        Some(api_types::Accept::Ssz) => Builder::new()
+                            .status(200)
+                            .body(result.data.as_ssz_bytes())
+                            .map(add_ssz_content_type_header)
+                            .map_err(|e| {
+                                warp_utils::reject::custom_server_error(format!(
+                                    "failed to create response: {}",
+                                    e
+                                ))
+                            }),
+                        _ => Ok(warp::reply::json(&result).into_response()),
+                    }
                 })
             },
         )
@@ -805,18 +820,34 @@ pub fn get_beacon_state_validator_balances<T: BeaconChainTypes>(
         .and(warp::path("validator_balances"))
         .and(warp::path::end())
         .and(multi_key_query::<eth2::types::ValidatorBalancesQuery>())
+        .and(warp::header::optional::<api_types::Accept>("accept"))
         .then(
             |state_id: StateId,
              task_spawner: TaskSpawner<T::EthSpec>,
              chain: Arc<BeaconChain<T>>,
-             query_res: Result<eth2::types::ValidatorBalancesQuery, warp::Rejection>| {
-                task_spawner.blocking_json_task(Priority::P1, move || {
+             query_res: Result<eth2::types::ValidatorBalancesQuery, warp::Rejection>,
+             accept_header: Option<api_types::Accept>| {
+                task_spawner.blocking_response_task(Priority::P1, move || {
                     let query = query_res?;
-                    crate::validators::get_beacon_state_validator_balances(
+                    let result = crate::validators::get_beacon_state_validator_balances(
                         state_id,
                         chain,
                         query.id.as_deref(),
-                    )
+                    )?;
+
+                    match accept_header {
+                        Some(api_types::Accept::Ssz) => Builder::new()
+                            .status(200)
+                            .body(result.data.as_ssz_bytes())
+                            .map(add_ssz_content_type_header)
+                            .map_err(|e| {
+                                warp_utils::reject::custom_server_error(format!(
+                                    "failed to create response: {}",
+                                    e
+                                ))
+                            }),
+                        _ => Ok(warp::reply::json(&result).into_response()),
+                    }
                 })
             },
         )

beacon_node/http_api/tests/tests.rs

@@ -984,6 +984,149 @@ impl ApiTester {
         self
     }
 
+    pub async fn test_beacon_states_validator_balances_ssz(self) -> Self {
+        for state_id in self.interesting_state_ids() {
+            let all_state_opt = state_id.state(&self.chain).ok();
+            let all_ssz_result = self
+                .client
+                .get_beacon_states_validator_balances_ssz(state_id.0, None)
+                .await
+                .unwrap();
+
+            if all_ssz_result.is_some() || all_state_opt.is_some() {
+                let ssz_bytes = all_ssz_result.expect("response should exist");
+                let result = Vec::<ValidatorBalanceData>::from_ssz_bytes(&ssz_bytes)
+                    .expect("should decode SSZ validator balances");
+
+                let (state, _, _) = all_state_opt.as_ref().expect("state should exist");
+                let expected: Vec<ValidatorBalanceData> = state
+                    .balances()
+                    .iter()
+                    .enumerate()
+                    .map(|(index, balance)| ValidatorBalanceData {
+                        index: index as u64,
+                        balance: *balance,
+                    })
+                    .collect();
+
+                assert_eq!(result, expected, "{:?}", state_id);
+            }
+
+            for validator_indices in self.interesting_validator_indices() {
+                let state_opt = state_id.state(&self.chain).ok();
+
+                let validators: Vec<Validator> = match state_opt.as_ref() {
+                    Some((state, _, _)) => state.validators().to_vec(),
+                    None => vec![],
+                };
+
+                let validator_index_ids: Vec<ValidatorId> = validator_indices
+                    .iter()
+                    .cloned()
+                    .map(ValidatorId::Index)
+                    .collect();
+
+                let validator_pubkey_ids: Vec<ValidatorId> = validator_indices
+                    .iter()
+                    .cloned()
+                    .map(|i| {
+                        ValidatorId::PublicKey(
+                            validators
+                                .get(i as usize)
+                                .map_or(PublicKeyBytes::empty(), |val| val.pubkey),
+                        )
+                    })
+                    .collect();
+
+                let ssz_result = match self
+                    .client
+                    .get_beacon_states_validator_balances_ssz(
+                        state_id.0,
+                        Some(validator_index_ids.as_slice()),
+                    )
+                    .await
+                {
+                    Ok(response) => response,
+                    Err(e) => panic!("query failed incorrectly: {e:?}"),
+                };
+
+                if ssz_result.is_none() && state_opt.is_none() {
+                    continue;
+                }
+
+                let ssz_bytes = ssz_result.expect("response should exist");
+                let get_result_index_ids = Vec::<ValidatorBalanceData>::from_ssz_bytes(&ssz_bytes)
+                    .expect("should decode SSZ validator balances");
+
+                let get_ssz_bytes_pubkey = self
+                    .client
+                    .get_beacon_states_validator_balances_ssz(
+                        state_id.0,
+                        Some(validator_pubkey_ids.as_slice()),
+                    )
+                    .await
+                    .unwrap()
+                    .expect("response should exist");
+                let get_result_pubkey_ids =
+                    Vec::<ValidatorBalanceData>::from_ssz_bytes(&get_ssz_bytes_pubkey)
+                        .expect("should decode SSZ validator balances");
+
+                let post_ssz_bytes_index = self
+                    .client
+                    .post_beacon_states_validator_balances_ssz(state_id.0, validator_index_ids)
+                    .await
+                    .unwrap()
+                    .expect("response should exist");
+                let post_result_index_ids =
+                    Vec::<ValidatorBalanceData>::from_ssz_bytes(&post_ssz_bytes_index)
+                        .expect("should decode SSZ validator balances");
+
+                let post_ssz_bytes_pubkey = self
+                    .client
+                    .post_beacon_states_validator_balances_ssz(state_id.0, validator_pubkey_ids)
+                    .await
+                    .unwrap()
+                    .expect("response should exist");
+                let post_result_pubkey_ids =
+                    Vec::<ValidatorBalanceData>::from_ssz_bytes(&post_ssz_bytes_pubkey)
+                        .expect("should decode SSZ validator balances");
+
+                let expected: Vec<ValidatorBalanceData> = {
+                    let (state, _, _) = state_opt.as_ref().expect("state should exist");
+                    // If validator_indices is empty, all balances are returned.
+                    if validator_indices.is_empty() {
+                        state
+                            .balances()
+                            .iter()
+                            .enumerate()
+                            .map(|(index, balance)| ValidatorBalanceData {
+                                index: index as u64,
+                                balance: *balance,
+                            })
+                            .collect()
+                    } else {
+                        let mut validators = Vec::with_capacity(validator_indices.len());
+                        for i in validator_indices {
+                            if i < state.balances().len() as u64 {
+                                validators.push(ValidatorBalanceData {
+                                    index: i,
+                                    balance: *state.balances().get(i as usize).unwrap(),
+                                });
+                            }
+                        }
+                        validators
+                    }
+                };
+
+                assert_eq!(get_result_index_ids, expected, "{:?}", state_id);
+                assert_eq!(get_result_pubkey_ids, expected, "{:?}", state_id);
+                assert_eq!(post_result_index_ids, expected, "{:?}", state_id);
+                assert_eq!(post_result_pubkey_ids, expected, "{:?}", state_id);
+            }
+        }
+        self
+    }
+
     pub async fn test_beacon_states_validator_identities(self) -> Self {
         for state_id in self.interesting_state_ids() {
             for validator_indices in self.interesting_validator_indices() {
@@ -9723,6 +9866,8 @@ async fn beacon_get_state_info() {
         .test_beacon_states_randao()
         .await
         .test_beacon_states_validator_identities_ssz()
+        .await
+        .test_beacon_states_validator_balances_ssz()
         .await;
 }
 

common/eth2/src/lib.rs

@@ -710,6 +710,36 @@ impl BeaconNodeHttpClient {
         self.get_opt(path).await
     }
 
+    /// `GET beacon/states/{state_id}/validator_balances?id`
+    ///
+    /// Returns `Ok(None)` on a 404 error.
+    pub async fn get_beacon_states_validator_balances_ssz(
+        &self,
+        state_id: StateId,
+        ids: Option<&[ValidatorId]>,
+    ) -> Result<Option<Vec<u8>>, Error> {
+        let mut path = self.eth_path(V1)?;
+
+        path.path_segments_mut()
+            .map_err(|()| Error::InvalidUrl(self.server.clone()))?
+            .push("beacon")
+            .push("states")
+            .push(&state_id.to_string())
+            .push("validator_balances");
+
+        if let Some(ids) = ids {
+            let id_string = ids
+                .iter()
+                .map(|i| i.to_string())
+                .collect::<Vec<_>>()
+                .join(",");
+            path.query_pairs_mut().append_pair("id", &id_string);
+        }
+
+        self.get_bytes_opt_accept_header(path, Accept::Ssz, self.timeouts.default)
+            .await
+    }
+
     /// TESTING ONLY: This request should fail with a 415 response code.
     pub async fn post_beacon_states_validator_balances_with_ssz_header(
         &self,
@@ -753,6 +783,29 @@ impl BeaconNodeHttpClient {
         self.post_with_opt_response(path, &request).await
     }
 
+    /// `POST beacon/states/{state_id}/validator_balances`
+    ///
+    /// Returns `Ok(None)` on a 404 error.
+    pub async fn post_beacon_states_validator_balances_ssz(
+        &self,
+        state_id: StateId,
+        ids: Vec<ValidatorId>,
+    ) -> Result<Option<Vec<u8>>, Error> {
+        let mut path = self.eth_path(V1)?;
+
+        path.path_segments_mut()
+            .map_err(|()| Error::InvalidUrl(self.server.clone()))?
+            .push("beacon")
+            .push("states")
+            .push(&state_id.to_string())
+            .push("validator_balances");
+
+        let request = ValidatorBalancesRequestBody { ids };
+
+        self.post_bytes_opt_accept_header(path, &request, Accept::Ssz, self.timeouts.default)
+            .await
+    }
+
     /// `POST beacon/states/{state_id}/validator_identities`
     ///
     ///  Returns `Ok(None)` on a 404 error.

common/eth2/src/types.rs

@@ -359,7 +359,7 @@ pub struct ValidatorData {
     pub validator: Validator,
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Encode, Decode)]
 pub struct ValidatorBalanceData {
     #[serde(with = "serde_utils::quoted_u64")]
     pub index: u64,