From 5ffaf89bc90ae4bd2154f8b8615afe83d3338b50 Mon Sep 17 00:00:00 2001 From: Adam Singer Date: Mon, 17 Jun 2024 14:37:19 -0700 Subject: [PATCH] [bug] Ensure OperationId is used at external protocol points (#1001) Execution server was leaking out `ActionInfoHashKey`, remove leakage and ensure `OperationId` is used for external communications. Removed conversion from string to `ActionInfoHashKey`. Added tests around `OperationId` string parsing. Fixes: https://github.com/TraceMachina/nativelink/issues/1000 --- nativelink-service/src/execution_server.rs | 15 ++- nativelink-util/BUILD.bazel | 1 + nativelink-util/src/action_messages.rs | 102 +++++++++------ nativelink-util/tests/operation_id_tests.rs | 130 ++++++++++++++++++++ 4 files changed, 205 insertions(+), 43 deletions(-) create mode 100644 nativelink-util/tests/operation_id_tests.rs diff --git a/nativelink-service/src/execution_server.rs b/nativelink-service/src/execution_server.rs index ebc9cbb19..a42f0ef03 100644 --- a/nativelink-service/src/execution_server.rs +++ b/nativelink-service/src/execution_server.rs @@ -31,7 +31,7 @@ use nativelink_scheduler::action_scheduler::ActionScheduler; use nativelink_store::ac_utils::get_and_decode_digest; use nativelink_store::store_manager::StoreManager; use nativelink_util::action_messages::{ - ActionInfo, ActionInfoHashKey, ActionState, DEFAULT_EXECUTION_PRIORITY, + ActionInfo, ActionInfoHashKey, ActionState, OperationId, DEFAULT_EXECUTION_PRIORITY, }; use nativelink_util::common::DigestInfo; use nativelink_util::digest_hasher::{make_ctx_for_hash_func, DigestHasherFunc}; @@ -238,17 +238,20 @@ impl ExecutionServer { &self, request: Request, ) -> Result, Status> { - let unique_qualifier = ActionInfoHashKey::try_from(request.into_inner().name.as_str()) - .err_tip(|| "Decoding operation name into ActionInfoHashKey")?; - let Some(instance_info) = self.instance_infos.get(&unique_qualifier.instance_name) else { + let operation_id = OperationId::try_from(request.into_inner().name.as_str()) + .err_tip(|| "Decoding operation name into OperationId")?; + let Some(instance_info) = self + .instance_infos + .get(&operation_id.unique_qualifier.instance_name) + else { return Err(Status::not_found(format!( "No scheduler with the instance name {}", - unique_qualifier.instance_name + operation_id.unique_qualifier.instance_name ))); }; let Some(rx) = instance_info .scheduler - .find_existing_action(&unique_qualifier) + .find_existing_action(&operation_id.unique_qualifier) .await else { return Err(Status::not_found("Failed to find existing task")); diff --git a/nativelink-util/BUILD.bazel b/nativelink-util/BUILD.bazel index 5f8d4c50a..2f982acd0 100644 --- a/nativelink-util/BUILD.bazel +++ b/nativelink-util/BUILD.bazel @@ -74,6 +74,7 @@ rust_test_suite( "tests/fastcdc_test.rs", "tests/fs_test.rs", "tests/health_utils_test.rs", + "tests/operation_id_tests.rs", "tests/proto_stream_utils_test.rs", "tests/resource_info_test.rs", "tests/retry_test.rs", diff --git a/nativelink-util/src/action_messages.rs b/nativelink-util/src/action_messages.rs index 5f2da86a3..b1b185b23 100644 --- a/nativelink-util/src/action_messages.rs +++ b/nativelink-util/src/action_messages.rs @@ -81,13 +81,73 @@ impl OperationId { impl TryFrom<&str> for OperationId { type Error = Error; + /// Attempts to convert a string slice into an `OperationId`. + /// + /// The input string `value` is expected to be in the format: + /// `//-//`. + /// + /// # Parameters + /// + /// - `value`: A `&str` representing the `OperationId` to be converted. + /// + /// # Returns + /// + /// - `Result`: Returns `Ok(OperationId)` if the conversion is + /// successful, or an `Err(Error)` if it fails. + /// + /// # Errors + /// + /// This function will return an error if the input string is not in the expected format or if + /// any of the components cannot be parsed correctly. The detailed error messages provide + /// insight into which part of the input string caused the failure. + /// + /// ## Example Usage + /// + /// ```no_run + /// use nativelink_util::action_messages::OperationId; + /// let operation_id_str = "main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/0/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52"; + /// let operation_id = OperationId::try_from(operation_id_str); + /// ``` + /// + /// In this example, `operation_id_str` is a string representing an `OperationId`. + /// The `try_from` method is used to convert it into an `OperationId` instance. + /// If any part of the string is incorrectly formatted, an error will be returned with a + /// descriptive message. fn try_from(value: &str) -> Result { let (unique_qualifier, id) = value - .split_once(':') - .err_tip(|| "Invalid Id string - {value}")?; + .rsplit_once('/') + .err_tip(|| format!("Invalid OperationId unique_qualifier / id fragment - {value}"))?; + let (instance_name, rest) = unique_qualifier + .split_once('/') + .err_tip(|| format!("Invalid ActionInfoHashKey instance name fragment - {value}"))?; + let (digest_function, rest) = rest + .split_once('/') + .err_tip(|| format!("Invalid ActionInfoHashKey digest function fragment - {value}"))?; + let (digest_hash, rest) = rest + .split_once('-') + .err_tip(|| format!("Invalid ActionInfoHashKey digest hash fragment - {value}"))?; + let (digest_size, salt) = rest + .split_once('/') + .err_tip(|| format!("Invalid ActionInfoHashKey digest size fragment - {value}"))?; + let digest = DigestInfo::try_new( + digest_hash, + digest_size + .parse::() + .err_tip(|| format!("Invalid ActionInfoHashKey size value fragment - {value}"))?, + ) + .err_tip(|| format!("Invalid DigestInfo digest hash - {value}"))?; + let salt = u64::from_str_radix(salt, 16) + .err_tip(|| format!("Invalid ActionInfoHashKey salt hex conversion - {value}"))?; + let unique_qualifier = ActionInfoHashKey { + instance_name: instance_name.to_string(), + digest_function: digest_function.try_into()?, + digest, + salt, + }; + let id = Uuid::parse_str(id).map_err(|e| make_input_err!("Failed to parse {e} as uuid"))?; Ok(Self { - unique_qualifier: ActionInfoHashKey::try_from(unique_qualifier)?, - id: Uuid::parse_str(id).map_err(|e| make_input_err!("Failed to parse {e} as uuid"))?, + unique_qualifier, + id, }) } } @@ -95,7 +155,7 @@ impl TryFrom<&str> for OperationId { impl std::fmt::Display for OperationId { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_fmt(format_args!( - "{}:{}", + "{}/{}", self.unique_qualifier.action_name(), self.id )) @@ -191,38 +251,6 @@ impl ActionInfoHashKey { } } -impl TryFrom<&str> for ActionInfoHashKey { - type Error = Error; - - fn try_from(value: &str) -> Result { - let (instance_name, other) = value - .split_once('/') - .err_tip(|| "Invalid ActionInfoHashKey string - {value}")?; - let (digest_function, other) = other - .split_once('/') - .err_tip(|| "Invalid ActionInfoHashKey string - {value}")?; - let (digest_hash, other) = other - .split_once('-') - .err_tip(|| "Invalid ActionInfoHashKey string - {value}")?; - let (digest_size, salt) = other - .split_once('/') - .err_tip(|| "Invalid ActionInfoHashKey string - {value}")?; - let digest = DigestInfo::try_new( - digest_hash, - digest_size - .parse::() - .err_tip(|| "Expected digest size to be a number for ActionInfoHashKey")?, - )?; - let salt = u64::from_str_radix(salt, 16).err_tip(|| "Expected salt to be a hex string")?; - Ok(Self { - instance_name: instance_name.to_string(), - digest_function: digest_function.try_into()?, - digest, - salt, - }) - } -} - /// Information needed to execute an action. This struct is used over bazel's proto `Action` /// for simplicity and offers a `salt`, which is useful to ensure during hashing (for dicts) /// to ensure we never match against another `ActionInfo` (when a task should never be cached). diff --git a/nativelink-util/tests/operation_id_tests.rs b/nativelink-util/tests/operation_id_tests.rs new file mode 100644 index 000000000..e1e8b5e30 --- /dev/null +++ b/nativelink-util/tests/operation_id_tests.rs @@ -0,0 +1,130 @@ +// Copyright 2024 The Native Link Authors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use nativelink_error::{Code, Error}; +use nativelink_macro::nativelink_test; +use nativelink_util::action_messages::OperationId; +use pretty_assertions::assert_eq; + +#[nativelink_test] +async fn parse_operation_id() -> Result<(), Error> { + let operation_id = OperationId::try_from("main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/0/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52").unwrap(); + assert_eq!( + operation_id.to_string(), + "main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/0/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52"); + assert_eq!( + operation_id.action_name(), + "main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/0" + ); + assert_eq!( + operation_id.id.to_string(), + "19b16cf8-a1ad-4948-aaac-b6f4eb7fca52" + ); + assert_eq!( + hex::encode(operation_id.get_hash()), + "5a36f0db39e27667c4b91937cd29c1df8799ba468f2de6810c6865be05517644" + ); + Ok(()) +} + +#[nativelink_test] +async fn parse_empty_failure() -> Result<(), Error> { + let operation_id = OperationId::try_from("").err().unwrap(); + assert_eq!(operation_id.code, Code::Internal); + assert_eq!(operation_id.messages.len(), 1); + assert_eq!( + operation_id.messages[0], + "Invalid OperationId unique_qualifier / id fragment - " + ); + + let operation_id = OperationId::try_from("/").err().unwrap(); + assert_eq!(operation_id.code, Code::Internal); + assert_eq!(operation_id.messages.len(), 1); + assert_eq!( + operation_id.messages[0], + "Invalid ActionInfoHashKey instance name fragment - /" + ); + + let operation_id = OperationId::try_from("main").err().unwrap(); + assert_eq!(operation_id.code, Code::Internal); + assert_eq!(operation_id.messages.len(), 1); + assert_eq!( + operation_id.messages[0], + "Invalid OperationId unique_qualifier / id fragment - main" + ); + + let operation_id = OperationId::try_from("main/nohashfn/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/0/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52").err().unwrap(); + assert_eq!(operation_id.code, Code::InvalidArgument); + assert_eq!(operation_id.messages.len(), 1); + assert_eq!( + operation_id.messages[0], + "Unknown or unsupported digest function for string conversion: \"NOHASHFN\"" + ); + + let operation_id = + OperationId::try_from("main/SHA256/badhash-211/0/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52") + .err() + .unwrap(); + assert_eq!(operation_id.messages.len(), 3); + assert_eq!(operation_id.code, Code::InvalidArgument); + assert_eq!(operation_id.messages[0], "Odd number of digits"); + assert_eq!(operation_id.messages[1], "Invalid sha256 hash: badhash"); + assert_eq!( + operation_id.messages[2], + "Invalid DigestInfo digest hash - main/SHA256/badhash-211/0/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52" + ); + + let operation_id = OperationId::try_from("main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-/0/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52").err().unwrap(); + assert_eq!(operation_id.messages.len(), 2); + assert_eq!(operation_id.code, Code::InvalidArgument); + assert_eq!( + operation_id.messages[0], + "cannot parse integer from empty string" + ); + assert_eq!(operation_id.messages[1], "Invalid ActionInfoHashKey size value fragment - main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-/0/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52"); + + let operation_id = OperationId::try_from("main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f--211/0/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52").err().unwrap(); + assert_eq!(operation_id.code, Code::InvalidArgument); + assert_eq!(operation_id.messages.len(), 2); + assert_eq!(operation_id.messages[0], "invalid digit found in string"); + assert_eq!(operation_id.messages[1], "Invalid ActionInfoHashKey size value fragment - main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f--211/0/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52"); + + let operation_id = OperationId::try_from("main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/x/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52").err().unwrap(); + assert_eq!(operation_id.messages.len(), 2); + assert_eq!(operation_id.code, Code::InvalidArgument); + assert_eq!(operation_id.messages[0], "invalid digit found in string"); + assert_eq!(operation_id.messages[1], "Invalid ActionInfoHashKey salt hex conversion - main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/x/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52"); + + let operation_id = OperationId::try_from("main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/-10/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52").err().unwrap(); + assert_eq!(operation_id.messages.len(), 2); + assert_eq!(operation_id.code, Code::InvalidArgument); + assert_eq!(operation_id.messages[0], "invalid digit found in string"); + assert_eq!(operation_id.messages[1], "Invalid ActionInfoHashKey salt hex conversion - main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/-10/19b16cf8-a1ad-4948-aaac-b6f4eb7fca52"); + + let operation_id = OperationId::try_from("main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/0/baduuid").err().unwrap(); + assert_eq!(operation_id.messages.len(), 1); + assert_eq!(operation_id.code, Code::InvalidArgument); + assert_eq!(operation_id.messages[0], "Failed to parse invalid character: expected an optional prefix of `urn:uuid:` followed by [0-9a-fA-F-], found `u` at 4 as uuid"); + + let operation_id = OperationId::try_from( + "main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/0", + ) + .err() + .unwrap(); + assert_eq!(operation_id.messages.len(), 1); + assert_eq!(operation_id.code, Code::Internal); + assert_eq!(operation_id.messages[0], "Invalid ActionInfoHashKey digest size fragment - main/SHA256/4a0885a39d5ba8da3123c02ff56b73196a8b23fd3c835e1446e74a3a3ff4313f-211/0"); + + Ok(()) +}