Change signing API to use pre-initialized key

Bug: 161471326
Test: atest ComposHostTestCases
Change-Id: I3a78bd400efceaed6ced5f5a07575bb4530e196a
diff --git a/compos/aidl/com/android/compos/ICompOsService.aidl b/compos/aidl/com/android/compos/ICompOsService.aidl
index ec4f0f6..aaba86c 100644
--- a/compos/aidl/com/android/compos/ICompOsService.aidl
+++ b/compos/aidl/com/android/compos/ICompOsService.aidl
@@ -22,6 +22,16 @@
 /** {@hide} */
 interface ICompOsService {
     /**
+     * Initializes the service with the supplied encrypted private key blob. The key cannot be
+     * changed once initialized, so once initiailzed, a repeated call will fail with
+     * EX_ILLEGAL_STATE.
+     *
+     * @param keyBlob The encrypted blob containing the private key, as returned by
+     *                generateSigningKey().
+     */
+    void initializeSigningKey(in byte[] keyBlob);
+
+    /**
      * Execute a command composed of the args, in a context that may be specified in the Metadata,
      * e.g. with file descriptors pre-opened. The service is responsible to decide what executables
      * it may run.
@@ -53,13 +63,12 @@
     boolean verifySigningKey(in byte[] keyBlob, in byte[] publicKey);
 
     /**
-     * Use the supplied encrypted private key to sign some data.
+     * Signs some data with the initialized key. The call will fail with EX_ILLEGAL_STATE if not
+     * yet initialized.
      *
-     * @param keyBlob The encrypted blob containing the private key, as returned by
-     *                generateSigningKey().
      * @param data The data to be signed. (Large data sizes may cause failure.)
      * @return the signature.
      */
     // STOPSHIP(b/193241041): We must not expose this from the PVM.
-    byte[] sign(in byte[] keyBlob, in byte[] data);
+    byte[] sign(in byte[] data);
 }
diff --git a/compos/compos_key_cmd/compos_key_cmd.cpp b/compos/compos_key_cmd/compos_key_cmd.cpp
index 04ba1d0..7ad1591 100644
--- a/compos/compos_key_cmd/compos_key_cmd.cpp
+++ b/compos/compos_key_cmd/compos_key_cmd.cpp
@@ -337,8 +337,7 @@
     return result;
 }
 
-static Result<void> signFile(ICompOsService* service, const std::vector<uint8_t>& key_blob,
-                             const std::string& file) {
+static Result<void> signFile(ICompOsService* service, const std::string& file) {
     unique_fd fd(TEMP_FAILURE_RETRY(open(file.c_str(), O_RDONLY | O_CLOEXEC)));
     if (!fd.ok()) {
         return ErrnoError() << "Failed to open";
@@ -386,7 +385,7 @@
     memcpy(to_be_signed->digest, digest->digest, digest->digest_size);
 
     std::vector<uint8_t> signature;
-    auto status = service->sign(key_blob, buffer, &signature);
+    auto status = service->sign(buffer, &signature);
     if (!status.isOk()) {
         return Error() << "Failed to sign: " << status.getDescription();
     }
@@ -420,8 +419,13 @@
         return blob.error();
     }
 
+    auto status = service->initializeSigningKey(blob.value());
+    if (!status.isOk()) {
+        return Error() << "Failed to initialize signing key: " << status.getDescription();
+    }
+
     for (auto& file : files) {
-        auto result = signFile(service.get(), blob.value(), file);
+        auto result = signFile(service.get(), file);
         if (!result.ok()) {
             return Error() << result.error() << ": " << file;
         }
@@ -519,7 +523,8 @@
                   << "  verify <blob file> <public key file> Verify that the content of the\n"
                   << "    specified private key blob and public key files are valid.\n "
                   << "  sign <blob file> <files to be signed> Generate signatures for one or\n"
-                  << "    more files using the supplied private key blob.\n"
+                  << "    more files using the supplied private key blob. Signature is stored in\n"
+                  << "    <filename>.signature\n"
                   << "  make-instance <image file> Create an empty instance image file for a VM.\n"
                   << "\n"
                   << "OPTIONS: --log <log file> (--cid <cid> | --start <image file>)\n"
diff --git a/compos/src/compsvc.rs b/compos/src/compsvc.rs
index 8fe4795..b55fb7c 100644
--- a/compos/src/compsvc.rs
+++ b/compos/src/compsvc.rs
@@ -22,6 +22,7 @@
 use log::{debug, warn};
 use std::ffi::CString;
 use std::path::PathBuf;
+use std::sync::{Arc, RwLock};
 
 use crate::compilation::{compile, CompilerOutput};
 use crate::compos_key_service::CompOsKeyService;
@@ -43,6 +44,7 @@
     let service = CompOsService {
         dex2oat_path: PathBuf::from(DEX2OAT_PATH),
         key_service: CompOsKeyService::new(rpc_binder)?,
+        key_blob: Arc::new(RwLock::new(Vec::new())),
     };
     Ok(BnCompOsService::new_binder(service, BinderFeatures::default()))
 }
@@ -50,11 +52,22 @@
 struct CompOsService {
     dex2oat_path: PathBuf,
     key_service: CompOsKeyService,
+    key_blob: Arc<RwLock<Vec<u8>>>,
 }
 
 impl Interface for CompOsService {}
 
 impl ICompOsService for CompOsService {
+    fn initializeSigningKey(&self, key_blob: &[u8]) -> BinderResult<()> {
+        let mut w = self.key_blob.write().unwrap();
+        if w.is_empty() {
+            *w = Vec::from(key_blob);
+            Ok(())
+        } else {
+            Err(new_binder_exception(ExceptionCode::ILLEGAL_STATE, "Cannot re-initialize the key"))
+        }
+    }
+
     fn execute(&self, args: &[String], metadata: &Metadata) -> BinderResult<i8> {
         let authfs_service = get_authfs_service()?;
         let output = compile(&self.dex2oat_path, args, authfs_service, metadata).map_err(|e| {
@@ -90,10 +103,15 @@
         })
     }
 
-    fn sign(&self, key_blob: &[u8], data: &[u8]) -> BinderResult<Vec<u8>> {
-        self.key_service
-            .do_sign(key_blob, data)
-            .map_err(|e| new_binder_exception(ExceptionCode::ILLEGAL_STATE, e.to_string()))
+    fn sign(&self, data: &[u8]) -> BinderResult<Vec<u8>> {
+        let key = &*self.key_blob.read().unwrap();
+        if key.is_empty() {
+            Err(new_binder_exception(ExceptionCode::ILLEGAL_STATE, "Key is not initialized"))
+        } else {
+            self.key_service
+                .do_sign(key, data)
+                .map_err(|e| new_binder_exception(ExceptionCode::ILLEGAL_STATE, e.to_string()))
+        }
     }
 }
 
diff --git a/compos/tests/java/android/compos/test/ComposKeyTestCase.java b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
index 6ef82f7..6a3f755 100644
--- a/compos/tests/java/android/compos/test/ComposKeyTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
@@ -120,6 +120,33 @@
                         TEST_ROOT + "test_key.pubkey",
                         TEST_ROOT + "test_key2.blob");
         assertThat(result.getStatus()).isEqualTo(CommandStatus.FAILED);
+
+        // Now, continue to test the signing operation. It's the best to do this in a new test
+        // method. Since we boot a VM for each test method, and booting a VM on cuttlefish/GCE is
+        // very slow, a new test method unfortunately makes the whole test module to exceed the
+        // timeout configured in the test infrastructure.
+
+        // Generate key - should succeed
+        android.run(
+                COMPOS_KEY_CMD_BIN,
+                "--cid " + mCid,
+                "generate",
+                TEST_ROOT + "test_key3.blob",
+                TEST_ROOT + "test_key3.pubkey");
+
+        // Generate some data to sign in a writable directory
+        android.run("echo something > /data/local/tmp/something.txt");
+
+        // Sign something - should succeed
+        android.run(
+                COMPOS_KEY_CMD_BIN,
+                "--cid " + mCid,
+                "sign",
+                TEST_ROOT + "test_key3.blob",
+                "/data/local/tmp/something.txt");
+
+        // Check existence of the output signature - should succeed
+        android.run("test -f /data/local/tmp/something.txt.signature");
     }
 
     private void startVm() throws Exception {