Skip to content

Commit 9efc5bd

Browse files
committed
Move structure filling and URL selection code
This change aims to move structure filling and URL selection related code to specific modules, with a reduce in main.rs lines of code as a consequence Signed-off-by: Sergio Arroutbi <sarroutb@redhat.com>
1 parent 0da07db commit 9efc5bd

File tree

3 files changed

+115
-124
lines changed

3 files changed

+115
-124
lines changed

keylime-push-model-agent/src/main.rs

Lines changed: 79 additions & 115 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// Copyright 2025 Keylime Authors
3-
use crate::struct_filler::StructureFiller;
43
use anyhow::{Context, Result};
54
use clap::Parser;
65
use log::{debug, error, info};
@@ -16,20 +15,79 @@ mod url_selector;
1615
const DEFAULT_TIMEOUT_MILLIS: &str = "5000";
1716
const HTTPS_PREFIX: &str = "https://";
1817
const DEFAULT_METHOD: &str = "POST";
18+
const DEFAULT_MESSAGE_TYPE: MessageType = MessageType::Attestation;
19+
const DEFAULT_MESSAGE_TYPE_STR: &str = "Attestation";
1920

20-
enum MessageType {
21+
pub enum MessageType {
2122
Attestation,
2223
EvidenceHandling,
2324
Session,
2425
}
25-
const DEFAULT_MESSAGE_TYPE: MessageType = MessageType::Attestation;
26-
const DEFAULT_MESSAGE_TYPE_STR: &str = "Attestation";
2726

2827
pub struct ResponseInformation {
2928
pub status_code: reqwest::StatusCode,
3029
pub body: String,
3130
}
3231

32+
#[derive(Parser, Debug)]
33+
#[command(author, version, about, long_about = None, ignore_errors = true)]
34+
struct Args {
35+
/// API version
36+
/// Default: "v3.0"
37+
#[arg(long, default_value = url_selector::DEFAULT_API_VERSION)]
38+
api_version: Option<String>,
39+
/// CA certificate file
40+
#[arg(long, default_value = "/var/lib/keylime/cv_ca/cacert.crt")]
41+
ca_certificate: String,
42+
/// Client certificate file
43+
#[arg(
44+
short,
45+
long,
46+
default_value = "/var/lib/keylime/cv_ca/client-cert.crt"
47+
)]
48+
certificate: String,
49+
/// Client private key file
50+
#[arg(
51+
short,
52+
long,
53+
default_value = "/var/lib/keylime/cv_ca/client-private.pem"
54+
)]
55+
key: String,
56+
/// json file
57+
#[arg(short, long, default_missing_value = "")]
58+
json_file: Option<String>,
59+
/// identifier
60+
/// Default: 12345678
61+
#[arg(long, default_value = "12345678")]
62+
agent_identifier: String,
63+
/// index
64+
/// Default: 1
65+
#[arg(long, default_value = "1")]
66+
attestation_index: Option<String>,
67+
/// insecure
68+
#[arg(long, action, default_missing_value = "true")]
69+
insecure: Option<bool>,
70+
/// Type of message
71+
/// Default: "Attestation"
72+
#[arg(long, default_value = DEFAULT_MESSAGE_TYPE_STR)]
73+
message_type: Option<String>,
74+
/// Method
75+
/// Default: "POST"
76+
#[arg(long, default_missing_value = DEFAULT_METHOD)]
77+
method: Option<String>,
78+
/// Session ID
79+
/// Default: 1
80+
#[arg(long, default_missing_value = "1", default_value = "1")]
81+
session_index: Option<String>,
82+
/// Timeout in milliseconds
83+
/// Default: 5000
84+
#[arg(long, default_value = DEFAULT_TIMEOUT_MILLIS)]
85+
timeout: u64,
86+
/// Verifier URL
87+
#[arg(short, long, default_value = "https://127.0.0.1:8881")]
88+
verifier_url: String,
89+
}
90+
3391
fn get_message_type(args: &Args) -> MessageType {
3492
if args.message_type.is_some() {
3593
match args.message_type.as_ref().unwrap().as_str() {
@@ -105,70 +163,35 @@ fn get_client(args: &Args) -> Result<reqwest::Client> {
105163
.context("Failed to build plain HTTP client")
106164
}
107165

108-
fn get_filler_request(args: &Args) -> Box<dyn StructureFiller> {
109-
if args.json_file.is_none() {
110-
return Box::new(struct_filler::FillerFromCode {});
111-
}
112-
Box::new(struct_filler::FillerFromFile {
113-
file_path: args.json_file.clone().unwrap(),
114-
})
115-
}
116-
117-
fn get_url_from_message_type(args: &Args) -> String {
166+
fn get_request_builder_from_method(
167+
args: &Args,
168+
) -> Result<reqwest::RequestBuilder> {
169+
let client = get_client(args)?;
118170
let url_args = url_selector::UrlArgs {
119171
verifier_url: args.verifier_url.clone(),
120172
api_version: args.api_version.clone(),
121173
session_index: args.session_index.clone(),
122174
agent_identifier: Some(args.agent_identifier.clone()),
123175
attestation_index: args.attestation_index.clone(),
124176
};
125-
match get_message_type(args) {
126-
MessageType::Attestation => {
127-
url_selector::get_attestation_request_url(&url_args)
128-
}
129-
MessageType::EvidenceHandling => {
130-
if args.attestation_index.is_some() {
131-
url_selector::get_evidence_handling_request_url_with_index(
132-
&url_args,
133-
)
134-
} else {
135-
url_selector::get_evidence_handling_request_url(&url_args)
136-
}
137-
}
138-
MessageType::Session => {
139-
url_selector::get_session_request_url(&url_args)
140-
}
141-
}
142-
}
143-
144-
fn get_request_builder_from_method(
145-
args: &Args,
146-
) -> Result<reqwest::RequestBuilder> {
147-
let client = get_client(args)?;
177+
let url = url_selector::get_url_from_message_type(
178+
&url_args,
179+
&get_message_type(args),
180+
);
148181
match args.method.as_deref() {
149-
Some("POST") => {
150-
Ok(client.post(get_url_from_message_type(args).as_str()))
151-
}
152-
Some("PUT") => {
153-
Ok(client.put(get_url_from_message_type(args).as_str()))
154-
}
155-
Some("DELETE") => {
156-
Ok(client.delete(get_url_from_message_type(args).as_str()))
157-
}
158-
Some("GET") => {
159-
Ok(client.get(get_url_from_message_type(args).as_str()))
160-
}
161-
Some("PATCH") => {
162-
Ok(client.patch(get_url_from_message_type(args).as_str()))
163-
}
164-
_ => Ok(client.post(get_url_from_message_type(args).as_str())),
182+
Some("POST") => Ok(client.post(url)),
183+
Some("PUT") => Ok(client.put(url)),
184+
Some("DELETE") => Ok(client.delete(url)),
185+
Some("GET") => Ok(client.get(url)),
186+
Some("PATCH") => Ok(client.patch(url)),
187+
_ => Ok(client.post(url)),
165188
}
166189
}
167190

168191
async fn send_push_model_request(args: &Args) -> Result<ResponseInformation> {
169-
let filler = get_filler_request(args);
170-
171-
let json_value = match get_message_type(args) {
192+
let filler = struct_filler::get_filler_request(args.json_file.clone());
193+
let message_type = get_message_type(args);
194+
let json_value = match message_type {
172195
MessageType::Attestation => {
173196
json_dump::dump_attestation_request_to_value(
174197
&filler.get_attestation_request(),
@@ -205,65 +228,6 @@ async fn send_push_model_request(args: &Args) -> Result<ResponseInformation> {
205228
Ok(rsp)
206229
}
207230

208-
#[derive(Parser, Debug)]
209-
#[command(author, version, about, long_about = None, ignore_errors = true)]
210-
struct Args {
211-
/// API version
212-
/// Default: "v3.0"
213-
#[arg(long, default_value = url_selector::DEFAULT_API_VERSION)]
214-
api_version: Option<String>,
215-
/// CA certificate file
216-
#[arg(long, default_value = "/var/lib/keylime/cv_ca/cacert.crt")]
217-
ca_certificate: String,
218-
/// Client certificate file
219-
#[arg(
220-
short,
221-
long,
222-
default_value = "/var/lib/keylime/cv_ca/client-cert.crt"
223-
)]
224-
certificate: String,
225-
/// Client private key file
226-
#[arg(
227-
short,
228-
long,
229-
default_value = "/var/lib/keylime/cv_ca/client-private.pem"
230-
)]
231-
key: String,
232-
/// json file
233-
#[arg(short, long, default_missing_value = "")]
234-
json_file: Option<String>,
235-
/// identifier
236-
/// Default: 12345678
237-
#[arg(long, default_value = "12345678")]
238-
agent_identifier: String,
239-
/// index
240-
/// Default: 1
241-
#[arg(long, default_value = "1")]
242-
attestation_index: Option<String>,
243-
/// insecure
244-
#[arg(long, action, default_missing_value = "true")]
245-
insecure: Option<bool>,
246-
/// Type of message
247-
/// Default: "Attestation"
248-
#[arg(long, default_value = DEFAULT_MESSAGE_TYPE_STR)]
249-
message_type: Option<String>,
250-
/// Method
251-
/// Default: "POST"
252-
#[arg(long, default_missing_value = DEFAULT_METHOD)]
253-
method: Option<String>,
254-
/// Session ID
255-
/// Default: 1
256-
#[arg(long, default_missing_value = "1", default_value = "1")]
257-
session_index: Option<String>,
258-
/// Timeout in milliseconds
259-
/// Default: 5000
260-
#[arg(long, default_value = DEFAULT_TIMEOUT_MILLIS)]
261-
timeout: u64,
262-
/// Verifier URL
263-
#[arg(short, long, default_value = "https://127.0.0.1:8881")]
264-
verifier_url: String,
265-
}
266-
267231
async fn run() -> Result<()> {
268232
let args = Args::parse();
269233
info!("Verifier URL: {}", args.verifier_url);

keylime-push-model-agent/src/struct_filler.rs

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,17 @@ pub trait StructureFiller {
1717
) -> structures::EvidenceHandlingRequest;
1818
}
1919

20+
pub fn get_filler_request(
21+
json_file: Option<String>,
22+
) -> Box<dyn StructureFiller> {
23+
if json_file.is_none() {
24+
return Box::new(FillerFromCode {});
25+
}
26+
Box::new(FillerFromFile {
27+
file_path: json_file.clone().unwrap(),
28+
})
29+
}
30+
2031
pub struct FillerFromCode;
2132
impl StructureFiller for FillerFromCode {
2233
fn get_attestation_request(&self) -> structures::AttestationRequest {

keylime-push-model-agent/src/url_selector.rs

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// Copyright 2025 Keylime Authors
33
pub const DEFAULT_API_VERSION: &str = "v3.0";
44
const DEFAULT_INDEX: &str = "1";
5+
use crate::MessageType;
56

67
pub struct UrlArgs {
78
pub verifier_url: String,
@@ -11,18 +12,35 @@ pub struct UrlArgs {
1112
pub session_index: Option<String>,
1213
}
1314

15+
pub fn get_attestation_index(args: &UrlArgs) -> String {
16+
if args.attestation_index.is_some() {
17+
return args.attestation_index.clone().unwrap();
18+
}
19+
DEFAULT_INDEX.to_string()
20+
}
21+
1422
fn get_api_version(args: &UrlArgs) -> String {
1523
if args.api_version.is_some() {
1624
return args.api_version.clone().unwrap();
1725
}
1826
DEFAULT_API_VERSION.to_string()
1927
}
2028

21-
pub fn get_attestation_index(args: &UrlArgs) -> String {
22-
if args.attestation_index.is_some() {
23-
return args.attestation_index.clone().unwrap();
29+
pub fn get_url_from_message_type(
30+
url_args: &UrlArgs,
31+
message_type: &MessageType,
32+
) -> String {
33+
match message_type {
34+
MessageType::Attestation => get_attestation_request_url(url_args),
35+
MessageType::EvidenceHandling => {
36+
if url_args.attestation_index.is_some() {
37+
get_evidence_handling_request_url_with_index(url_args)
38+
} else {
39+
get_evidence_handling_request_url(url_args)
40+
}
41+
}
42+
MessageType::Session => get_session_request_url(url_args),
2443
}
25-
DEFAULT_INDEX.to_string()
2644
}
2745

2846
fn get_index_suffix(args: &UrlArgs) -> String {
@@ -33,7 +51,7 @@ fn get_index_suffix(args: &UrlArgs) -> String {
3351
"".to_string()
3452
}
3553

36-
pub fn get_attestation_request_url(args: &UrlArgs) -> String {
54+
fn get_attestation_request_url(args: &UrlArgs) -> String {
3755
let id = args.agent_identifier.clone().unwrap();
3856
let verifier_url = args.verifier_url.clone();
3957
let api_version = get_api_version(args);
@@ -45,7 +63,7 @@ pub fn get_attestation_request_url(args: &UrlArgs) -> String {
4563
format!("{verifier_url}/{api_version}/agents/{id}/attestations")
4664
}
4765

48-
pub fn get_evidence_handling_request_url(args: &UrlArgs) -> String {
66+
fn get_evidence_handling_request_url(args: &UrlArgs) -> String {
4967
let id = args.agent_identifier.clone().unwrap();
5068
let verifier_url = args.verifier_url.clone();
5169
let api_version = get_api_version(args);
@@ -57,9 +75,7 @@ pub fn get_evidence_handling_request_url(args: &UrlArgs) -> String {
5775
format!("{verifier_url}/{api_version}/agents/{id}/attestations")
5876
}
5977

60-
pub fn get_evidence_handling_request_url_with_index(
61-
args: &UrlArgs,
62-
) -> String {
78+
fn get_evidence_handling_request_url_with_index(args: &UrlArgs) -> String {
6379
let id = args.agent_identifier.clone().unwrap();
6480
let verifier_url = args.verifier_url.clone();
6581
let api_version = get_api_version(args);

0 commit comments

Comments
 (0)