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")); + } }