Merge "Store apk root hash to the instance disk"
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 47271a7..52c00b8 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -52,7 +52,6 @@
         "microdroid_build_prop",
         "microdroid_init_rc",
         "microdroid_launcher",
-        "microdroid_manager",
 
         "ueventd.rc",
         "libbinder",
@@ -112,6 +111,7 @@
                 "apkdmverity",
                 "authfs",
                 "authfs_service",
+                "microdroid_manager",
                 "zipfuse",
 
                 // TODO(b/184872979): Needed by authfs. Remove once the Rust API is created.
diff --git a/microdroid_manager/Android.bp b/microdroid_manager/Android.bp
index 5fae7b1..0b1bc33 100644
--- a/microdroid_manager/Android.bp
+++ b/microdroid_manager/Android.bp
@@ -15,6 +15,8 @@
         "libapkverify",
         "libbinder_rpc_unstable_bindgen",
         "libbinder_rs",
+        "libbyteorder",
+        "libidsig",
         "libkernlog",
         "liblibc",
         "liblog_rust",
@@ -22,15 +24,23 @@
         "libmicrodroid_payload_config",
         "libnix",
         "libprotobuf",
+        "libring",
         "librustutils",
         "libserde",
         "libserde_json",
+        "libuuid",
         "libvsock",
+        "librand",
     ],
     shared_libs: [
         "libbinder_rpc_unstable",
     ],
     init_rc: ["microdroid_manager.rc"],
+    multilib: {
+        lib32: {
+            enabled: false,
+        },
+    },
 }
 
 rust_binary {
diff --git a/microdroid_manager/src/instance.rs b/microdroid_manager/src/instance.rs
new file mode 100644
index 0000000..3ef3b63
--- /dev/null
+++ b/microdroid_manager/src/instance.rs
@@ -0,0 +1,309 @@
+// Copyright 2021, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Provides routines to read/write on the instance disk.
+//!
+//! Instance disk is a disk where the identity of a VM instance is recorded. The identity usually
+//! includes certificates of the VM payload that is trusted, but not limited to it. Instance disk
+//! is empty when a VM is first booted. The identity data is filled in during the first boot, and
+//! then encrypted and signed. Subsequent boots decrypts and authenticates the data and uses the
+//! identity data to further verify the payload (e.g. against the certificate).
+//!
+//! Instance disk consists of a disk header and one or more partitions each of which consists of a
+//! header and payload. Each header (both the disk header and a partition header) is 512 bytes
+//! long. Payload is just next to the header and its size can be arbitrary. Headers are located at
+//! 512 bytes boundaries. So, when the size of a payload is not multiple of 512, there exists a gap
+//! between the end of the payload and the start of the next partition (if there is any).
+//!
+//! Each partition is identified by a UUID. A partition is created for a program loader that
+//! participates in the boot chain of the VM. Each program loader is expected to locate the
+//! partition that corresponds to the loader using the UUID that is assigned to the loader.
+//!
+//! The payload of a partition is encrypted/signed by a key that is unique to the loader and to the
+//! VM as well. Failing to decrypt/authenticate a partition by a loader stops the boot process.
+
+use anyhow::{anyhow, bail, Context, Result};
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use ring::aead::{Aad, Algorithm, LessSafeKey, Nonce, UnboundKey, AES_256_GCM};
+use ring::hkdf::{Salt, HKDF_SHA256};
+use std::fs::{File, OpenOptions};
+use std::io::{Read, Seek, SeekFrom, Write};
+use uuid::Uuid;
+
+/// Path to the instance disk inside the VM
+const INSTANCE_IMAGE_PATH: &str = "/dev/block/by-name/vm-instance";
+
+/// Magic string in the instance disk header
+const DISK_HEADER_MAGIC: &str = "Android-VM-instance";
+
+/// Version of the instance disk format
+const DISK_HEADER_VERSION: u16 = 1;
+
+/// Size of the headers in the instance disk
+const DISK_HEADER_SIZE: u64 = 512;
+const PARTITION_HEADER_SIZE: u64 = 512;
+
+/// UUID of the partition that microdroid manager uses
+const MICRODROID_PARTITION_UUID: &str = "cf9afe9a-0662-11ec-a329-c32663a09d75";
+
+/// Encryption algorithm used to cipher payload
+static ENCRYPT_ALG: &Algorithm = &AES_256_GCM;
+
+/// Handle to the instance disk
+pub struct InstanceDisk {
+    file: File,
+}
+
+/// Information from a partition header
+struct PartitionHeader {
+    uuid: Uuid,
+    payload_size: u64, // in bytes
+}
+
+/// Offset of a partition in the instance disk
+type PartitionOffset = u64;
+
+impl InstanceDisk {
+    /// Creates handle to instance disk
+    pub fn new() -> Result<Self> {
+        let mut file = OpenOptions::new()
+            .read(true)
+            .write(true)
+            .open(INSTANCE_IMAGE_PATH)
+            .with_context(|| format!("Failed to open {}", INSTANCE_IMAGE_PATH))?;
+
+        // Check if this file is a valid instance disk by examining the header (the first block)
+        let mut magic = [0; DISK_HEADER_MAGIC.len()];
+        file.read_exact(&mut magic)?;
+        if magic != DISK_HEADER_MAGIC.as_bytes() {
+            bail!("invalid magic: {:?}", magic);
+        }
+
+        let version = file.read_u16::<LittleEndian>()?;
+        if version == 0 {
+            bail!("invalid version: {}", version);
+        }
+        if version > DISK_HEADER_VERSION {
+            bail!("unsupported version: {}", version);
+        }
+
+        Ok(Self { file })
+    }
+
+    /// Reads the identity data that was written by microdroid manager. The returned data is
+    /// plaintext, although it is stored encrypted. In case when the partition for microdroid
+    /// manager doesn't exist, which can happen if it's the first boot, `Ok(None)` is returned.
+    pub fn read_microdroid_data(&mut self) -> Result<Option<Box<[u8]>>> {
+        let (header, offset) = self.locate_microdroid_header()?;
+        if header.is_none() {
+            return Ok(None);
+        }
+        let header = header.unwrap();
+        let payload_offset = offset + PARTITION_HEADER_SIZE;
+        self.file.seek(SeekFrom::Start(payload_offset))?;
+
+        // Read the 12-bytes nonce (unencrypted)
+        let mut nonce = [0; 12];
+        self.file.read_exact(&mut nonce)?;
+        let nonce = Nonce::assume_unique_for_key(nonce);
+
+        // Read the encrypted payload
+        let payload_size = header.payload_size - 12; // we already have read the nonce
+        let mut data = vec![0; payload_size as usize];
+        self.file.read_exact(&mut data)?;
+
+        // Read the header as well because it's part of the signed data (though not encrypted).
+        let mut header = [0; PARTITION_HEADER_SIZE as usize];
+        self.file.seek(SeekFrom::Start(offset))?;
+        self.file.read_exact(&mut header)?;
+
+        // Decrypt and authenticate the data (along with the header). The data is decrypted in
+        // place. `open_in_place` returns slice to the decrypted part in the buffer.
+        let plaintext_len = get_key().open_in_place(nonce, Aad::from(&header), &mut data)?.len();
+        // Truncate to remove the tag
+        data.truncate(plaintext_len);
+
+        Ok(Some(data.into_boxed_slice()))
+    }
+
+    /// Writes identity data to the partition for microdroid manager. The partition is appended
+    /// if it doesn't exist. The data is stored encrypted.
+    pub fn write_microdroid_data(&mut self, data: &[u8]) -> Result<()> {
+        let (header, offset) = self.locate_microdroid_header()?;
+
+        // By encrypting and signing the data, tag will be appended. The tag also becomes part of
+        // the encrypted payload which will be written. In addition, a 12-bytes nonce will be
+        // prepended (non-encrypted).
+        let payload_size = (data.len() + ENCRYPT_ALG.tag_len() + 12) as u64;
+
+        // If the partition exists, make sure we don't change the partition size. If not (i.e.
+        // partition is not found), write the header at the empty place.
+        if let Some(header) = header {
+            if header.payload_size != payload_size {
+                bail!("Can't change payload size from {} to {}", header.payload_size, payload_size);
+            }
+        } else {
+            let uuid = Uuid::parse_str(MICRODROID_PARTITION_UUID)?;
+            self.write_header_at(offset, &uuid, payload_size)?;
+        }
+
+        // Read the header as it is used as additionally authenticated data (AAD).
+        let mut header = [0; PARTITION_HEADER_SIZE as usize];
+        self.file.seek(SeekFrom::Start(offset))?;
+        self.file.read_exact(&mut header)?;
+
+        // Generate a nonce randomly and recorde it on the disk first.
+        let nonce = Nonce::assume_unique_for_key(rand::random::<[u8; 12]>());
+        self.file.seek(SeekFrom::Start(offset + PARTITION_HEADER_SIZE))?;
+        self.file.write_all(nonce.as_ref())?;
+
+        // Then encrypt and sign the data. The non-encrypted input data is copied to a vector
+        // because it is encrypted in place, and also the tag is appended.
+        let mut data = data.to_vec();
+        get_key().seal_in_place_append_tag(nonce, Aad::from(&header), &mut data)?;
+
+        // Persist the encrypted payload data
+        self.file.write_all(&data)?;
+        self.file.flush()?;
+
+        Ok(())
+    }
+
+    /// Read header at `header_offset` and parse it into a `PartitionHeader`.
+    fn read_header_at(&mut self, header_offset: u64) -> Result<PartitionHeader> {
+        assert!(
+            header_offset % PARTITION_HEADER_SIZE == 0,
+            "header offset {} is not aligned to 512 bytes",
+            header_offset
+        );
+
+        let mut uuid = [0; 16];
+        self.file.seek(SeekFrom::Start(header_offset))?;
+        self.file.read_exact(&mut uuid)?;
+        let uuid = Uuid::from_bytes(uuid);
+        let payload_size = self.file.read_u64::<LittleEndian>()?;
+
+        Ok(PartitionHeader { uuid, payload_size })
+    }
+
+    /// Write header at `header_offset`
+    fn write_header_at(
+        &mut self,
+        header_offset: u64,
+        uuid: &Uuid,
+        payload_size: u64,
+    ) -> Result<()> {
+        self.file.seek(SeekFrom::Start(header_offset))?;
+        self.file.write_all(uuid.as_bytes())?;
+        self.file.write_u64::<LittleEndian>(payload_size)?;
+        Ok(())
+    }
+
+    /// Locate the header of the partition for microdroid manager. A pair of `PartitionHeader` and
+    /// the offset of the partition in the disk is returned. If the partition is not found,
+    /// `PartitionHeader` is `None` and the offset points to the empty partition that can be used
+    /// for the partition.
+    fn locate_microdroid_header(&mut self) -> Result<(Option<PartitionHeader>, PartitionOffset)> {
+        let microdroid_uuid = Uuid::parse_str(MICRODROID_PARTITION_UUID)?;
+
+        // the first partition header is located just after the disk header
+        let mut header_offset = DISK_HEADER_SIZE;
+        loop {
+            let header = self.read_header_at(header_offset)?;
+            if header.uuid == microdroid_uuid {
+                // found a matching header
+                return Ok((Some(header), header_offset));
+            } else if header.uuid == Uuid::nil() {
+                // found an empty space
+                return Ok((None, header_offset));
+            }
+            // Move to the next partition. Be careful about overflow.
+            let payload_size = round_to_multiple(header.payload_size, PARTITION_HEADER_SIZE)?;
+            let part_size = payload_size
+                .checked_add(PARTITION_HEADER_SIZE)
+                .ok_or_else(|| anyhow!("partition too large"))?;
+            header_offset = header_offset
+                .checked_add(part_size)
+                .ok_or_else(|| anyhow!("next partition at invalid offset"))?;
+        }
+    }
+}
+
+/// Round `n` up to the nearest multiple of `unit`
+fn round_to_multiple(n: u64, unit: u64) -> Result<u64> {
+    assert!((unit & (unit - 1)) == 0, "{} is not power of two", unit);
+    let ret = (n + unit - 1) & !(unit - 1);
+    if ret < n {
+        bail!("overflow")
+    }
+    Ok(ret)
+}
+
+struct ZeroOnDropKey(LessSafeKey);
+
+impl Drop for ZeroOnDropKey {
+    fn drop(&mut self) {
+        // Zeroize the key by overwriting it with a key constructed from zeros of same length
+        // This works because the raw key bytes are allocated inside the struct, not on the heap
+        let zero = [0; 32];
+        let zero_key = LessSafeKey::new(UnboundKey::new(ENCRYPT_ALG, &zero).unwrap());
+        unsafe {
+            ::std::ptr::write_volatile::<LessSafeKey>(&mut self.0, zero_key);
+        }
+    }
+}
+
+impl std::ops::Deref for ZeroOnDropKey {
+    type Target = LessSafeKey;
+    fn deref(&self) -> &LessSafeKey {
+        &self.0
+    }
+}
+
+/// Returns the key that is used to encrypt the microdroid manager partition. It is derived from
+/// the sealing CDI of the previous stage, which is Android Boot Loader (ABL).
+fn get_key() -> ZeroOnDropKey {
+    // Sealing CDI from the previous stage. For now, this is hardcoded.
+    // TODO(jiyong): actually read this from the previous stage
+    const SEALING_CDI: [u8; 32] = [10; 32];
+
+    // Derive a key from the Sealing CDI
+    // Step 1 is extraction: https://datatracker.ietf.org/doc/html/rfc5869#section-2.2 where a
+    // pseduo random key (PRK) is extracted from (Input Keying Material - IKM, which is secret) and
+    // optional salt.
+    let salt = Salt::new(HKDF_SHA256, &[]); // use 0 as salt
+    let prk = salt.extract(&SEALING_CDI); // Sealing CDI as IKM
+
+    // Step 2 is expansion: https://datatracker.ietf.org/doc/html/rfc5869#section-2.3 where the PRK
+    // (optionally with the `info` which gives contextual information) is expanded into the output
+    // keying material (OKM). Note that the process fails only when the size of OKM is longer than
+    // 255 * SHA256_HASH_SIZE (32), which isn't the case here.
+    let info = [b"microdroid_manager_key".as_ref()];
+    let okm = prk.expand(&info, HKDF_SHA256).unwrap(); // doesn't fail as explained above
+    let mut key = [0; 32];
+    okm.fill(&mut key).unwrap(); // doesn't fail as explained above
+
+    // The term LessSafe might be misleading here. LessSafe here just means that the API can
+    // possibly accept same nonces for different messages. However, since we encrypt/decrypt only a
+    // single message (the microdroid_manager partition payload) with a randomly generated nonce,
+    // this is safe enough.
+    let ret = ZeroOnDropKey(LessSafeKey::new(UnboundKey::new(ENCRYPT_ALG, &key).unwrap()));
+
+    // Don't forget to zeroize the raw key array as well
+    unsafe {
+        ::std::ptr::write_volatile::<[u8; 32]>(&mut key, [0; 32]);
+    }
+
+    ret
+}
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index cfd8ded..278a14d 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -14,14 +14,17 @@
 
 //! Microdroid Manager
 
+mod instance;
 mod ioutil;
 mod metadata;
 
+use crate::instance::InstanceDisk;
 use anyhow::{anyhow, bail, Context, Result};
 use apkverify::verify;
 use binder::unstable_api::{new_spibinder, AIBinder};
 use binder::{FromIBinder, Strong};
-use log::{error, info, warn};
+use idsig::V4Signature;
+use log::{debug, error, info, warn};
 use microdroid_payload_config::{Task, TaskType, VmPayloadConfig};
 use nix::ioctl_read_bad;
 use rustutils::system_properties::PropertyWatcher;
@@ -94,6 +97,18 @@
         return Err(err);
     }
 
+    let mut instance = InstanceDisk::new()?;
+    // TODO(jiyong): the data should have an internal structure
+    if let Some(data) = instance.read_microdroid_data().context("Failed to read identity data")? {
+        debug!("read apk root hash: {}", to_hex_string(&data));
+        //TODO(jiyong) apkdmverity should use this root hash instead of the one read from the idsig
+        //file, if the root hash is found in the instance image.
+    } else {
+        let data = get_apk_roothash()?;
+        debug!("write apk root hash: {}", to_hex_string(&data));
+        instance.write_microdroid_data(data.as_ref()).context("Failed to write identity data")?;
+    }
+
     let service = get_vms_rpc_binder().expect("cannot connect to VirtualMachineService");
 
     if !metadata.payload_config_path.is_empty() {
@@ -129,6 +144,12 @@
     Ok(())
 }
 
+fn get_apk_roothash() -> Result<Box<[u8]>> {
+    let mut idsig = File::open("/dev/block/by-name/microdroid-apk-idsig")?;
+    let idsig = V4Signature::from(&mut idsig)?;
+    Ok(idsig.hashing_info.raw_root_hash)
+}
+
 fn load_config(path: &Path) -> Result<VmPayloadConfig> {
     info!("loading config from {:?}...", path);
     let file = ioutil::wait_for_file(path, WAIT_TIMEOUT)?;
@@ -214,3 +235,7 @@
 
     Ok(path)
 }
+
+fn to_hex_string(buf: &[u8]) -> String {
+    buf.iter().map(|b| format!("{:02X}", b)).collect()
+}