Implement GET_AUTH_TOKEN_KEY

Implement authentication token requests, and ensure we only respond
to requests from the Gatekeeper TA.

Test: make ta + qemu
Change-Id: Ice37bbdac06e832faaf3dd8ca7947d8ea5f66cd5
diff --git a/src/client_id.rs b/src/client_id.rs
new file mode 100644
index 0000000..60bd8da
--- /dev/null
+++ b/src/client_id.rs
@@ -0,0 +1,50 @@
+//
+// Copyright (C) 2024 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.
+
+use optee_utee::Uuid;
+use optee_utee_sys::{
+    TEE_GetPropertyAsIdentity, TEE_Identity, TEE_LOGIN_PUBLIC, TEE_LOGIN_TRUSTED_APP,
+    TEE_PROPSET_CURRENT_CLIENT, TEE_SUCCESS, TEE_UUID,
+};
+
+// Identity of the current-session TA's client.
+pub enum ClientId {
+    TA(Uuid),
+    Other,
+}
+
+impl ClientId {
+    pub fn detect() -> Self {
+        let prop_name = c"gpd.client.identity";
+        let mut detected = TEE_Identity {
+            login: TEE_LOGIN_PUBLIC,
+            uuid: TEE_UUID { timeLow: 0, timeMid: 0, timeHiAndVersion: 0, clockSeqAndNode: [0; 8] },
+        };
+        // SAFETY: TEE_GetPropertyAsIdentity expects *const c_char for the name and
+        // *mut TEE_Identity for the value, and both pointers are valid.
+        match unsafe {
+            TEE_GetPropertyAsIdentity(
+                TEE_PROPSET_CURRENT_CLIENT,
+                prop_name.as_ptr() as _,
+                &mut detected,
+            )
+        } {
+            TEE_SUCCESS if detected.login == TEE_LOGIN_TRUSTED_APP => {
+                ClientId::TA(Uuid::from_raw(detected.uuid))
+            }
+            _ => ClientId::Other,
+        }
+    }
+}
diff --git a/src/invoke.rs b/src/invoke.rs
index fe4de78..b484615 100644
--- a/src/invoke.rs
+++ b/src/invoke.rs
@@ -28,6 +28,8 @@
 pub const PTA_SYSTEM_DERIVE_TA_UNIQUE_KEY: u32 = 1;
 pub const PTA_SYSTEM_UUID: &'static str = "3a2f8978-5dc0-11e8-9c2d-fa7ae01bbebc";
 
+pub const TA_GATEKEEPER_UUID: &'static str = "4d573443-6a56-4272-ac6f-2425af9ef9bb";
+
 #[allow(dead_code)]
 pub enum Direction {
     In,
diff --git a/src/main.rs b/src/main.rs
index 31198c1..ad41e56 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -20,6 +20,7 @@
 #[macro_use]
 extern crate alloc;
 
+mod client_id;
 mod command;
 mod config;
 mod crypto;
@@ -32,9 +33,13 @@
 use alloc::boxed::Box;
 use alloc::collections::vec_deque::VecDeque;
 use alloc::vec::Vec;
+use client_id::ClientId;
 use core::ffi::c_void;
+use kmr_ta::device::DeviceHmac;
 use log::info;
-use optee_utee::{ta_close_session, ta_create, ta_destroy, ta_invoke_command, ta_open_session};
+use optee_utee::{
+    ta_close_session, ta_create, ta_destroy, ta_invoke_command, ta_open_session, Uuid,
+};
 use optee_utee::{Error, ErrorKind, ParamType, Parameters, Result};
 use spin::Mutex;
 use spin::Once;
@@ -166,7 +171,7 @@
     params: &mut Parameters,
 ) -> Result<()> {
     match cmd_id {
-        command::GET_AUTH_TOKEN_KEY => Err(Error::new(ErrorKind::NotSupported)),
+        command::GET_AUTH_TOKEN_KEY => handle_authtoken_key_request(sess_ctx, params),
         command::WRITE => handle_ipc_request(sess_ctx, params),
         command::READ => handle_ipc_response(sess_ctx, params),
         #[cfg(feature = "dev")]
@@ -190,9 +195,14 @@
         return Err(Error::new(ErrorKind::Busy));
     }
 
+    if params.0.raw().is_null() {
+        return Err(Error::new(ErrorKind::BadParameters));
+    }
+    // SAFETY: as_memref() dereferences the raw TEE_Param pointer, which is non-null.
+    let mut p = unsafe { params.0.as_memref()? };
+
     // Check MemrefInput here because as_memref() does not differentiate
     // between input and output.
-    let mut p = unsafe { params.0.as_memref()? };
     if let ParamType::MemrefInput = p.param_type() {
         let req = p.buffer();
         // Deny requests that are too big.
@@ -216,9 +226,14 @@
         return Err(Error::new(ErrorKind::NoData));
     }
 
+    if params.0.raw().is_null() {
+        return Err(Error::new(ErrorKind::BadParameters));
+    }
+    // SAFETY: as_memref() dereferences the raw TEE_Param pointer, which is non-null.
+    let mut p = unsafe { params.0.as_memref()? };
+
     // Check MemrefOutput here because as_memref() does not differentiate
     // between input and output.
-    let mut p = unsafe { params.0.as_memref()? };
     if let ParamType::MemrefOutput = p.param_type() {
         // Take all packets out. Unwrap() never panics.
         let mut packets = keymint_ctx.ipc_resp.take().unwrap();
@@ -238,3 +253,46 @@
         Err(Error::new(ErrorKind::BadParameters))
     }
 }
+
+fn client_is_gatekeeper_ta() -> bool {
+    let gatekeeper_uuid = Uuid::parse_str(invoke::TA_GATEKEEPER_UUID).unwrap();
+
+    match ClientId::detect() {
+        ClientId::TA(uuid) => uuid == gatekeeper_uuid,
+        _ => false,
+    }
+}
+
+fn handle_authtoken_key_request(
+    _keymint_ctx: &mut SessionContext,
+    params: &mut Parameters,
+) -> Result<()> {
+    if !client_is_gatekeeper_ta() {
+        return Err(Error::new(ErrorKind::AccessDenied));
+    }
+
+    if params.1.raw().is_null() {
+        return Err(Error::new(ErrorKind::BadParameters));
+    }
+    // SAFETY: as_memref() dereferences the raw TEE_Param pointer, which is non-null.
+    let mut p = unsafe { params.1.as_memref()? };
+
+    // Check MemrefOutput here because as_memref() does not differentiate
+    // between input and output.
+    if let ParamType::MemrefOutput = p.param_type() {
+        let hmac_key = match software::keys::AuthTokenHmac::try_new() {
+            Ok(hmac) => hmac.get_hmac_key().unwrap(),
+            _ => return Err(Error::new(ErrorKind::Generic)),
+        };
+
+        if hmac_key.0.len() > p.buffer().len() {
+            return Err(Error::new(ErrorKind::ShortBuffer));
+        }
+
+        p.set_updated_size(hmac_key.0.len());
+        p.buffer().copy_from_slice(hmac_key.0.as_slice());
+        Ok(())
+    } else {
+        Err(Error::new(ErrorKind::BadParameters))
+    }
+}