Skip to content

Commit

Permalink
Optimize the test case code of pki module.
Browse files Browse the repository at this point in the history
  • Loading branch information
@wa5i
wa5i committed Sep 18, 2024
1 parent e388354 commit 4bce448
Showing 1 changed file with 40 additions and 234 deletions.
274 changes: 40 additions & 234 deletions src/modules/pki/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -186,9 +186,7 @@ mod test {
assert!(test_write_api(&core, token, format!("{}roles/{}", path, role_name).as_str(), true, Some(role_data)).is_ok());
}

fn generate_root(core: &Core, token: &str, path: &str, exported: bool, is_ok: bool) {
let key_type = "rsa";
let key_bits = 4096;
fn generate_root(core: &Core, token: &str, path: &str, exported: bool, key_type: &str, key_bits: u32, is_ok: bool) {
let common_name = "test-ca";
let req_data = json!({
"common_name": common_name,
Expand Down Expand Up @@ -266,7 +264,7 @@ mod test {
assert_eq!(resp_ca_pem.unwrap_err(), RvError::ErrPkiCaNotConfig);
}

fn issue_cert_by_generate_root(core: &Core, token: &str, path: &str, role_name: &str) {
fn issue_cert_by_generate_root(core: &Core, token: &str, path: &str, role_name: &str, key_type: &str, key_bits: u32) {
let dns_sans = vec!["test.com", "a.test.com", "b.test.com"];
let issue_data = json!({
"ttl": "10d",
Expand All @@ -293,9 +291,9 @@ mod test {
for alt_name in alt_names {
assert!(dns_sans.contains(&alt_name.dnsname().unwrap()));
}
assert_eq!(cert_data["private_key_type"].as_str().unwrap(), "rsa");
assert_eq!(cert_data["private_key_type"].as_str().unwrap(), key_type);
let priv_key = PKey::private_key_from_pem(cert_data["private_key"].as_str().unwrap().as_bytes()).unwrap();
assert_eq!(priv_key.bits(), 4096);
assert_eq!(priv_key.bits(), key_bits);
assert!(priv_key.public_eq(&cert.public_key().unwrap()));
let serial_number = cert.serial_number().to_bn().unwrap();
let serial_number_hex = serial_number.to_hex_str().unwrap();
Expand Down Expand Up @@ -438,8 +436,8 @@ x/+V28hUf8m8P2NxP5ALaDZagdaMfzjGZo3O3wDv33Cds0P5GMGQYnRXDxcZN/2L
assert_eq!(role_data["ttl"].as_u64().unwrap(), 60 * 24 * 60 * 60);
assert_eq!(role_data["max_ttl"].as_u64().unwrap(), 365 * 24 * 60 * 60);
assert_eq!(role_data["not_before_duration"].as_u64().unwrap(), 30);
assert_eq!(role_data["key_type"].as_str().unwrap(), key_type);
assert_eq!(role_data["key_bits"].as_u64().unwrap(), key_bits);
assert_eq!(role_data["key_type"].as_str().unwrap(), "rsa");
assert_eq!(role_data["key_bits"].as_u64().unwrap(), 4096);
assert_eq!(role_data["country"].as_str().unwrap(), "CN");
assert_eq!(role_data["province"].as_str().unwrap(), "Beijing");
assert_eq!(role_data["locality"].as_str().unwrap(), "Beijing");
Expand Down Expand Up @@ -490,9 +488,9 @@ x/+V28hUf8m8P2NxP5ALaDZagdaMfzjGZo3O3wDv33Cds0P5GMGQYnRXDxcZN/2L
for alt_name in alt_names {
assert!(dns_sans.contains(&alt_name.dnsname().unwrap()));
}
assert_eq!(cert_data["private_key_type"].as_str().unwrap(), key_type);
assert_eq!(cert_data["private_key_type"].as_str().unwrap(), "rsa");
let priv_key = PKey::private_key_from_pem(cert_data["private_key"].as_str().unwrap().as_bytes()).unwrap();
assert_eq!(priv_key.bits(), key_bits);
assert_eq!(priv_key.bits(), 4096);
assert!(priv_key.public_eq(&cert.public_key().unwrap()));
let serial_number = cert.serial_number().to_bn().unwrap();
let serial_number_hex = serial_number.to_hex_str().unwrap();
Expand All @@ -516,13 +514,13 @@ x/+V28hUf8m8P2NxP5ALaDZagdaMfzjGZo3O3wDv33Cds0P5GMGQYnRXDxcZN/2L
let resp_lowercase = test_read_api(
&core,
token,
format!("{}/cert/{}", path, serial_number_hex.to_lowercase().as_str()).as_str(),
format!("{}cert/{}", path, serial_number_hex.to_lowercase().as_str()).as_str(),
true,
);
let resp_uppercase = test_read_api(
&core,
token,
format!("{}/cert/{}", path, serial_number_hex.to_uppercase().as_str()).as_str(),
format!("{}cert/{}", path, serial_number_hex.to_uppercase().as_str()).as_str(),
true,
);
println!("resp_uppercase: {:?}", resp_uppercase);
Expand Down Expand Up @@ -561,10 +559,35 @@ x/+V28hUf8m8P2NxP5ALaDZagdaMfzjGZo3O3wDv33Cds0P5GMGQYnRXDxcZN/2L
// config role
config_role(&core, token, path, role_name);

generate_root(&core, token, path, true, true);
generate_root(&core, token, path, false, true);
generate_root(&core, token, path, true, "rsa", 4096, true);
generate_root(&core, token, path, false, "rsa", 4096, true);

issue_cert_by_generate_root(&core, token, path, role_name, "rsa", 4096);

delete_root(&core, token, path, true);
}

#[cfg(feature = "crypto_adaptor_tongsuo")]
fn test_pki_sm2_generate_root() {
let (root_token, c) = test_rusty_vault_init("test_pki_generate_root");
let token = &root_token;
let core = c.read().unwrap();
let path = "sm2pki/";
let role_name = "test";

// mount pki backend to path: pki/
test_mount_api(&core, token, "pki", path);

// config ca
config_ca(&core, token, path);

// config role
config_role(&core, token, path, role_name);

generate_root(&core, token, path, true, "sm2", 256, true);
generate_root(&core, token, path, false, "sm2", 256, true);

issue_cert_by_generate_root(&core, token, path, role_name);
issue_cert_by_generate_root(&core, token, path, role_name, "sm2", 256);

delete_root(&core, token, path, true);
}
Expand Down Expand Up @@ -823,7 +846,7 @@ x/+V28hUf8m8P2NxP5ALaDZagdaMfzjGZo3O3wDv33Cds0P5GMGQYnRXDxcZN/2L
let (root_token, c) = test_rusty_vault_init("test_pki_generate_key");
let token = &root_token;
let core = c.read().unwrap();
let path = "pki/";
let path = "pki";

// mount pki backend to path: pki/
test_mount_api(&core, token, "pki", path);
Expand Down Expand Up @@ -922,7 +945,7 @@ x/+V28hUf8m8P2NxP5ALaDZagdaMfzjGZo3O3wDv33Cds0P5GMGQYnRXDxcZN/2L
let (root_token, c) = test_rusty_vault_init("test_pki_import_key");
let token = &root_token;
let core = c.read().unwrap();
let path = "pki/";
let path = "pki";

// mount pki backend to path: pki/
test_mount_api(&core, token, "pki", path);
Expand Down Expand Up @@ -1613,222 +1636,5 @@ xxxxxxxxxxxxxx
test_pki_encrypt_decrypt(&core, token, path, "aes-ecb-128-import", "rusty_vault test".as_bytes(), true);
test_pki_encrypt_decrypt(&core, token, path, "aes-ecb-192-import", "rusty_vault test".as_bytes(), true);
test_pki_encrypt_decrypt(&core, token, path, "aes-ecb-256-import", "rusty_vault test".as_bytes(), true);
<<<<<<< HEAD
}

fn test_pki_generate_root(core: Arc<RwLock<Core>>, token: &str, path: &str, key_type: &str, key_bits: u32, exported: bool, is_ok: bool) {
let core = core.read().unwrap();

let common_name = "test-ca";
let req_data = json!({
"common_name": common_name,
"ttl": "365d",
"country": "cn",
"key_type": key_type,
"key_bits": key_bits,
})
.as_object()
.unwrap()
.clone();
println!("generate root req_data: {:?}, is_ok: {}", req_data, is_ok);
let resp = test_write_api(
&core,
token,
format!("{}/root/generate/{}", path, if exported { "exported" } else { "internal" }).as_str(),
is_ok,
Some(req_data),
);
if !is_ok {
return;
}
let resp_body = resp.unwrap();
assert!(resp_body.is_some());
let data = resp_body.unwrap().data;
assert!(data.is_some());
let key_data = data.unwrap();
println!("generate root result: {:?}", key_data);

let resp_ca_pem = test_read_api(&core, token, &format!("{}/ca/pem", path), true);
let resp_ca_pem_cert_data = resp_ca_pem.unwrap().unwrap().data.unwrap();

let ca_cert = X509::from_pem(resp_ca_pem_cert_data["certificate"].as_str().unwrap().as_bytes()).unwrap();
let subject = ca_cert.subject_name();
let cn = subject.entries_by_nid(Nid::COMMONNAME).next().unwrap();
assert_eq!(cn.data().as_slice(), common_name.as_bytes());

let not_after = Asn1Time::days_from_now(365).unwrap();
let ttl_diff = ca_cert.not_after().diff(&not_after);
assert!(ttl_diff.is_ok());
let ttl_diff = ttl_diff.unwrap();
assert_eq!(ttl_diff.days, 0);

if exported {
assert!(key_data["private_key_type"].as_str().is_some());
assert_eq!(key_data["private_key_type"].as_str().unwrap(), key_type);
assert!(key_data["private_key"].as_str().is_some());
let private_key_pem = key_data["private_key"].as_str().unwrap();
match key_type {
"rsa" => {
let rsa_key = Rsa::private_key_from_pem(private_key_pem.as_bytes());
assert!(rsa_key.is_ok());
assert_eq!(rsa_key.unwrap().size() * 8, key_bits);
}
"sm2" | "ec" => {
let ec_key = EcKey::private_key_from_pem(private_key_pem.as_bytes());
assert!(ec_key.is_ok());
assert_eq!(ec_key.unwrap().group().degree(), key_bits);
}
_ => {}
}
} else {
assert!(key_data.get("private_key").is_none());
}
}

fn test_pki_delete_root(core: Arc<RwLock<Core>>, token: &str, path: &str, is_ok: bool) {
let core = core.read().unwrap();

let resp = test_delete_api(&core, token, &format!("{}/root", path), is_ok);
if !is_ok {
return;
}
assert!(resp.is_ok());

let resp_ca_pem = test_read_api(&core, token, &format!("{}/ca/pem", path), false);
assert_eq!(resp_ca_pem.unwrap_err(), RvError::ErrPkiCaNotConfig);
}

fn test_pki_issue_cert_by_generate_root(core: Arc<RwLock<Core>>, token: &str, path: &str, key_type: &str, key_bits: u32) {
let core = core.read().unwrap();

let dns_sans = vec!["test.com", "a.test.com", "b.test.com"];
let issue_data = json!({
"ttl": "10d",
"common_name": "test.com",
"alt_names": "a.test.com,b.test.com",
})
.as_object()
.unwrap()
.clone();

// issue cert
let resp = test_write_api(&core, token, &format!("{}/issue/test", path), true, Some(issue_data));
assert!(resp.is_ok());
let resp_body = resp.unwrap();
assert!(resp_body.is_some());
let data = resp_body.unwrap().data;
assert!(data.is_some());
let cert_data = data.unwrap();
let cert = X509::from_pem(cert_data["certificate"].as_str().unwrap().as_bytes()).unwrap();
let alt_names = cert.subject_alt_names();
assert!(alt_names.is_some());
let alt_names = alt_names.unwrap();
assert_eq!(alt_names.len(), dns_sans.len());
for alt_name in alt_names {
assert!(dns_sans.contains(&alt_name.dnsname().unwrap()));
}
assert_eq!(cert_data["private_key_type"].as_str().unwrap(), key_type);
let priv_key = PKey::private_key_from_pem(cert_data["private_key"].as_str().unwrap().as_bytes()).unwrap();
assert_eq!(priv_key.bits(), key_bits);
assert!(priv_key.public_eq(&cert.public_key().unwrap()));
let serial_number = cert.serial_number().to_bn().unwrap();
let serial_number_hex = serial_number.to_hex_str().unwrap();
assert_eq!(
cert_data["serial_number"].as_str().unwrap().replace(":", "").to_lowercase().as_str(),
serial_number_hex.to_lowercase().as_str()
);
let expiration_time = Asn1Time::from_unix(cert_data["expiration"].as_i64().unwrap()).unwrap();
let ttl_compare = cert.not_after().compare(&expiration_time);
assert!(ttl_compare.is_ok());
assert_eq!(ttl_compare.unwrap(), std::cmp::Ordering::Equal);
let now_timestamp = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
let expiration_ttl = cert_data["expiration"].as_u64().unwrap();
let ttl = expiration_ttl - now_timestamp;
let expect_ttl = 10 * 24 * 60 * 60;
assert!(ttl <= expect_ttl);
assert!((ttl + 10) > expect_ttl);

let authority_key_id = cert.authority_key_id();
assert!(authority_key_id.is_some());

println!("authority_key_id: {:?}", authority_key_id.unwrap().as_slice());

let resp_ca_pem = test_read_api(&core, token, &format!("{}/ca/pem", path), true);
let resp_ca_pem_cert_data = resp_ca_pem.unwrap().unwrap().data.unwrap();

let ca_cert = X509::from_pem(resp_ca_pem_cert_data["certificate"].as_str().unwrap().as_bytes()).unwrap();
let subject = ca_cert.subject_name();
let cn = subject.entries_by_nid(Nid::COMMONNAME).next().unwrap();
assert_eq!(cn.data().as_slice(), "test-ca".as_bytes());
println!("ca subject_key_id: {:?}", ca_cert.subject_key_id().unwrap().as_slice());
assert_eq!(ca_cert.subject_key_id().unwrap().as_slice(), authority_key_id.unwrap().as_slice());
}

#[test]
fn test_pki_module() {
let dir = env::temp_dir().join("rusty_vault_pki_module");
assert!(fs::create_dir(&dir).is_ok());
defer! (
assert!(fs::remove_dir_all(&dir).is_ok());
);

let mut root_token = String::new();
println!("root_token: {:?}", root_token);

let mut conf: HashMap<String, Value> = HashMap::new();
conf.insert("path".to_string(), Value::String(dir.to_string_lossy().into_owned()));

let backend = storage::new_backend("file", &conf).unwrap();

let barrier = storage::barrier_aes_gcm::AESGCMBarrier::new(Arc::clone(&backend));

let c = Arc::new(RwLock::new(Core { physical: backend, barrier: Arc::new(barrier), ..Default::default() }));

{
let mut core = c.write().unwrap();
assert!(core.config(Arc::clone(&c), None).is_ok());

let seal_config = SealConfig { secret_shares: 10, secret_threshold: 5 };

let result = core.init(&seal_config);
assert!(result.is_ok());
let init_result = result.unwrap();
println!("init_result: {:?}", init_result);

let mut unsealed = false;
for i in 0..seal_config.secret_threshold {
let key = &init_result.secret_shares[i as usize];
let unseal = core.unseal(key);
assert!(unseal.is_ok());
unsealed = unseal.unwrap();
}

root_token = init_result.root_token;

assert!(unsealed);
}

{
println!("root_token: {:?}", root_token);
test_pki_mount(Arc::clone(&c), &root_token, "pki");
test_pki_config_ca(Arc::clone(&c), &root_token, "pki");
test_pki_config_role(Arc::clone(&c), &root_token, "pki", "rsa", 4096);
test_pki_issue_cert(Arc::clone(&c), &root_token, "pki", "rsa", 4096);
test_pki_generate_key(Arc::clone(&c), &root_token, "pki");
test_pki_import_key(Arc::clone(&c), &root_token, "pki");
test_pki_generate_root(Arc::clone(&c), &root_token, "pki", "rsa", 4096, true, true);
test_pki_generate_root(Arc::clone(&c), &root_token, "pki", "rsa", 4096, false, true);
test_pki_issue_cert_by_generate_root(Arc::clone(&c), &root_token, "pki", "rsa", 4096);
test_pki_delete_root(Arc::clone(&c), &root_token, "pki", true);

#[cfg(feature = "crypto_adaptor_tongsuo")]
{
test_pki_mount(Arc::clone(&c), &root_token, "sm2pki");
test_pki_generate_root(Arc::clone(&c), &root_token, "sm2pki", "sm2", 256, true, true);
test_pki_config_role(Arc::clone(&c), &root_token, "sm2pki", "sm2", 256);
test_pki_issue_cert_by_generate_root(Arc::clone(&c), &root_token, "sm2pki", "sm2", 256);
test_pki_delete_root(Arc::clone(&c), &root_token, "sm2pki", true);
}
}
}
}

0 comments on commit 4bce448

Please sign in to comment.