From ae38e6d12c5e5586cf935983bf27501ef1f7daa5 Mon Sep 17 00:00:00 2001 From: Dmitry Porokh Date: Fri, 24 Jul 2026 11:47:00 -0700 Subject: [PATCH] feat(admin-cli): generate RPC commands from protobuf Add custom protobuf options for CLI command and argument metadata, then compile annotated unary RPCs into clap argument types and dispatch implementations. Migrate TPM CA deletion and primary DPU selection as initial examples. Signed-off-by: Dmitry Porokh --- Cargo.lock | 2 + Cargo.toml | 1 + crates/admin-cli/src/managed_host/mod.rs | 16 +- .../src/managed_host/set_primary_dpu/args.rs | 52 --- .../src/managed_host/set_primary_dpu/cmd.rs | 25 -- .../src/managed_host/set_primary_dpu/mod.rs | 32 -- crates/admin-cli/src/managed_host/tests.rs | 15 + crates/admin-cli/src/tpm_ca/delete/args.rs | 31 -- crates/admin-cli/src/tpm_ca/delete/cmd.rs | 23 -- crates/admin-cli/src/tpm_ca/delete/mod.rs | 31 -- crates/admin-cli/src/tpm_ca/mod.rs | 11 +- crates/admin-cli/src/tpm_ca/tests.rs | 5 +- crates/macros/src/lib.rs | 25 +- crates/rpc/Cargo.toml | 7 +- crates/rpc/build.rs | 5 +- crates/rpc/proto/admin_cli.proto | 36 ++ crates/rpc/proto/forge.proto | 36 +- crates/rpc/src/admin_cli.rs | 8 + crates/rpc/src/protos/mod.rs | 3 + crates/tonic-client-wrapper/Cargo.toml | 4 + crates/tonic-client-wrapper/src/codegen.rs | 377 +++++++++++++++++- 21 files changed, 534 insertions(+), 211 deletions(-) delete mode 100644 crates/admin-cli/src/managed_host/set_primary_dpu/args.rs delete mode 100644 crates/admin-cli/src/managed_host/set_primary_dpu/cmd.rs delete mode 100644 crates/admin-cli/src/managed_host/set_primary_dpu/mod.rs delete mode 100644 crates/admin-cli/src/tpm_ca/delete/args.rs delete mode 100644 crates/admin-cli/src/tpm_ca/delete/cmd.rs delete mode 100644 crates/admin-cli/src/tpm_ca/delete/mod.rs create mode 100644 crates/rpc/proto/admin_cli.proto diff --git a/Cargo.lock b/Cargo.lock index db075175b3..330ce7a5a7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11815,7 +11815,9 @@ dependencies = [ "lazy_static", "prettyplease", "proc-macro2", + "prost", "prost-types", + "protobuf", "quote", "syn 2.0.117", "temp-dir", diff --git a/Cargo.toml b/Cargo.toml index 371be4f86f..6592434721 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -185,6 +185,7 @@ prometheus = "0.14.0" prost = "0.14.3" prost-build = "0.14.3" prost-types = "0.14.3" +protobuf = "3.7.2" pwhash = "1.0.0" quick-xml = "0.39" quote = { default-features = false, version = "1.0" } diff --git a/crates/admin-cli/src/managed_host/mod.rs b/crates/admin-cli/src/managed_host/mod.rs index 464f6f777c..16473013e1 100644 --- a/crates/admin-cli/src/managed_host/mod.rs +++ b/crates/admin-cli/src/managed_host/mod.rs @@ -20,7 +20,6 @@ mod maintenance; mod power_options; mod quarantine; mod reset_host_reprovisioning; -mod set_primary_dpu; mod set_primary_interface; mod show; mod start_updates; @@ -61,7 +60,20 @@ pub enum Cmd { #[clap(about = "Set the primary interface (boot device) for the managed host")] SetPrimaryInterface(set_primary_interface::Args), #[clap(about = "Deprecated: use set-primary-interface. Sets the primary DPU.")] - SetPrimaryDpu(set_primary_dpu::Args), + #[command(after_long_help = "\ +EXAMPLES: + +Set the primary DPU for a host: + $ nico-admin-cli managed-host set-primary-dpu 12345678-1234-5678-90ab-cdef01234567 \ + abcdef01-2345-6789-abcd-ef0123456789 + +Set the primary DPU and reboot the host afterward: + $ nico-admin-cli managed-host set-primary-dpu 12345678-1234-5678-90ab-cdef01234567 \ + abcdef01-2345-6789-abcd-ef0123456789 --reboot + +")] + #[rpc] + SetPrimaryDpu(rpc::forge::set_primary_dpu::Args), #[clap(about = "Download debug bundle with logs for a specific host")] DebugBundle(debug_bundle::Args), } diff --git a/crates/admin-cli/src/managed_host/set_primary_dpu/args.rs b/crates/admin-cli/src/managed_host/set_primary_dpu/args.rs deleted file mode 100644 index 6a457e57c5..0000000000 --- a/crates/admin-cli/src/managed_host/set_primary_dpu/args.rs +++ /dev/null @@ -1,52 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 carbide_uuid::machine::MachineId; -use clap::Parser; -use rpc::forge as forgerpc; - -#[derive(Parser, Debug)] -#[command(after_long_help = "\ -EXAMPLES: - -Set the primary DPU for a host: - $ nico-admin-cli managed-host set-primary-dpu 12345678-1234-5678-90ab-cdef01234567 \ - abcdef01-2345-6789-abcd-ef0123456789 - -Set the primary DPU and reboot the host afterward: - $ nico-admin-cli managed-host set-primary-dpu 12345678-1234-5678-90ab-cdef01234567 \ - abcdef01-2345-6789-abcd-ef0123456789 --reboot - -")] -pub struct Args { - #[clap(help = "ID of the host machine")] - pub host_machine_id: MachineId, - #[clap(help = "ID of the DPU machine to make primary")] - pub dpu_machine_id: MachineId, - #[clap(long, help = "Reboot the host after the update")] - pub reboot: bool, -} - -impl From for forgerpc::SetPrimaryDpuRequest { - fn from(args: Args) -> Self { - Self { - host_machine_id: Some(args.host_machine_id), - dpu_machine_id: Some(args.dpu_machine_id), - reboot: args.reboot, - } - } -} diff --git a/crates/admin-cli/src/managed_host/set_primary_dpu/cmd.rs b/crates/admin-cli/src/managed_host/set_primary_dpu/cmd.rs deleted file mode 100644 index 311f38aa85..0000000000 --- a/crates/admin-cli/src/managed_host/set_primary_dpu/cmd.rs +++ /dev/null @@ -1,25 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 super::args::Args; -use crate::errors::CarbideCliResult; -use crate::rpc::ApiClient; - -pub async fn set_primary_dpu(api_client: &ApiClient, args: Args) -> CarbideCliResult<()> { - api_client.0.set_primary_dpu(args).await?; - Ok(()) -} diff --git a/crates/admin-cli/src/managed_host/set_primary_dpu/mod.rs b/crates/admin-cli/src/managed_host/set_primary_dpu/mod.rs deleted file mode 100644 index 2c7c83308c..0000000000 --- a/crates/admin-cli/src/managed_host/set_primary_dpu/mod.rs +++ /dev/null @@ -1,32 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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. - */ - -pub mod args; -pub mod cmd; - -pub use args::Args; - -use crate::cfg::run::Run; -use crate::cfg::runtime::RuntimeContext; -use crate::errors::CarbideCliResult; - -impl Run for Args { - async fn run(self, ctx: &mut RuntimeContext) -> CarbideCliResult<()> { - cmd::set_primary_dpu(&ctx.api_client, self).await?; - Ok(()) - } -} diff --git a/crates/admin-cli/src/managed_host/tests.rs b/crates/admin-cli/src/managed_host/tests.rs index 8da71da257..397c28b531 100644 --- a/crates/admin-cli/src/managed_host/tests.rs +++ b/crates/admin-cli/src/managed_host/tests.rs @@ -249,6 +249,21 @@ fn parse_set_primary_dpu() { match cmd { Cmd::SetPrimaryDpu(args) => { assert!(!args.reboot); + let request: rpc::forge::SetPrimaryDpuRequest = args.into(); + assert_eq!( + request + .host_machine_id + .expect("host ID must be populated") + .to_string(), + TEST_MACHINE_ID + ); + assert_eq!( + request + .dpu_machine_id + .expect("DPU ID must be populated") + .to_string(), + TEST_MACHINE_ID + ); } _ => panic!("expected SetPrimaryDpu variant"), } diff --git a/crates/admin-cli/src/tpm_ca/delete/args.rs b/crates/admin-cli/src/tpm_ca/delete/args.rs deleted file mode 100644 index ab0b42137b..0000000000 --- a/crates/admin-cli/src/tpm_ca/delete/args.rs +++ /dev/null @@ -1,31 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 clap::Parser; - -#[derive(Parser, Debug)] -#[command(after_long_help = "\ -EXAMPLES: - -Delete a TPM CA certificate by its id (from `tpm-ca show`): - $ nico-admin-cli tpm-ca delete --ca-id 42 - -")] -pub struct Args { - #[clap(short, long, help = "TPM CA id obtained from the show command")] - pub ca_id: i32, -} diff --git a/crates/admin-cli/src/tpm_ca/delete/cmd.rs b/crates/admin-cli/src/tpm_ca/delete/cmd.rs deleted file mode 100644 index 8fd750db8c..0000000000 --- a/crates/admin-cli/src/tpm_ca/delete/cmd.rs +++ /dev/null @@ -1,23 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 crate::errors::CarbideCliResult; -use crate::rpc::ApiClient; - -pub async fn delete(ca_cert_id: i32, api_client: &ApiClient) -> CarbideCliResult<()> { - Ok(api_client.0.tpm_delete_ca_cert(ca_cert_id).await?) -} diff --git a/crates/admin-cli/src/tpm_ca/delete/mod.rs b/crates/admin-cli/src/tpm_ca/delete/mod.rs deleted file mode 100644 index 1d43022abd..0000000000 --- a/crates/admin-cli/src/tpm_ca/delete/mod.rs +++ /dev/null @@ -1,31 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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. - */ - -pub mod args; -pub mod cmd; - -pub use args::Args; - -use crate::cfg::run::Run; -use crate::cfg::runtime::RuntimeContext; -use crate::errors::CarbideCliResult; - -impl Run for Args { - async fn run(self, ctx: &mut RuntimeContext) -> CarbideCliResult<()> { - cmd::delete(self.ca_id, &ctx.api_client).await - } -} diff --git a/crates/admin-cli/src/tpm_ca/mod.rs b/crates/admin-cli/src/tpm_ca/mod.rs index b1b21ba94b..ed5385e50c 100644 --- a/crates/admin-cli/src/tpm_ca/mod.rs +++ b/crates/admin-cli/src/tpm_ca/mod.rs @@ -17,7 +17,6 @@ mod add; mod add_bulk; -mod delete; mod show; mod show_unmatched_ek; @@ -33,7 +32,15 @@ pub enum Cmd { #[clap(about = "Show all TPM CA certificates")] Show(show::Args), #[clap(about = "Delete TPM CA certificate with a given id")] - Delete(delete::Args), + #[command(after_long_help = "\ +EXAMPLES: + +Delete a TPM CA certificate by its id (from `tpm-ca show`): + $ nico-admin-cli tpm-ca delete --ca-id 42 + +")] + #[rpc] + Delete(rpc::forge::tpm_delete_ca_cert::Args), #[clap(about = "Add TPM CA certificate encoded in DER/CER/PEM format in a given file")] Add(add::Args), #[clap(about = "Show TPM EK certificates for which there is no CA match")] diff --git a/crates/admin-cli/src/tpm_ca/tests.rs b/crates/admin-cli/src/tpm_ca/tests.rs index f7b03ae3e1..12c05093c6 100644 --- a/crates/admin-cli/src/tpm_ca/tests.rs +++ b/crates/admin-cli/src/tpm_ca/tests.rs @@ -81,7 +81,10 @@ fn parse_delete() { run = |argv| { Cmd::try_parse_from(argv.iter().copied()) .map(|cmd| match cmd { - Cmd::Delete(args) => args.ca_id, + Cmd::Delete(args) => { + let request: rpc::forge::TpmCaCertId = args.into(); + request.ca_cert_id + } _ => panic!("expected Delete variant"), }) .map_err(drop) diff --git a/crates/macros/src/lib.rs b/crates/macros/src/lib.rs index ecb002f64a..bf8bdad96e 100644 --- a/crates/macros/src/lib.rs +++ b/crates/macros/src/lib.rs @@ -50,7 +50,7 @@ type AttributeArgs = syn::punctuated::Punctuated; /// SubGroup(sub::Cmd), /// } /// ``` -#[proc_macro_derive(Dispatch, attributes(dispatch))] +#[proc_macro_derive(Dispatch, attributes(dispatch, rpc))] pub fn derive_dispatch(input: TokenStream) -> TokenStream { let input = syn::parse_macro_input!(input as DeriveInput); match expand_dispatch(input) { @@ -74,15 +74,29 @@ fn expand_dispatch(input: DeriveInput) -> syn::Result { let mut run_arms = Vec::new(); let mut dispatch_arms = Vec::new(); + let mut rpc_arms = Vec::new(); for variant in &data.variants { let variant_name = &variant.ident; let is_dispatch = variant.attrs.iter().any(|a| a.path().is_ident("dispatch")); + let is_rpc = variant.attrs.iter().any(|a| a.path().is_ident("rpc")); - if is_dispatch { + if is_dispatch && is_rpc { + return Err(syn::Error::new_spanned( + variant, + "a Dispatch variant cannot be both #[dispatch] and #[rpc]", + )); + } else if is_dispatch { dispatch_arms.push(quote! { #name::#variant_name(cmd) => cmd.dispatch(ctx).await, }); + } else if is_rpc { + rpc_arms.push(quote! { + #name::#variant_name(args) => args + .execute(&ctx.api_client.0) + .await + .map_err(Into::into), + }); } else { run_arms.push(quote! { #name::#variant_name(args) => args.run(&mut ctx).await, @@ -95,6 +109,11 @@ fn expand_dispatch(input: DeriveInput) -> syn::Result { } else { quote! { use crate::cfg::dispatch::Dispatch as _; } }; + let rpc_import = if rpc_arms.is_empty() { + quote! {} + } else { + quote! { use ::rpc::admin_cli::CliRpcCommand as _; } + }; let output = quote! { impl crate::cfg::dispatch::Dispatch for #name { @@ -104,9 +123,11 @@ fn expand_dispatch(input: DeriveInput) -> syn::Result { ) -> crate::errors::CarbideCliResult<()> { use crate::cfg::run::Run; #dispatch_import + #rpc_import match self { #(#run_arms)* #(#dispatch_arms)* + #(#rpc_arms)* } } } diff --git a/crates/rpc/Cargo.toml b/crates/rpc/Cargo.toml index 004cd6a593..0c6f423f35 100644 --- a/crates/rpc/Cargo.toml +++ b/crates/rpc/Cargo.toml @@ -26,7 +26,12 @@ name = "rpc" [features] sqlx = ["dep:sqlx"] -cli = ["dep:clap", "dep:prettytable-rs", "dep:serde_yaml", "dep:csv"] +cli = [ + "dep:clap", + "dep:prettytable-rs", + "dep:serde_yaml", + "dep:csv", +] model = [ "dep:carbide-api-model", "dep:once_cell", diff --git a/crates/rpc/build.rs b/crates/rpc/build.rs index 97cc54a021..4c845b7bc1 100644 --- a/crates/rpc/build.rs +++ b/crates/rpc/build.rs @@ -30,9 +30,8 @@ fn main() -> Result<(), Box> { let derive_prost_builder = "#[cfg_attr(feature = \"test-support\", derive(carbide_prost_builder::Builder))]"; - tonic_prost_build::configure() - .file_descriptor_set_path(reflection) + .file_descriptor_set_path(&reflection) .type_attribute( ".google.protobuf.Timestamp", "#[derive(serde::Serialize, serde::Deserialize)]", @@ -1171,6 +1170,8 @@ fn main() -> Result<(), Box> { client_wrapper_generator.write_rpc_client_wrapper(out_dir.join("forge_api_client.rs"))?; client_wrapper_generator .write_rpc_convenience_converters(out_dir.join("convenience_converters.rs"))?; + client_wrapper_generator + .write_admin_cli_commands(&reflection, out_dir.join("admin_cli_commands.rs"))?; // enable the code generator for the nmx-c proto let nmx_c_client_wrapper = codegen::CodeGenerator::new(codegen::Config { diff --git a/crates/rpc/proto/admin_cli.proto b/crates/rpc/proto/admin_cli.proto new file mode 100644 index 0000000000..4680194a92 --- /dev/null +++ b/crates/rpc/proto/admin_cli.proto @@ -0,0 +1,36 @@ +syntax = "proto3"; + +package admin_cli; + +import "google/protobuf/descriptor.proto"; + +// Marks a unary RPC as eligible for generated nico-admin-cli plumbing. +message Command { + // Responses may only be ignored when this is explicitly acknowledged. + bool discard_response = 1; +} + +enum ArgumentKind { + POSITIONAL = 0; + OPTION = 1; +} + +// Describes how one request field is exposed by clap. +message Argument { + ArgumentKind kind = 1; + string long = 2; + string short = 3; + string help = 4; + + // A required protobuf message is represented by Option in prost. This + // unwraps that representation in the generated CLI argument structure. + bool required = 5; +} + +extend google.protobuf.MethodOptions { + Command command = 51000; +} + +extend google.protobuf.FieldOptions { + Argument argument = 51001; +} diff --git a/crates/rpc/proto/forge.proto b/crates/rpc/proto/forge.proto index 0a39a5259f..09c11baba9 100644 --- a/crates/rpc/proto/forge.proto +++ b/crates/rpc/proto/forge.proto @@ -6,6 +6,7 @@ package forge; // message types! import "common.proto"; +import "admin_cli.proto"; import "dns.proto"; import "google/protobuf/duration.proto"; import "google/protobuf/empty.proto"; @@ -702,7 +703,11 @@ service Forge { rpc TpmAddCaCert(TpmCaCert) returns (TpmCaAddedCaStatus); rpc TpmShowCaCerts(google.protobuf.Empty) returns (TpmCaCertDetailCollection); rpc TpmShowUnmatchedEkCerts(google.protobuf.Empty) returns (TpmEkCertStatusCollection); - rpc TpmDeleteCaCert(TpmCaCertId) returns (google.protobuf.Empty); + rpc TpmDeleteCaCert(TpmCaCertId) returns (google.protobuf.Empty) { + option (admin_cli.command) = { + discard_response: true + }; + } rpc RedfishBrowse(RedfishBrowseRequest) returns (RedfishBrowseResponse); rpc RedfishListActions(RedfishListActionsRequest) returns (RedfishListActionsResponse); @@ -813,7 +818,11 @@ service Forge { rpc RemediationApplied(RemediationAppliedRequest) returns (google.protobuf.Empty); // end DPU Remediation APIs - rpc SetPrimaryDpu(SetPrimaryDpuRequest) returns (google.protobuf.Empty); + rpc SetPrimaryDpu(SetPrimaryDpuRequest) returns (google.protobuf.Empty) { + option (admin_cli.command) = { + discard_response: true + }; + } // Make any host interface (DPU or not) the primary / boot interface. This is // the generic form of SetPrimaryDpu, which it deprecates. rpc SetPrimaryInterface(SetPrimaryInterfaceRequest) returns (google.protobuf.Empty); @@ -1249,7 +1258,12 @@ message TpmCaAddedCaStatus{ } message TpmCaCertId{ - int32 ca_cert_id = 1; + int32 ca_cert_id = 1 [(admin_cli.argument) = { + kind: OPTION + long: "ca-id" + short: "c" + help: "TPM CA id obtained from the show command" + }]; } message TpmEkCertStatus{ @@ -8331,9 +8345,19 @@ message RemediationApplicationStatus { // end DPU remediation models message SetPrimaryDpuRequest { - common.MachineId host_machine_id = 1; - common.MachineId dpu_machine_id = 2; - bool reboot=3; + common.MachineId host_machine_id = 1 [(admin_cli.argument) = { + help: "ID of the host machine" + required: true + }]; + common.MachineId dpu_machine_id = 2 [(admin_cli.argument) = { + help: "ID of the DPU machine to make primary" + required: true + }]; + bool reboot = 3 [(admin_cli.argument) = { + kind: OPTION + long: "reboot" + help: "Reboot the host after the update" + }]; } message SetPrimaryInterfaceRequest { diff --git a/crates/rpc/src/admin_cli.rs b/crates/rpc/src/admin_cli.rs index a4da4fdfa6..4bad990d8f 100644 --- a/crates/rpc/src/admin_cli.rs +++ b/crates/rpc/src/admin_cli.rs @@ -25,6 +25,14 @@ pub use output::OutputFormat; +#[cfg(feature = "cli")] +pub trait CliRpcCommand { + fn execute( + self, + client: &crate::forge_api_client::ForgeApiClient, + ) -> impl std::future::Future>; +} + pub mod output { use std::fmt; diff --git a/crates/rpc/src/protos/mod.rs b/crates/rpc/src/protos/mod.rs index 745ff557d6..6c6c536ff0 100644 --- a/crates/rpc/src/protos/mod.rs +++ b/crates/rpc/src/protos/mod.rs @@ -35,6 +35,9 @@ pub mod scout_firmware_upgrade { #[rustfmt::skip] pub mod forge { include!(concat!(env!("OUT_DIR"), "/forge.rs")); + + #[cfg(feature = "cli")] + include!(concat!(env!("OUT_DIR"), "/admin_cli_commands.rs")); } #[allow(non_snake_case, unknown_lints, clippy::all)] diff --git a/crates/tonic-client-wrapper/Cargo.toml b/crates/tonic-client-wrapper/Cargo.toml index e408d399bd..8c2d1765b0 100644 --- a/crates/tonic-client-wrapper/Cargo.toml +++ b/crates/tonic-client-wrapper/Cargo.toml @@ -30,7 +30,9 @@ codegen = [ "dep:syn", "dep:prettyplease", "dep:heck", + "dep:prost", "dep:prost-types", + "dep:protobuf", "dep:lazy_static", ] @@ -44,7 +46,9 @@ quote = { optional = true, workspace = true } syn = { optional = true, workspace = true } prettyplease = { optional = true, workspace = true } heck = { optional = true, workspace = true } +prost = { optional = true, workspace = true } prost-types = { optional = true, workspace = true } +protobuf = { optional = true, workspace = true } lazy_static = { optional = true, workspace = true } [dev-dependencies] diff --git a/crates/tonic-client-wrapper/src/codegen.rs b/crates/tonic-client-wrapper/src/codegen.rs index a3e5c89035..12fe5858ba 100644 --- a/crates/tonic-client-wrapper/src/codegen.rs +++ b/crates/tonic-client-wrapper/src/codegen.rs @@ -21,8 +21,10 @@ use std::path::Path; use heck::{ToSnakeCase, ToUpperCamelCase}; use proc_macro2::{LexError, TokenStream}; +use prost::Message as _; use prost_types::field_descriptor_proto::Label; -use prost_types::{FileDescriptorProto, MethodDescriptorProto}; +use prost_types::{FieldDescriptorProto, FileDescriptorProto, MethodDescriptorProto}; +use protobuf::Message as _; use quote::{TokenStreamExt, quote}; use crate::utils::{base_types, field_is_optional, resolve_field_primitive_type}; @@ -38,6 +40,10 @@ pub enum Error { Lex(#[from] LexError), #[error("invalid protobuf type: {0}")] InvalidProtobufType(String), + #[error("invalid admin CLI annotation: {0}")] + InvalidAdminCliAnnotation(String), + #[error("failed to decode protobuf descriptor set: {0}")] + DescriptorDecode(String), #[error(transparent)] Io(#[from] std::io::Error), #[error("syntax error in generated code: {0}")] @@ -273,6 +279,214 @@ impl CodeGenerator { Ok(()) } + /// Generate clap argument types and RPC dispatch implementations from the + /// `admin_cli.command` and `admin_cli.argument` protobuf options. + pub fn write_admin_cli_commands, Q: AsRef>( + &self, + descriptor_set: P, + out: Q, + ) -> Result<()> { + let descriptor_bytes = fs::read(descriptor_set)?; + let raw_descriptors = + protobuf::descriptor::FileDescriptorSet::parse_from_bytes(&descriptor_bytes) + .map_err(|error| Error::DescriptorDecode(error.to_string()))?; + let field_options = collect_admin_cli_field_options(&raw_descriptors)?; + + let mut commands = TokenStream::new(); + for file in &raw_descriptors.file { + for service in &file.service { + for method in &service.method { + let Some(command) = decode_admin_cli_command(method.options.as_ref())? else { + continue; + }; + + let method_name = method.name(); + if method.client_streaming() || method.server_streaming() { + return Err(invalid_cli_annotation(format!( + "{method_name}: streaming RPCs are not supported" + ))); + } + if !command.discard_response { + return Err(invalid_cli_annotation(format!( + "{method_name}: discard_response must be explicitly enabled" + ))); + } + + let input_name = method.input_type(); + let input = self.message_types.get(input_name).ok_or_else(|| { + invalid_cli_annotation(format!( + "{method_name}: input message {input_name} was not found" + )) + })?; + commands.append_all(self.make_admin_cli_command( + method_name, + input, + &field_options, + )?); + } + } + } + + write_token_stream_if_not_up_to_date(commands, out) + } + + fn make_admin_cli_command( + &self, + method_name: &str, + input: &MessageWithPackage, + field_options: &HashMap<(String, i32), AdminCliArgumentOption>, + ) -> Result { + let command_module: TokenStream = method_name.to_snake_case().parse()?; + let request_type: TokenStream = self + .convert_protobuf_type_to_rust_type(&input.qualified_name())? + .parse()?; + let rpc_method: TokenStream = method_name.to_snake_case().parse()?; + + let mut clap_fields = TokenStream::new(); + let mut request_fields = TokenStream::new(); + let mut seen_longs = HashSet::new(); + let mut seen_shorts = HashSet::new(); + + for field in &input.message.field { + let Some(option) = field_options.get(&(input.qualified_name(), field.number())) else { + continue; + }; + let field_name: TokenStream = field.name().to_snake_case().parse()?; + let field_type = self.admin_cli_field_type(method_name, field, option.required)?; + let help = &option.help; + if help.is_empty() { + return Err(invalid_cli_annotation(format!( + "{method_name}.{}: help must not be empty", + field.name() + ))); + } + + let clap_attribute = match AdminCliArgumentKind::try_from(option.kind).map_err( + |_| { + invalid_cli_annotation(format!( + "{method_name}.{}: unknown argument kind {}", + field.name(), + option.kind + )) + }, + )? { + AdminCliArgumentKind::Positional => { + if !option.long.is_empty() || !option.short.is_empty() { + return Err(invalid_cli_annotation(format!( + "{method_name}.{}: positional arguments cannot have long or short names", + field.name() + ))); + } + quote! { #[arg(help = #help)] } + } + AdminCliArgumentKind::Option => { + if option.long.is_empty() { + return Err(invalid_cli_annotation(format!( + "{method_name}.{}: option arguments require a long name", + field.name() + ))); + } + if !seen_longs.insert(option.long.clone()) { + return Err(invalid_cli_annotation(format!( + "{method_name}: duplicate --{} option", + option.long + ))); + } + let long = &option.long; + if option.short.is_empty() { + quote! { #[arg(long = #long, help = #help)] } + } else { + let mut chars = option.short.chars(); + let short = chars.next().expect("short option is not empty"); + if chars.next().is_some() || !seen_shorts.insert(short) { + return Err(invalid_cli_annotation(format!( + "{method_name}.{}: short option must be one unique character", + field.name() + ))); + } + quote! { #[arg(long = #long, short = #short, help = #help)] } + } + } + }; + clap_fields.extend(quote! { + #clap_attribute + pub #field_name: #field_type, + }); + + let request_value = if option.required && field_is_optional(field) { + quote! { Some(args.#field_name) } + } else { + quote! { args.#field_name } + }; + request_fields.extend(quote! { #field_name: #request_value, }); + } + + Ok(quote! { + pub mod #command_module { + #[derive(::clap::Args, Debug)] + pub struct Args { + #clap_fields + } + + impl From for #request_type { + fn from(args: Args) -> Self { + Self { + #request_fields + ..Default::default() + } + } + } + + impl crate::admin_cli::CliRpcCommand for Args { + async fn execute( + self, + client: &crate::forge_api_client::ForgeApiClient, + ) -> Result<(), ::tonic::Status> { + client.#rpc_method(self).await?; + Ok(()) + } + } + } + }) + } + + fn admin_cli_field_type( + &self, + method_name: &str, + field: &FieldDescriptorProto, + required: bool, + ) -> Result { + if field.oneof_index.is_some() { + return Err(invalid_cli_annotation(format!( + "{method_name}.{}: oneof fields are not supported", + field.name() + ))); + } + if field.label == Some(Label::Repeated as i32) { + return Err(invalid_cli_annotation(format!( + "{method_name}.{}: repeated fields are not supported", + field.name() + ))); + } + + let base_type = if let Some(primitive) = resolve_field_primitive_type(field) { + primitive + } else if let Some(type_name) = &field.type_name { + self.convert_protobuf_type_to_rust_type(type_name)? + } else { + return Err(invalid_cli_annotation(format!( + "{method_name}.{}: unsupported protobuf field type", + field.name() + ))); + }; + + if field_is_optional(field) && !required { + format!("Option<{base_type}>").parse().map_err(Into::into) + } else { + base_type.parse().map_err(Into::into) + } + } + fn make_convenience_converter( &self, message_with_package: &MessageWithPackage, @@ -619,6 +833,113 @@ impl MessageWithPackage { } } +const ADMIN_CLI_COMMAND_EXTENSION: u32 = 51_000; +const ADMIN_CLI_ARGUMENT_EXTENSION: u32 = 51_001; + +#[derive(Clone, PartialEq, prost::Message)] +struct AdminCliCommandOption { + #[prost(bool, tag = "1")] + discard_response: bool, +} + +#[derive(Clone, PartialEq, prost::Message)] +struct AdminCliArgumentOption { + #[prost(enumeration = "AdminCliArgumentKind", tag = "1")] + kind: i32, + #[prost(string, tag = "2")] + long: String, + #[prost(string, tag = "3")] + short: String, + #[prost(string, tag = "4")] + help: String, + #[prost(bool, tag = "5")] + required: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, prost::Enumeration)] +enum AdminCliArgumentKind { + Positional = 0, + Option = 1, +} + +fn invalid_cli_annotation(message: String) -> Error { + Error::InvalidAdminCliAnnotation(message) +} + +fn decode_admin_cli_command( + options: Option<&protobuf::descriptor::MethodOptions>, +) -> Result> { + let Some(payload) = unknown_length_delimited(options, ADMIN_CLI_COMMAND_EXTENSION)? else { + return Ok(None); + }; + AdminCliCommandOption::decode(payload) + .map(Some) + .map_err(|error| invalid_cli_annotation(format!("invalid command option: {error}"))) +} + +fn decode_admin_cli_argument( + options: Option<&protobuf::descriptor::FieldOptions>, +) -> Result> { + let Some(payload) = unknown_length_delimited(options, ADMIN_CLI_ARGUMENT_EXTENSION)? else { + return Ok(None); + }; + AdminCliArgumentOption::decode(payload) + .map(Some) + .map_err(|error| invalid_cli_annotation(format!("invalid argument option: {error}"))) +} + +trait HasUnknownFields { + fn unknown_fields(&self) -> &protobuf::UnknownFields; +} + +impl HasUnknownFields for protobuf::descriptor::MethodOptions { + fn unknown_fields(&self) -> &protobuf::UnknownFields { + self.special_fields.unknown_fields() + } +} + +impl HasUnknownFields for protobuf::descriptor::FieldOptions { + fn unknown_fields(&self) -> &protobuf::UnknownFields { + self.special_fields.unknown_fields() + } +} + +fn unknown_length_delimited( + options: Option<&T>, + extension: u32, +) -> Result> { + let Some(value) = options.and_then(|options| options.unknown_fields().get(extension)) else { + return Ok(None); + }; + match value { + protobuf::UnknownValueRef::LengthDelimited(payload) => Ok(Some(payload)), + _ => Err(invalid_cli_annotation(format!( + "option extension {extension} has an invalid wire type" + ))), + } +} + +fn collect_admin_cli_field_options( + descriptors: &protobuf::descriptor::FileDescriptorSet, +) -> Result> { + let mut result = HashMap::new(); + for file in &descriptors.file { + for message in &file.message_type { + let qualified_name = if file.package().is_empty() { + format!(".{}", message.name()) + } else { + format!(".{}.{}", file.package(), message.name()) + }; + for field in &message.field { + if let Some(option) = decode_admin_cli_argument(field.options.as_ref())? { + result.insert((qualified_name.clone(), field.number()), option); + } + } + } + } + Ok(result) +} + fn write_token_stream_if_not_up_to_date>( token_stream: TokenStream, out: T, @@ -1011,4 +1332,58 @@ mod tests { ); } } + + #[test] + fn test_decode_admin_cli_options() { + let expected = AdminCliCommandOption { + discard_response: true, + }; + let mut options = protobuf::descriptor::MethodOptions::new(); + options + .special_fields + .mut_unknown_fields() + .add_length_delimited(ADMIN_CLI_COMMAND_EXTENSION, expected.encode_to_vec()); + + assert_eq!( + decode_admin_cli_command(Some(&options)).unwrap(), + Some(expected) + ); + } + + #[test] + fn test_make_admin_cli_command() { + let generator = test_generator(include_str!("test_fixtures/test.proto")); + let input = generator.message_types.get(".MultiRequest").unwrap(); + let field_options = HashMap::from([ + ( + (".MultiRequest".to_string(), 1), + AdminCliArgumentOption { + help: "First value".to_string(), + ..Default::default() + }, + ), + ( + (".MultiRequest".to_string(), 2), + AdminCliArgumentOption { + kind: AdminCliArgumentKind::Option as i32, + long: "second".to_string(), + short: "s".to_string(), + help: "Second value".to_string(), + ..Default::default() + }, + ), + ]); + let generated = generator + .make_admin_cli_command("MultiRpc", input, &field_options) + .unwrap() + .to_string(); + + assert!(generated.contains("pub mod multi_rpc")); + assert!(generated.contains("pub struct Args")); + assert!(generated.contains("arg (help = \"First value\")")); + assert!(generated.contains("long = \"second\"")); + assert!(generated.contains("short = 's'")); + assert!(generated.contains("impl From < Args > for crate :: test :: MultiRequest")); + assert!(generated.contains("client . multi_rpc (self) . await")); + } }