Fix UT

Author: cYKatherine

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

@@ -7,8 +7,14 @@
 @Builder(toBuilder = true)
 public class MaaTokenPayload {
     public static final String SEV_SNP_VM_TYPE = "sevsnpvm";
+    public static final String AZURE_CC_PROTOCOL = "azure-cc";
+    public static final String AZURE_CC_AKS_PROTOCOL = "azure-cc-aks";
+    // the `x-ms-compliance-status` value for ACI CC
+    public static final String AZURE_COMPLIANT_UVM = "azure-compliant-uvm";
+    // the `x-ms-compliance-status` value for AKS CC
+    public static final String AZURE_COMPLIANT_UVM_AKS = "azure-signed-katacc-uvm";
 
-    private String azure_compliant_uvm;
+    private String azureProtocol;
     private String attestationType;
     private String complianceStatus;
     private boolean vmDebuggable;
@@ -21,6 +27,11 @@ public boolean isSevSnpVM(){
     }
 
     public boolean isUtilityVMCompliant(){
-        return azure_compliant_uvm.equalsIgnoreCase(complianceStatus);
+        if (azureProtocol == AZURE_CC_PROTOCOL) {
+            return AZURE_COMPLIANT_UVM.equalsIgnoreCase(complianceStatus);
+        } else if (azureProtocol == AZURE_CC_AKS_PROTOCOL) {
+            return AZURE_COMPLIANT_UVM_AKS.equalsIgnoreCase(complianceStatus);
+        }
+        return false;
     }
 }

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

@@ -15,12 +15,6 @@
 import static com.uid2.shared.secure.JwtUtils.tryGetField;
 
 public class MaaTokenSignatureValidator implements IMaaTokenSignatureValidator {
-    // the `x-ms-compliance-status` value for ACI CC
-    public static final String AZURE_COMPLIANT_UVM = "azure-compliant-uvm";
-
-    // the `x-ms-compliance-status` value for AKS CC
-    public static final String AZURE_COMPLIANT_UVM_AKS = "azure-signed-katacc-uvm";
-
     // set to true to facilitate local test.
     public static final boolean BYPASS_SIGNATURE_CHECK = false;
 
@@ -82,12 +76,7 @@ public MaaTokenPayload validate(String tokenString, String protocol) throws Atte
 
         var tokenPayloadBuilder = MaaTokenPayload.builder();
 
-        if (protocol == "azure-cc") {
-            tokenPayloadBuilder.azure_compliant_uvm(AZURE_COMPLIANT_UVM);
-        } else if(protocol == "azure-cc-aks") {
-            tokenPayloadBuilder.azure_compliant_uvm(AZURE_COMPLIANT_UVM_AKS);
-        }
-
+        tokenPayloadBuilder.azureProtocol(protocol);
         tokenPayloadBuilder.attestationType(tryGetField(rawPayload, "x-ms-attestation-type", String.class));
         tokenPayloadBuilder.complianceStatus(tryGetField(rawPayload, "x-ms-compliance-status", String.class));
         tokenPayloadBuilder.vmDebuggable(tryGetField(rawPayload, "x-ms-sevsnpvm-is-debuggable", Boolean.class));

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

@@ -55,7 +55,7 @@ private static byte[] encodeStringUnicodeAttestationEndpoint(String data) {
 
     @BeforeEach
     public void setup() throws AttestationException {
-        when(alwaysPassTokenValidator.validate(any())).thenReturn(VALID_TOKEN_PAYLOAD);
+        when(alwaysPassTokenValidator.validate(any(), any())).thenReturn(VALID_TOKEN_PAYLOAD);
         when(alwaysPassPolicyValidator.validate(any(), any())).thenReturn(ENCLAVE_ID);
     }
 
@@ -72,7 +72,7 @@ public void testHappyPath() throws AttestationException {
     @Test
     public void testSignatureCheckFailed_ClientError() throws AttestationException {
         var errorStr = "token signature validation failed";
-        when(alwaysFailTokenValidator.validate(any())).thenThrow(new AttestationClientException(errorStr, AttestationFailure.BAD_PAYLOAD));
+        when(alwaysFailTokenValidator.validate(any(), any())).thenThrow(new AttestationClientException(errorStr, AttestationFailure.BAD_PAYLOAD));
         var provider = new AzureCCCoreAttestationService(alwaysFailTokenValidator, alwaysPassPolicyValidator);
         provider.registerEnclave(ENCLAVE_ID);
         attest(provider, ar -> {
@@ -84,7 +84,7 @@ public void testSignatureCheckFailed_ClientError() throws AttestationException {
 
     @Test
     public void testSignatureCheckFailed_ServerError() throws AttestationException {
-        when(alwaysFailTokenValidator.validate(any())).thenThrow(new AttestationException("unknown server error"));
+        when(alwaysFailTokenValidator.validate(any(), any())).thenThrow(new AttestationException("unknown server error"));
         var provider = new AzureCCCoreAttestationService(alwaysFailTokenValidator, alwaysPassPolicyValidator);
         provider.registerEnclave(ENCLAVE_ID);
         attest(provider, ar -> {

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

@@ -37,6 +37,6 @@ public void testE2E() throws AttestationException {
         var maaToken = "<Placeholder>";
         var maaServerUrl = "https://sharedeus.eus.attest.azure.net";
         var validator = new MaaTokenSignatureValidator(maaServerUrl);
-        var token = validator.validate(maaToken);
+        var token = validator.validate(maaToken, "azure-cc");
     }
 }

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

@@ -30,7 +30,7 @@ public static MaaTokenPayload validateAndParseToken(JsonObject payload, Clock cl
         var tokenVerifier = new MaaTokenSignatureValidator(MAA_BASE_URL, keyProvider, clock);
 
         // validate token
-        return tokenVerifier.validate(token);
+        return tokenVerifier.validate(token, "azure-cc");
     }
 
     private static class MockKeyProvider implements IPublicKeyProvider {

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

@@ -97,6 +97,7 @@ private MaaTokenPayload generateBasicPayload() {
                 .vmDebuggable(false)
                 .runtimeData(generateBasicRuntimeData())
                 .ccePolicyDigest(CCE_POLICY_DIGEST)
+                .azureProtocol("azure-cc")
                 .build();
     }