[PPML] Implement bi-attestation in AttestationCLI.scala (intel-analyt…

…ics#5648) * Implement bi-attestation in AttestationCLI.scala * Refine * Refine * Refine * Refine * Fix problems according to comments Co-authored-by: xiangyuT <[email protected]>
ForJadeForest · Sep 20, 2022 · adc8bbd · adc8bbd
1 parent fbbb383
commit adc8bbd
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/scala/ppml/src/main/scala/com/intel/analytics/bigdl/ppml/attestation/AttestationCLI.scala b/scala/ppml/src/main/scala/com/intel/analytics/bigdl/ppml/attestation/AttestationCLI.scala
@@ -32,7 +32,8 @@ object AttestationCLI {
         case class CmdParams(appID: String = "test",
                              appKey: String = "test",
                              asType: String = ATTESTATION_CONVENTION.MODE_EHSM_KMS,
-                             asURL: String = "127.0.0.1",
+                             asURL: String = "127.0.0.1:9000",
+                             challenge: String = "",
                              userReport: String = "ppml")
 
         val cmdParser = new OptionParser[CmdParams]("PPML Attestation Quote Generation Cmd tool") {
@@ -43,11 +44,14 @@ object AttestationCLI {
               .text("app key for this app")
               .action((x, c) => c.copy(appKey = x))
             opt[String]('u', "asURL")
-              .text("attestation service url, default is 127.0.0.1")
+              .text("attestation service url, default is 127.0.0.1:9000")
               .action((x, c) => c.copy(asURL = x))
             opt[String]('t', "asType")
               .text("attestation service type, default is EHSMKeyManagementService")
-              .action((x, c) => c.copy(asURL = x))
+              .action((x, c) => c.copy(asType = x))
+            opt[String]('c', "challenge")
+              .text("challenge to attestation service, default is '' which skip bi-attestation")
+              .action((x, c) => c.copy(challenge = x))
             opt[String]('p', "userReport")
               .text("userReportDataPath, default is test")
               .action((x, c) => c.copy(userReport = x))
@@ -69,6 +73,20 @@ object AttestationCLI {
                 new DummyAttestationService()
             case _ => throw new AttestationRuntimeException("Wrong Attestation service type")
         }
+
+        val challengeString = params.challenge
+        if (challengeString.length() > 0) {
+            val asQuote = as.getQuoteFromServer(challengeString)
+            // System.out.print(asQuote)
+            val quoteVerifier = new SGXDCAPQuoteVerifierImpl()
+            val verifyQuoteResult = quoteVerifier.verifyQuote(asQuote.getBytes())
+            if (verifyQuoteResult == 0) {
+              System.out.println("Quote Verification Success!")
+            } else {
+              System.out.println("Quote Verification Fail! Application killed")
+              System.exit(1)
+            }
+        }
         val attResult = as.attestWithServer(Base64.getEncoder.encodeToString(quote))
         // System.out.print(as.attestWithServer(quote))
         if (attResult._1) {

diff --git a/...ml/src/main/scala/com/intel/analytics/bigdl/ppml/attestation/EHSMAttestationService.scala b/...ml/src/main/scala/com/intel/analytics/bigdl/ppml/attestation/EHSMAttestationService.scala
@@ -20,6 +20,7 @@ package com.intel.analytics.bigdl.ppml.attestation
 import com.intel.analytics.bigdl.dllib.utils.Log4Error
 import com.intel.analytics.bigdl.ppml.utils.EHSMParams
 import com.intel.analytics.bigdl.ppml.utils.HTTPUtil.postRequest
+import java.util.Base64
 import org.apache.logging.log4j.LogManager
 import org.json.JSONObject
 
@@ -68,7 +69,8 @@ class EHSMAttestationService(kmsServerIP: String, kmsServerPort: String,
     if (challenge != postResult.getString(RES_CHALLENGE)) {
       Log4Error.invalidOperationError(false, "Challenge not matched")
     }
-    postResult.getString(RES_QUOTE)
+    val quote = Base64.getDecoder().decode(postResult.getString(RES_QUOTE))
+    new String(quote)
   }
 
   override def attestWithServer(quote: String): (Boolean, String) = {