Skip to content

Commit 82e517f

Browse files
committed
[DRAFT]: Code refactoring
Signed-off-by: Sergio Arroutbi <[email protected]>
1 parent be33d8b commit 82e517f

File tree

2 files changed

+143
-91
lines changed

2 files changed

+143
-91
lines changed

keylime-agent/src/agent_activation.rs

+125
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// Copyright 2025 Keylime Authors
3+
use crate::config;
4+
use crate::error::{Error, Result};
5+
use base64::{engine::general_purpose, Engine as _};
6+
use keylime::{
7+
crypto::{self, x509::CertificateBuilder},
8+
registrar_client::RegistrarClientBuilder,
9+
tpm::{self, IAKResult, IDevIDResult},
10+
};
11+
use log::{error, info};
12+
use openssl::x509::X509;
13+
use tss_esapi::handles::KeyHandle;
14+
use tss_esapi::structures::PublicBuffer;
15+
use tss_esapi::traits::Marshall;
16+
17+
#[derive(Debug)]
18+
pub struct AgentActivation {
19+
pub ak: tpm::AKResult,
20+
pub ek_result: tpm::EKResult,
21+
pub api_versions: Vec<String>,
22+
pub agent: config::AgentConfig,
23+
pub agent_uuid: String,
24+
pub mtls_cert: Option<X509>,
25+
pub device_id: Option<keylime::device_id::DeviceID>,
26+
pub attest: Option<tss_esapi::structures::Attest>,
27+
pub signature: Option<tss_esapi::structures::Signature>,
28+
pub ak_handle: KeyHandle,
29+
}
30+
31+
pub async fn activate_agent(
32+
mut aa: AgentActivation,
33+
mut ctx: &mut tpm::Context<'_>,
34+
) -> Result<()> {
35+
let iak_pub;
36+
let idevid_pub;
37+
let ak_pub = &PublicBuffer::try_from(aa.ak.public)?.marshall()?;
38+
let ek_pub =
39+
&PublicBuffer::try_from(aa.ek_result.public.clone())?.marshall()?;
40+
41+
// Create a RegistrarClientBuilder and set the parameters
42+
let mut builder = RegistrarClientBuilder::new()
43+
.ak_pub(ak_pub)
44+
.ek_pub(ek_pub)
45+
.enabled_api_versions(
46+
aa.api_versions.iter().map(|ver| ver.as_ref()).collect(),
47+
)
48+
.registrar_ip(aa.agent.registrar_ip.clone())
49+
.registrar_port(aa.agent.registrar_port)
50+
.uuid(&aa.agent_uuid)
51+
.ip(aa.agent.contact_ip.clone())
52+
.port(aa.agent.contact_port);
53+
54+
if let Some(mtls_cert) = aa.mtls_cert {
55+
builder = builder.mtls_cert(mtls_cert);
56+
}
57+
58+
// If the certificate is not None add it to the builder
59+
if let Some(ek_cert) = aa.ek_result.ek_cert {
60+
builder = builder.ek_cert(ek_cert);
61+
}
62+
63+
// Set the IAK/IDevID related fields, if enabled
64+
if aa.agent.enable_iak_idevid {
65+
let (Some(dev_id), Some(attest), Some(signature)) =
66+
(&aa.device_id, aa.attest, aa.signature)
67+
else {
68+
error!("IDevID and IAK are enabled but could not be generated");
69+
return Err(Error::Configuration(
70+
config::KeylimeConfigError::Generic(
71+
"IDevID and IAK are enabled but could not be generated"
72+
.to_string(),
73+
),
74+
));
75+
};
76+
77+
iak_pub =
78+
PublicBuffer::try_from(dev_id.iak_pubkey.clone())?.marshall()?;
79+
idevid_pub = PublicBuffer::try_from(dev_id.idevid_pubkey.clone())?
80+
.marshall()?;
81+
builder = builder
82+
.iak_attest(attest.marshall()?)
83+
.iak_sign(signature.marshall()?)
84+
.iak_pub(&iak_pub)
85+
.idevid_pub(&idevid_pub);
86+
87+
// If the IAK certificate was provided, set it
88+
if let Some(iak_cert) = dev_id.iak_cert.clone() {
89+
builder = builder.iak_cert(iak_cert);
90+
}
91+
92+
// If the IDevID certificate was provided, set it
93+
if let Some(idevid_cert) = dev_id.idevid_cert.clone() {
94+
builder = builder.idevid_cert(idevid_cert);
95+
}
96+
}
97+
98+
// Build the registrar client
99+
let mut registrar_client = builder.build().await?;
100+
101+
// Request keyblob material
102+
let keyblob = registrar_client.register_agent().await?;
103+
104+
info!("SUCCESS: Agent {} registered", &aa.agent_uuid);
105+
106+
let key = ctx.activate_credential(
107+
keyblob,
108+
aa.ak_handle,
109+
aa.ek_result.key_handle,
110+
)?;
111+
112+
// Flush EK if we created it
113+
if aa.agent.ek_handle.is_empty() {
114+
ctx.flush_context(aa.ek_result.key_handle.into())?;
115+
}
116+
117+
let mackey = general_purpose::STANDARD.encode(key.value());
118+
let auth_tag =
119+
crypto::compute_hmac(mackey.as_bytes(), aa.agent_uuid.as_bytes())?;
120+
let auth_tag = hex::encode(&auth_tag);
121+
122+
registrar_client.activate_agent(&auth_tag).await?;
123+
info!("SUCCESS: Agent {} activated", &aa.agent_uuid);
124+
Ok(())
125+
}

keylime-agent/src/main.rs

+18-91
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
// missing_docs: there is many functions missing documentations for now
3232
#![allow(unused, missing_docs)]
3333

34+
mod agent_activation;
3435
mod agent_handler;
3536
mod api;
3637
mod common;
@@ -46,6 +47,7 @@ mod revocation;
4647
mod secure_mount;
4748

4849
use actix_web::{dev::Service, http, middleware, rt, web, App, HttpServer};
50+
use agent_activation::AgentActivation;
4951
use base64::{engine::general_purpose, Engine as _};
5052
use clap::{Arg, Command as ClapApp};
5153
use common::*;
@@ -611,98 +613,23 @@ async fn main() -> Result<()> {
611613
warn!("mTLS disabled, Tenant and Verifier will reach out to agent via HTTP");
612614
}
613615

614-
{
615-
// Declare here as these must live longer than the builder
616-
let iak_pub;
617-
let idevid_pub;
618-
let ak_pub = &PublicBuffer::try_from(ak.public)?.marshall()?;
619-
let ek_pub =
620-
&PublicBuffer::try_from(ek_result.public.clone())?.marshall()?;
621-
622-
// Create a RegistrarClientBuilder and set the parameters
623-
let mut builder = RegistrarClientBuilder::new()
624-
.ak_pub(ak_pub)
625-
.ek_pub(ek_pub)
626-
.enabled_api_versions(
627-
api_versions.iter().map(|ver| ver.as_ref()).collect(),
628-
)
629-
.registrar_ip(config.agent.registrar_ip.clone())
630-
.registrar_port(config.agent.registrar_port)
631-
.uuid(&agent_uuid)
632-
.ip(config.agent.contact_ip.clone())
633-
.port(config.agent.contact_port);
634-
635-
if let Some(mtls_cert) = mtls_cert {
636-
builder = builder.mtls_cert(mtls_cert);
637-
}
638-
639-
// If the certificate is not None add it to the builder
640-
if let Some(ek_cert) = ek_result.ek_cert {
641-
builder = builder.ek_cert(ek_cert);
642-
}
643-
644-
// Set the IAK/IDevID related fields, if enabled
645-
if config.agent.enable_iak_idevid {
646-
let (Some(dev_id), Some(attest), Some(signature)) =
647-
(&device_id, attest, signature)
648-
else {
649-
error!(
650-
"IDevID and IAK are enabled but could not be generated"
651-
);
652-
return Err(Error::Configuration(config::KeylimeConfigError::Generic(
653-
"IDevID and IAK are enabled but could not be generated"
654-
.to_string(),
655-
)));
656-
};
657-
658-
iak_pub = PublicBuffer::try_from(dev_id.iak_pubkey.clone())?
659-
.marshall()?;
660-
idevid_pub =
661-
PublicBuffer::try_from(dev_id.idevid_pubkey.clone())?
662-
.marshall()?;
663-
builder = builder
664-
.iak_attest(attest.marshall()?)
665-
.iak_sign(signature.marshall()?)
666-
.iak_pub(&iak_pub)
667-
.idevid_pub(&idevid_pub);
668-
669-
// If the IAK certificate was provided, set it
670-
if let Some(iak_cert) = dev_id.iak_cert.clone() {
671-
builder = builder.iak_cert(iak_cert);
672-
}
673-
674-
// If the IDevID certificate was provided, set it
675-
if let Some(idevid_cert) = dev_id.idevid_cert.clone() {
676-
builder = builder.idevid_cert(idevid_cert);
677-
}
678-
}
679-
680-
// Build the registrar client
681-
let mut registrar_client = builder.build().await?;
682-
683-
// Request keyblob material
684-
let keyblob = registrar_client.register_agent().await?;
685-
686-
info!("SUCCESS: Agent {} registered", &agent_uuid);
687-
688-
let key = ctx.activate_credential(
689-
keyblob,
690-
ak_handle,
691-
ek_result.key_handle,
692-
)?;
693-
694-
// Flush EK if we created it
695-
if config.agent.ek_handle.is_empty() {
696-
ctx.flush_context(ek_result.key_handle.into())?;
616+
let aa = AgentActivation {
617+
ak,
618+
ek_result,
619+
api_versions: api_versions.clone(),
620+
agent: config.agent.clone(),
621+
agent_uuid: agent_uuid.clone(),
622+
mtls_cert,
623+
device_id,
624+
attest,
625+
signature,
626+
ak_handle,
627+
};
628+
match agent_activation::activate_agent(aa, &mut ctx).await {
629+
Ok(()) => (),
630+
Err(e) => {
631+
error!("Failed to activate agent: {}", e);
697632
}
698-
699-
let mackey = general_purpose::STANDARD.encode(key.value());
700-
let auth_tag =
701-
crypto::compute_hmac(mackey.as_bytes(), agent_uuid.as_bytes())?;
702-
let auth_tag = hex::encode(&auth_tag);
703-
704-
registrar_client.activate_agent(&auth_tag).await?;
705-
info!("SUCCESS: Agent {} activated", &agent_uuid);
706633
}
707634

708635
let (mut payload_tx, mut payload_rx) =

0 commit comments

Comments
 (0)