Skip to content

Commit 25ef824

Browse files
cYKatherinecYKatherine
committed
Fix UT
1 parent 27d93d5 commit 25ef824

File tree

8 files changed

+220
-21
lines changed

8 files changed

+220
-21
lines changed

‎src/main/java/com/uid2/shared/secure/azurecc/MaaTokenPayload.java‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,14 @@
77
@Builder(toBuilder = true)
88
public class MaaTokenPayload {
99
public static final String SEV_SNP_VM_TYPE = "sevsnpvm";
10+
public static final String AZURE_CC_PROTOCOL = "azure-cc";
11+
public static final String AZURE_CC_AKS_PROTOCOL = "azure-cc-aks";
12+
// the `x-ms-compliance-status` value for ACI CC
13+
public static final String AZURE_COMPLIANT_UVM = "azure-compliant-uvm";
14+
// the `x-ms-compliance-status` value for AKS CC
15+
public static final String AZURE_COMPLIANT_UVM_AKS = "azure-signed-katacc-uvm";
1016

11-
private String azure_compliant_uvm;
17+
private String azureProtocol;
1218
private String attestationType;
1319
private String complianceStatus;
1420
private boolean vmDebuggable;
@@ -21,6 +27,11 @@ public boolean isSevSnpVM(){
2127
}
2228

2329
public boolean isUtilityVMCompliant(){
24-
return azure_compliant_uvm.equalsIgnoreCase(complianceStatus);
30+
if (azureProtocol == AZURE_CC_PROTOCOL) {
31+
return AZURE_COMPLIANT_UVM.equalsIgnoreCase(complianceStatus);
32+
} else if (azureProtocol == AZURE_CC_AKS_PROTOCOL) {
33+
return AZURE_COMPLIANT_UVM_AKS.equalsIgnoreCase(complianceStatus);
34+
}
35+
return false;
2536
}
2637
}

‎src/main/java/com/uid2/shared/secure/azurecc/MaaTokenSignatureValidator.java‎

Lines changed: 1 addition & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,6 @@
1515
import static com.uid2.shared.secure.JwtUtils.tryGetField;
1616

1717
public class MaaTokenSignatureValidator implements IMaaTokenSignatureValidator {
18-
// the `x-ms-compliance-status` value for ACI CC
19-
public static final String AZURE_COMPLIANT_UVM = "azure-compliant-uvm";
20-
21-
// the `x-ms-compliance-status` value for AKS CC
22-
public static final String AZURE_COMPLIANT_UVM_AKS = "azure-signed-katacc-uvm";
23-
2418
// set to true to facilitate local test.
2519
public static final boolean BYPASS_SIGNATURE_CHECK = false;
2620

@@ -82,12 +76,7 @@ public MaaTokenPayload validate(String tokenString, String protocol) throws Atte
8276

8377
var tokenPayloadBuilder = MaaTokenPayload.builder();
8478

85-
if (protocol == "azure-cc") {
86-
tokenPayloadBuilder.azure_compliant_uvm(AZURE_COMPLIANT_UVM);
87-
} else if(protocol == "azure-cc-aks") {
88-
tokenPayloadBuilder.azure_compliant_uvm(AZURE_COMPLIANT_UVM_AKS);
89-
}
90-
79+
tokenPayloadBuilder.azureProtocol(protocol);
9180
tokenPayloadBuilder.attestationType(tryGetField(rawPayload, "x-ms-attestation-type", String.class));
9281
tokenPayloadBuilder.complianceStatus(tryGetField(rawPayload, "x-ms-compliance-status", String.class));
9382
tokenPayloadBuilder.vmDebuggable(tryGetField(rawPayload, "x-ms-sevsnpvm-is-debuggable", Boolean.class));

‎src/test/java/com/uid2/shared/secure/AzureCCAksCoreAttestationServiceTest.java‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
package com.uid2.shared.secure;
2+
3+
import com.uid2.shared.secure.azurecc.IMaaTokenSignatureValidator;
4+
import com.uid2.shared.secure.azurecc.IPolicyValidator;
5+
import com.uid2.shared.secure.azurecc.MaaTokenPayload;
6+
import com.uid2.shared.secure.azurecc.RuntimeData;
7+
import io.vertx.core.AsyncResult;
8+
import io.vertx.core.Handler;
9+
import org.junit.jupiter.api.BeforeEach;
10+
import org.junit.jupiter.api.Test;
11+
import org.junit.jupiter.api.extension.ExtendWith;
12+
import org.mockito.Mock;
13+
import org.mockito.junit.jupiter.MockitoExtension;
14+
import org.mockito.junit.jupiter.MockitoSettings;
15+
import org.mockito.quality.Strictness;
16+
17+
import java.nio.ByteBuffer;
18+
import java.nio.charset.StandardCharsets;
19+
import java.util.Arrays;
20+
import java.util.Base64;
21+
22+
import static org.junit.jupiter.api.Assertions.*;
23+
import static org.mockito.ArgumentMatchers.any;
24+
import static org.mockito.Mockito.when;
25+
26+
@ExtendWith(MockitoExtension.class)
27+
@MockitoSettings(strictness = Strictness.LENIENT)
28+
class AzureCCAksCoreAttestationServiceTest {
29+
private static byte[] encodeStringUnicodeAttestationEndpoint(String data) {
30+
// buffer.array() may include extra empty bytes at the end. This returns only the bytes that have data
31+
ByteBuffer buffer = StandardCharsets.UTF_8.encode(data);
32+
return Arrays.copyOf(buffer.array(), buffer.limit());
33+
}
34+
private static final String ATTESTATION_REQUEST = "test-attestation-request";
35+
36+
private static final String PUBLIC_KEY = "test-public-key";
37+
38+
private static final String ENCLAVE_ID = "test-enclave";
39+
private static final String ATTESTATION_URL = "https://example.com";
40+
private static final String ENCODED_ATTESTATION_URL = Base64.getEncoder().encodeToString(encodeStringUnicodeAttestationEndpoint(ATTESTATION_URL));
41+
private static final RuntimeData RUNTIME_DATA = RuntimeData.builder().attestationUrl(ENCODED_ATTESTATION_URL).build();
42+
43+
private static final MaaTokenPayload VALID_TOKEN_PAYLOAD = MaaTokenPayload.builder().runtimeData(RUNTIME_DATA).build();
44+
@Mock
45+
private IMaaTokenSignatureValidator alwaysPassTokenValidator;
46+
47+
@Mock
48+
private IMaaTokenSignatureValidator alwaysFailTokenValidator;
49+
50+
@Mock
51+
private IPolicyValidator alwaysPassPolicyValidator;
52+
53+
@Mock
54+
private IPolicyValidator alwaysFailPolicyValidator;
55+
56+
@BeforeEach
57+
public void setup() throws AttestationException {
58+
when(alwaysPassTokenValidator.validate(any(), any())).thenReturn(VALID_TOKEN_PAYLOAD);
59+
when(alwaysPassPolicyValidator.validate(any(), any())).thenReturn(ENCLAVE_ID);
60+
}
61+
62+
@Test
63+
public void testHappyPath() throws AttestationException {
64+
var provider = new AzureCCAksCoreAttestationService(alwaysPassTokenValidator, alwaysPassPolicyValidator);
65+
provider.registerEnclave(ENCLAVE_ID);
66+
attest(provider, ar -> {
67+
assertTrue(ar.succeeded());
68+
assertTrue(ar.result().isSuccess());
69+
});
70+
}
71+
72+
@Test
73+
public void testSignatureCheckFailed_ClientError() throws AttestationException {
74+
var errorStr = "token signature validation failed";
75+
when(alwaysFailTokenValidator.validate(any(), any())).thenThrow(new AttestationClientException(errorStr, AttestationFailure.BAD_PAYLOAD));
76+
var provider = new AzureCCAksCoreAttestationService(alwaysFailTokenValidator, alwaysPassPolicyValidator);
77+
provider.registerEnclave(ENCLAVE_ID);
78+
attest(provider, ar -> {
79+
assertTrue(ar.succeeded());
80+
assertFalse(ar.result().isSuccess());
81+
assertEquals(errorStr, ar.result().getReason());
82+
});
83+
}
84+
85+
@Test
86+
public void testSignatureCheckFailed_ServerError() throws AttestationException {
87+
when(alwaysFailTokenValidator.validate(any(), any())).thenThrow(new AttestationException("unknown server error"));
88+
var provider = new AzureCCAksCoreAttestationService(alwaysFailTokenValidator, alwaysPassPolicyValidator);
89+
provider.registerEnclave(ENCLAVE_ID);
90+
attest(provider, ar -> {
91+
assertFalse(ar.succeeded());
92+
assertTrue(ar.cause() instanceof AttestationException);
93+
});
94+
}
95+
96+
@Test
97+
public void testPolicyCheckSuccess_ClientError() throws AttestationException {
98+
var errorStr = "policy validation failed";
99+
when(alwaysFailPolicyValidator.validate(any(), any())).thenThrow(new AttestationClientException(errorStr, AttestationFailure.BAD_PAYLOAD));
100+
var provider = new AzureCCAksCoreAttestationService(alwaysFailTokenValidator, alwaysFailPolicyValidator);
101+
provider.registerEnclave(ENCLAVE_ID);
102+
attest(provider, ar -> {
103+
assertTrue(ar.succeeded());
104+
assertFalse(ar.result().isSuccess());
105+
assertEquals(errorStr, ar.result().getReason());
106+
});
107+
}
108+
109+
@Test
110+
public void testPolicyCheckFailed_ServerError() throws AttestationException {
111+
when(alwaysFailPolicyValidator.validate(any(), any())).thenThrow(new AttestationException("unknown server error"));
112+
var provider = new AzureCCAksCoreAttestationService(alwaysFailTokenValidator, alwaysFailPolicyValidator);
113+
provider.registerEnclave(ENCLAVE_ID);
114+
attest(provider, ar -> {
115+
assertFalse(ar.succeeded());
116+
assertTrue(ar.cause() instanceof AttestationException);
117+
});
118+
}
119+
120+
@Test
121+
public void testEnclaveNotRegistered() throws AttestationException {
122+
var provider = new AzureCCAksCoreAttestationService(alwaysFailTokenValidator, alwaysPassPolicyValidator);
123+
attest(provider, ar -> {
124+
assertTrue(ar.succeeded());
125+
assertFalse(ar.result().isSuccess());
126+
assertEquals(AttestationFailure.FORBIDDEN_ENCLAVE, ar.result().getFailure());
127+
});
128+
}
129+
130+
private static void attest(ICoreAttestationService provider, Handler<AsyncResult<AttestationResult>> handler) {
131+
provider.attest(
132+
ATTESTATION_REQUEST.getBytes(StandardCharsets.UTF_8),
133+
PUBLIC_KEY.getBytes(StandardCharsets.UTF_8),
134+
handler);
135+
}
136+
}

‎src/test/java/com/uid2/shared/secure/AzureCCCoreAttestationServiceTest.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ private static byte[] encodeStringUnicodeAttestationEndpoint(String data) {
5555

5656
@BeforeEach
5757
public void setup() throws AttestationException {
58-
when(alwaysPassTokenValidator.validate(any())).thenReturn(VALID_TOKEN_PAYLOAD);
58+
when(alwaysPassTokenValidator.validate(any(), any())).thenReturn(VALID_TOKEN_PAYLOAD);
5959
when(alwaysPassPolicyValidator.validate(any(), any())).thenReturn(ENCLAVE_ID);
6060
}
6161

@@ -72,7 +72,7 @@ public void testHappyPath() throws AttestationException {
7272
@Test
7373
public void testSignatureCheckFailed_ClientError() throws AttestationException {
7474
var errorStr = "token signature validation failed";
75-
when(alwaysFailTokenValidator.validate(any())).thenThrow(new AttestationClientException(errorStr, AttestationFailure.BAD_PAYLOAD));
75+
when(alwaysFailTokenValidator.validate(any(), any())).thenThrow(new AttestationClientException(errorStr, AttestationFailure.BAD_PAYLOAD));
7676
var provider = new AzureCCCoreAttestationService(alwaysFailTokenValidator, alwaysPassPolicyValidator);
7777
provider.registerEnclave(ENCLAVE_ID);
7878
attest(provider, ar -> {
@@ -84,7 +84,7 @@ public void testSignatureCheckFailed_ClientError() throws AttestationException {
8484

8585
@Test
8686
public void testSignatureCheckFailed_ServerError() throws AttestationException {
87-
when(alwaysFailTokenValidator.validate(any())).thenThrow(new AttestationException("unknown server error"));
87+
when(alwaysFailTokenValidator.validate(any(), any())).thenThrow(new AttestationException("unknown server error"));
8888
var provider = new AzureCCCoreAttestationService(alwaysFailTokenValidator, alwaysPassPolicyValidator);
8989
provider.registerEnclave(ENCLAVE_ID);
9090
attest(provider, ar -> {

‎src/test/java/com/uid2/shared/secure/azurecc/MaaTokenSignatureValidatorTest.java‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,28 @@ public void testPayload() throws Exception {
2222
var expectedLocation = "East US";
2323
var expectedPublicKey = "abc";
2424

25-
var tokenPayload = validateAndParseToken(payload, clock);
25+
var tokenPayload = validateAndParseToken(payload, clock, "azure-cc");
26+
assertEquals(true, tokenPayload.isSevSnpVM());
27+
assertEquals(true, tokenPayload.isUtilityVMCompliant());
28+
assertEquals(false, tokenPayload.isVmDebuggable());
29+
assertEquals(expectedCcePolicy, tokenPayload.getCcePolicyDigest());
30+
assertEquals(expectedLocation, tokenPayload.getRuntimeData().getLocation());
31+
assertEquals(expectedPublicKey, tokenPayload.getRuntimeData().getPublicKey());
32+
}
33+
34+
@Test
35+
public void testAksPayload() throws Exception {
36+
// expire at 1695313895
37+
var payloadPath = "/com.uid2.shared/test/secure/azurecc/jwt_payload_aks.json";
38+
var payload = loadFromJson(payloadPath);
39+
var clock = new TestClock();
40+
clock.setCurrentTimeMs(1695313893000L);
41+
42+
var expectedCcePolicy = "fef932e0103f6132437e8a1223f32efc4bea63342f893b5124645224ef29ba73";
43+
var expectedLocation = "East US";
44+
var expectedPublicKey = "abc";
45+
46+
var tokenPayload = validateAndParseToken(payload, clock, "azure-cc-aks");
2647
assertEquals(true, tokenPayload.isSevSnpVM());
2748
assertEquals(true, tokenPayload.isUtilityVMCompliant());
2849
assertEquals(false, tokenPayload.isVmDebuggable());
@@ -37,6 +58,6 @@ public void testE2E() throws AttestationException {
3758
var maaToken = "<Placeholder>";
3859
var maaServerUrl = "https://sharedeus.eus.attest.azure.net";
3960
var validator = new MaaTokenSignatureValidator(maaServerUrl);
40-
var token = validator.validate(maaToken);
61+
var token = validator.validate(maaToken, "azure-cc");
4162
}
4263
}

‎src/test/java/com/uid2/shared/secure/azurecc/MaaTokenUtils.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
public class MaaTokenUtils {
1515
public static final String MAA_BASE_URL = "https://sharedeus.eus.attest.azure.net";
1616

17-
public static MaaTokenPayload validateAndParseToken(JsonObject payload, Clock clock) throws Exception{
17+
public static MaaTokenPayload validateAndParseToken(JsonObject payload, Clock clock, String protocol) throws Exception{
1818
var gen = KeyPairGenerator.getInstance(Const.Name.AsymetricEncryptionKeyClass);
1919
gen.initialize(2048, new SecureRandom());
2020
var keyPair = gen.generateKeyPair();
@@ -30,7 +30,7 @@ public static MaaTokenPayload validateAndParseToken(JsonObject payload, Clock cl
3030
var tokenVerifier = new MaaTokenSignatureValidator(MAA_BASE_URL, keyProvider, clock);
3131

3232
// validate token
33-
return tokenVerifier.validate(token);
33+
return tokenVerifier.validate(token, protocol);
3434
}
3535

3636
private static class MockKeyProvider implements IPublicKeyProvider {

‎src/test/java/com/uid2/shared/secure/azurecc/PolicyValidatorTest.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ private MaaTokenPayload generateBasicPayload() {
9797
.vmDebuggable(false)
9898
.runtimeData(generateBasicRuntimeData())
9999
.ccePolicyDigest(CCE_POLICY_DIGEST)
100+
.azureProtocol("azure-cc")
100101
.build();
101102
}
102103

‎src/test/resources/com.uid2.shared/test/secure/azurecc/jwt_payload_aks.json‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
{
2+
"exp": 1695313895,
3+
"iat": 1695285095,
4+
"iss": "https://sharedeus.eus.attest.azure.net",
5+
"jti": "3b16f2ab4492417aae4cc9a5e6506ca2519659c0d8fdc2bf442fe01aa9b8e46c",
6+
"nbf": 1695285095,
7+
"nonce": "7394904505194784658",
8+
"x-ms-attestation-type": "sevsnpvm",
9+
"x-ms-compliance-status": "azure-signed-katacc-uvm",
10+
"x-ms-policy-hash": "9NY0VnTQ-IiBriBplVUpFbczcDaEBUwsiFYAzHu_gco",
11+
"x-ms-runtime": {
12+
"location": "East US",
13+
"publicKey": "abc"
14+
},
15+
"x-ms-sevsnpvm-authorkeydigest": "000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
16+
"x-ms-sevsnpvm-bootloader-svn": 3,
17+
"x-ms-sevsnpvm-familyId": "01000000000000000000000000000000",
18+
"x-ms-sevsnpvm-guestsvn": 2,
19+
"x-ms-sevsnpvm-hostdata": "fef932e0103f6132437e8a1223f32efc4bea63342f893b5124645224ef29ba73",
20+
"x-ms-sevsnpvm-idkeydigest": "ebeeeabce075eeaba3d9ea24d8495137a2877c0d20ac6ea73fc6d2f8aeb50de132150e0a0752664919bcebbf2e8c5807",
21+
"x-ms-sevsnpvm-imageId": "02000000000000000000000000000000",
22+
"x-ms-sevsnpvm-is-debuggable": false,
23+
"x-ms-sevsnpvm-launchmeasurement": "03fea02823189b25d0623a5c81f97c8ba4d2fbc48c914a55ce525f90454ddcec303743dac2fc013f0846912d1412f6df",
24+
"x-ms-sevsnpvm-microcode-svn": 115,
25+
"x-ms-sevsnpvm-migration-allowed": false,
26+
"x-ms-sevsnpvm-reportdata": "4e7d4a413745ddea79f05d20d9ac7add3659ac783ef24684127bbbb3e50fc63c0000000000000000000000000000000000000000000000000000000000000000",
27+
"x-ms-sevsnpvm-reportid": "d137a83c2d42d81dd42d39ad95ef9023de63216ddaaf2c368a8c41a636ddb2a9",
28+
"x-ms-sevsnpvm-smt-allowed": true,
29+
"x-ms-sevsnpvm-snpfw-svn": 8,
30+
"x-ms-sevsnpvm-tee-svn": 0,
31+
"x-ms-sevsnpvm-uvm-endorsement": {
32+
"x-ms-sevsnpvm-guestsvn": "100",
33+
"x-ms-sevsnpvm-launchmeasurement": "03fea02823189b25d0623a5c81f97c8ba4d2fbc48c914a55ce525f90454ddcec303743dac2fc013f0846912d1412f6df"
34+
},
35+
"x-ms-sevsnpvm-uvm-endorsement-headers": {
36+
"feed": "ContainerPlat-AMD-UVM",
37+
"iss": "did:x509:0:sha256:I__iuL25oXEVFdTP_aBLx_eT1RPHbCQ_ECBQfYZpt9s::eku:1.3.6.1.4.1.311.76.59.1.2"
38+
},
39+
"x-ms-sevsnpvm-vmpl": 0,
40+
"x-ms-ver": "1.0"
41+
}

0 commit comments

Comments
 (0)