lib/src/server.rs

/*
 * Copyright (C) 2025 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.
 */

//! AuthMgr BE server
use crate::authorization_service::{
    AuthMgrAuthorizationRPCService, AuthMgrGlobalState, RawConnectionInformation,
};
use authmgr_common::{
    signed_connection_request::TEMP_AUTHMGR_FE_TRANSPORT_ID, CMD_RAW, CMD_RPC, TOKEN_LENGTH,
};
use log::{debug, error};
use rpcbinder::{RpcServer, RpcServerConnection};
use std::ffi::CStr;
use std::sync::Mutex;
use std::sync::{Arc, Weak};
use tipc::raw::{EventLoop, HandleSetWrapper};
use tipc::{
    ConnectResult, Deserialize, Handle, MessageResult, PortCfg, Serialize, Serializer, TipcError,
    UnbufferedService, Uuid,
};
use trusty_std::alloc::TryAllocFrom;

/// Port for the AuthMgr main service
pub(crate) const AUTHMGR_SERVICE_PORT: &CStr = c"com.android.trusty.rust.authmgr.V1";

/// Maximum message size.
/// TODO: determine the size
const MAX_MSG_SIZE: usize = 4000;

/// Represents the main AuthMgr service encapsulating the global state which is also shared with the
/// `AuthMgrAuthorizationRPCService`.
pub struct AuthMgrService {
    global_state: Arc<Mutex<AuthMgrGlobalState>>,
    rpc_service: Arc<RpcServer>,
    handle_set: Weak<HandleSetWrapper<AuthMgrServices>>,
}

impl AuthMgrService {
    fn new(
        global_state: Arc<Mutex<AuthMgrGlobalState>>,
        rpc_service: Arc<RpcServer>,
        handle_set: Weak<HandleSetWrapper<AuthMgrServices>>,
    ) -> Self {
        AuthMgrService { global_state, rpc_service, handle_set }
    }
}

// SAFETY: Out of the three fields of `AuthMgrService`:
// `AuthMgrGlobalState` and `HandleSetWrapper` are heap allocated in the main_loop() of this TA and
// therefore, are not bound to any particular thread.
// `RpcServer` implements (unsafe) `Send`.
unsafe impl Send for AuthMgrService {}

pub struct AuthMgrConnection {
    uuid: Uuid,
    // In on_connect, we do not know whether this is a connection request for the RPC binder service
    // or a raw connection request. Until the first message arrives, we mark the connection as
    // pending_to_be_routed = true
    pending_to_be_routed: Arc<Mutex<bool>>,
}

pub struct AuthMgrMessage(pub Vec<u8>);

impl Deserialize for AuthMgrMessage {
    type Error = TipcError;
    const MAX_SERIALIZED_SIZE: usize = MAX_MSG_SIZE;

    fn deserialize(bytes: &[u8], _handles: &mut [Option<Handle>]) -> Result<Self, TipcError> {
        Ok(AuthMgrMessage(Vec::try_alloc_from(bytes)?))
    }
}

impl<'s> Serialize<'s> for AuthMgrMessage {
    fn serialize<'a: 's, S: Serializer<'s>>(
        &'a self,
        serializer: &mut S,
    ) -> Result<S::Ok, S::Error> {
        serializer.serialize_bytes(self.0.as_slice())
    }
}

impl UnbufferedService for AuthMgrService {
    type Connection = AuthMgrConnection;

    fn on_connect(
        &self,
        _port: &PortCfg,
        handle: &Handle,
        peer: &Uuid,
    ) -> Result<ConnectResult<Self::Connection>, TipcError> {
        debug!("Accepted AthMgr BE connection from uuid: {:?}, handle: {:?}", peer, handle);
        Ok(ConnectResult::Accept(AuthMgrConnection {
            uuid: peer.clone(),
            pending_to_be_routed: Arc::new(Mutex::new(true)),
        }))
    }

    fn on_message(
        &self,
        connection: &Self::Connection,
        handle: &Handle,
        _buffer: &mut [u8],
    ) -> Result<MessageResult, TipcError> {
        let mut connection_pending_to_be_routed = connection.pending_to_be_routed.lock().unwrap();
        debug!(
            "AuthMgr BE received a message from uuid: {:?} for a pending_to_be_routed: {:?}, handle: {:?}.",
            connection.uuid,
            connection_pending_to_be_routed,
            handle
        );
        if *connection_pending_to_be_routed {
            let mut temp_buffer = [0u8; MAX_MSG_SIZE];
            let msg: AuthMgrMessage = handle.recv(&mut temp_buffer)?;
            if msg.0.is_empty() {
                error!("Expected a command byte, got an empty payload.");
                return Err(TipcError::InvalidData);
            }
            match msg.0[0] {
                CMD_RAW => {
                    debug!("It's a request for a new raw connection.");
                    if msg.0.len() != TOKEN_LENGTH + 1 {
                        error!(
                            "Expected a token of length: {:?}, got a payload of size: {:?}.",
                            TOKEN_LENGTH,
                            msg.0.len()
                        );
                        return Err(TipcError::InvalidData);
                    }
                    let token: [u8; TOKEN_LENGTH] =
                        (msg.0[1..TOKEN_LENGTH + 1]).try_into().map_err(|e| {
                            error!("Failed to read the token: {:?}", e);
                            TipcError::InvalidData
                        })?;
                    let client_connection = RawConnectionInformation::new(
                        handle.try_clone()?,
                        // TODO: b/392905377. This is temporary, until the VM-ID is available.
                        TEMP_AUTHMGR_FE_TRANSPORT_ID,
                    );
                    let mut global_state = self.global_state.lock().unwrap();
                    global_state
                        .authmgr_core
                        .init_connection_for_client(Box::new(client_connection), token)
                        .map_err(|e| {
                            error!("Failed to call init_connection_for_client: {:?}", e);
                            TipcError::UnknownError
                        })?;
                    *connection_pending_to_be_routed = false;
                    // Close the handle registered in the handle set
                    Ok(MessageResult::CloseConnection)
                }
                CMD_RPC => {
                    debug!("It's a request for a RPC binder connection.");
                    let authmgr_service_port_cfg = PortCfg::new_raw(AUTHMGR_SERVICE_PORT.into())
                        .allow_ta_connect()
                        .allow_ns_connect();
                    let rpc_connection = match self.rpc_service.on_connect(
                        &authmgr_service_port_cfg,
                        handle,
                        &connection.uuid,
                    )? {
                        ConnectResult::Accept(conn) => conn,
                        ConnectResult::CloseConnection => {
                            return Ok(MessageResult::CloseConnection)
                        }
                    };
                    debug!("AuthMgr RPC service.on_connect returned. Adding to the handle set.");
                    self.handle_set.upgrade().ok_or(TipcError::InvalidData)?.add_connection(
                        AuthMgrServicesConnection::RpcService(rpc_connection),
                        handle.try_clone()?,
                        Arc::new(AuthMgrServices::RpcService(Arc::clone(&self.rpc_service))),
                    )?;
                    *connection_pending_to_be_routed = false;
                    // Close the handle registered in the handle set
                    Ok(MessageResult::CloseConnection)
                }
                _ => {
                    error!("Unknown command.");
                    Err(TipcError::InvalidData)
                }
            }
        } else {
            error!(
                "No messages are expected on the same connection
                    after the initial command byte is sent."
            );
            Err(TipcError::InvalidData)
        }
    }
    // TODO: implement on_disconnect
}

enum AuthMgrServices {
    InitService(Arc<AuthMgrService>),
    RpcService(Arc<RpcServer>),
}

enum AuthMgrServicesConnection {
    InitService(AuthMgrConnection),
    RpcService(RpcServerConnection),
}

impl From<AuthMgrConnection> for AuthMgrServicesConnection {
    fn from(v: AuthMgrConnection) -> AuthMgrServicesConnection {
        AuthMgrServicesConnection::InitService(v)
    }
}

impl From<RpcServerConnection> for AuthMgrServicesConnection {
    fn from(v: RpcServerConnection) -> AuthMgrServicesConnection {
        AuthMgrServicesConnection::RpcService(v)
    }
}

impl<'a> TryFrom<&'a AuthMgrServicesConnection> for &'a AuthMgrConnection {
    type Error = TipcError;
    fn try_from(v: &'a AuthMgrServicesConnection) -> Result<&'a AuthMgrConnection, Self::Error> {
        match v {
            AuthMgrServicesConnection::InitService(amgr_conn) => Ok(amgr_conn),
            _ => Err(TipcError::InvalidData),
        }
    }
}

impl<'a> TryFrom<&'a AuthMgrServicesConnection> for &'a RpcServerConnection {
    type Error = TipcError;
    fn try_from(v: &'a AuthMgrServicesConnection) -> Result<&'a RpcServerConnection, Self::Error> {
        match v {
            AuthMgrServicesConnection::RpcService(rpc_conn) => Ok(rpc_conn),
            _ => Err(TipcError::InvalidData),
        }
    }
}

impl UnbufferedService for AuthMgrServices {
    type Connection = AuthMgrServicesConnection;

    fn on_connect(
        &self,
        port: &PortCfg,
        handle: &Handle,
        peer: &Uuid,
    ) -> Result<ConnectResult<Self::Connection>, TipcError> {
        match self {
            AuthMgrServices::InitService(authmgr_service) => {
                match authmgr_service.on_connect(port, handle, peer) {
                    Ok(conn_result) => match conn_result {
                        ConnectResult::Accept(conn) => Ok(ConnectResult::Accept(conn.into())),
                        ConnectResult::CloseConnection => Ok(ConnectResult::CloseConnection),
                    },
                    Err(e) => Err(e),
                }
            }

            AuthMgrServices::RpcService(authmgr_rpc_service) => {
                match authmgr_rpc_service.on_connect(port, handle, peer) {
                    Ok(conn_result) => match conn_result {
                        ConnectResult::Accept(conn) => Ok(ConnectResult::Accept(conn.into())),
                        ConnectResult::CloseConnection => Ok(ConnectResult::CloseConnection),
                    },
                    Err(e) => Err(e),
                }
            }
        }
    }

    fn on_message(
        &self,
        connection: &Self::Connection,
        handle: &Handle,
        buffer: &mut [u8],
    ) -> Result<MessageResult, TipcError> {
        match self {
            AuthMgrServices::InitService(authmgr_service) => authmgr_service.on_message(
                connection.try_into().map_err(|_| TipcError::InvalidData)?,
                handle,
                buffer,
            ),
            AuthMgrServices::RpcService(authmgr_rpc_service) => authmgr_rpc_service.on_message(
                connection.try_into().map_err(|_| TipcError::InvalidData)?,
                handle,
                buffer,
            ),
        }
    }
}

pub fn main_loop() -> Result<(), TipcError> {
    let handle_set_wrapper = Arc::new(HandleSetWrapper::<AuthMgrServices>::new()?);
    let global_state = AuthMgrGlobalState::new().map_err(|e| {
        error!("Could not create AuthMgr global state due to: {:?}", e);
        TipcError::UnknownError
    })?;
    let global_state = Arc::new(Mutex::new(global_state));
    let gs_clone = Arc::clone(&global_state);
    let authmgr_service_port_cfg =
        PortCfg::new_raw(AUTHMGR_SERVICE_PORT.into()).allow_ta_connect().allow_ns_connect();

    let cb_per_session = move |_uuid| {
        AuthMgrAuthorizationRPCService::new_authorization_session(Arc::clone(&global_state))
    };
    let authmgr_authorization_rpc_server = RpcServer::new_per_session(cb_per_session);
    let authmgr_service = AuthMgrService::new(
        gs_clone,
        Arc::new(authmgr_authorization_rpc_server),
        Arc::downgrade(&handle_set_wrapper),
    );

    let _port_wrapper = handle_set_wrapper.add_port(
        &authmgr_service_port_cfg,
        Arc::new(AuthMgrServices::InitService(Arc::new(authmgr_service))),
    )?;
    let event_loop = EventLoop::new(handle_set_wrapper.clone());
    event_loop.run()
}