Change signing API to use pre-initialized key

Bug: 161471326
Test: atest ComposHostTestCases
Change-Id: I3a78bd400efceaed6ced5f5a07575bb4530e196a
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()))
+        }
     }
 }