| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you 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 crate::error::AuthenticationError; |
| use crate::error::FrontendServiceError; |
| |
| use anyhow::Result; |
| use std::net::{IpAddr, Ipv6Addr}; |
| use std::sync::Arc; |
| use teaclave_proto::teaclave_access_control_service::{ |
| AuthorizeApiRequest, TeaclaveAccessControlClient, |
| }; |
| use teaclave_proto::teaclave_authentication_service::{ |
| TeaclaveAuthenticationInternalClient, UserAuthenticateRequest, |
| }; |
| use teaclave_proto::teaclave_common::UserCredential; |
| use teaclave_proto::teaclave_frontend_service::{ |
| ApproveTaskRequest, AssignDataRequest, CancelTaskRequest, CreateTaskRequest, |
| CreateTaskResponse, DeleteFunctionRequest, DisableFunctionRequest, GetFunctionRequest, |
| GetFunctionResponse, GetFunctionUsageStatsRequest, GetFunctionUsageStatsResponse, |
| GetInputFileRequest, GetInputFileResponse, GetOutputFileRequest, GetOutputFileResponse, |
| GetTaskRequest, GetTaskResponse, InvokeTaskRequest, ListFunctionsRequest, |
| ListFunctionsResponse, QueryAuditLogsRequest, QueryAuditLogsResponse, RegisterFunctionRequest, |
| RegisterFunctionResponse, RegisterFusionOutputRequest, RegisterFusionOutputResponse, |
| RegisterInputFileRequest, RegisterInputFileResponse, RegisterInputFromOutputRequest, |
| RegisterInputFromOutputResponse, RegisterOutputFileRequest, RegisterOutputFileResponse, |
| TeaclaveFrontend, UpdateFunctionRequest, UpdateFunctionResponse, UpdateInputFileRequest, |
| UpdateInputFileResponse, UpdateOutputFileRequest, UpdateOutputFileResponse, |
| }; |
| use teaclave_proto::teaclave_management_service::TeaclaveManagementClient; |
| use teaclave_rpc::transport::Channel; |
| use teaclave_rpc::Request; |
| use teaclave_service_enclave_utils::bail; |
| use teaclave_types::{Entry, EntryBuilder, TeaclaveServiceResponseResult, UserAuthClaims}; |
| use tokio::sync::Mutex; |
| |
| macro_rules! authentication_and_forward_to_management { |
| ($service: ident, $request: ident, $func: ident) => {{ |
| let function_name = stringify!($func).to_owned(); |
| let ip_option = $request.remote_addr().map(|s| s.ip()); |
| let ip = match ip_option { |
| Some(IpAddr::V4(ip_v4)) => ip_v4.to_ipv6_compatible(), |
| Some(IpAddr::V6(ip_v6)) => ip_v6, |
| None => Ipv6Addr::UNSPECIFIED, |
| }; |
| |
| let builder = EntryBuilder::new().ip(ip); |
| |
| let claims = match $service.authenticate(&$request).await { |
| Ok(claims) => { |
| if $service |
| .check_api_privilege( |
| claims.get_role().to_string().split('-').next().unwrap(), |
| stringify!($func), |
| ) |
| .await |
| { |
| claims |
| } else { |
| log::debug!( |
| "User is not authorized to access func: {}", |
| stringify!($func) |
| ); |
| |
| let entry = builder |
| .message(String::from("authenticate to ") + &function_name) |
| .result(false) |
| .build(); |
| $service.push_log(entry).await; |
| |
| bail!(FrontendServiceError::PermissionDenied); |
| } |
| } |
| Err(e) => { |
| log::debug!( |
| "User is not authenticated to access func: {}", |
| stringify!($func) |
| ); |
| |
| let entry = builder |
| .message( |
| String::from("authenticate to ") + &function_name + ": " + &e.to_string(), |
| ) |
| .result(false) |
| .build(); |
| $service.push_log(entry).await; |
| |
| bail!(e); |
| } |
| }; |
| |
| let user = claims.to_string(); |
| let builder = builder.user(user); |
| |
| let client = $service.management_client.clone(); |
| let mut client = client.lock().await; |
| let meta = $request.metadata().clone(); |
| let message = $request.get_ref().to_owned(); |
| |
| let mut request = Request::new(message); |
| let metadata = request.metadata_mut(); |
| *metadata = meta; |
| metadata.insert("role", claims.role.parse().unwrap()); |
| |
| let response = match client.$func(request).await { |
| Err(e) => { |
| let entry = builder |
| .clone() |
| .message(function_name.clone() + ":" + &e.to_string()) |
| .result(false) |
| .build(); |
| $service.push_log(entry).await; |
| return Err(e); |
| } |
| Ok(r) => r, |
| }; |
| |
| let entry = builder.message(function_name).result(true).build(); |
| $service.push_log(entry).await; |
| Ok(response) |
| }}; |
| } |
| |
| #[derive(Clone)] |
| pub(crate) struct TeaclaveFrontendService { |
| authentication_client: Arc<Mutex<TeaclaveAuthenticationInternalClient<Channel>>>, |
| management_client: Arc<Mutex<TeaclaveManagementClient<Channel>>>, |
| access_control_client: Arc<Mutex<TeaclaveAccessControlClient<Channel>>>, |
| audit_log_buffer: Arc<Mutex<Vec<Entry>>>, |
| } |
| |
| impl TeaclaveFrontendService { |
| pub(crate) async fn new( |
| authentication_client: Arc<Mutex<TeaclaveAuthenticationInternalClient<Channel>>>, |
| management_client: Arc<Mutex<TeaclaveManagementClient<Channel>>>, |
| access_control_client: Arc<Mutex<TeaclaveAccessControlClient<Channel>>>, |
| audit_log_buffer: Arc<Mutex<Vec<Entry>>>, |
| ) -> Result<Self> { |
| Ok(Self { |
| authentication_client, |
| management_client, |
| access_control_client, |
| audit_log_buffer, |
| }) |
| } |
| |
| pub async fn push_log(&self, entry: Entry) { |
| let mut buffer_lock = self.audit_log_buffer.lock().await; |
| buffer_lock.push(entry); |
| } |
| |
| async fn check_api_privilege(&self, user_role: &str, api: &str) -> bool { |
| let request = AuthorizeApiRequest { |
| user_role: user_role.to_owned(), |
| api: api.to_owned(), |
| }; |
| |
| let mut acs_client = self.access_control_client.lock().await; |
| let result = acs_client.authorize_api(request).await; |
| result.map(|r| r.into_inner().accept).unwrap_or(false) |
| } |
| } |
| |
| #[teaclave_rpc::async_trait] |
| impl TeaclaveFrontend for TeaclaveFrontendService { |
| async fn register_input_file( |
| &self, |
| request: Request<RegisterInputFileRequest>, |
| ) -> TeaclaveServiceResponseResult<RegisterInputFileResponse> { |
| authentication_and_forward_to_management!(self, request, register_input_file) |
| } |
| |
| async fn update_input_file( |
| &self, |
| request: Request<UpdateInputFileRequest>, |
| ) -> TeaclaveServiceResponseResult<UpdateInputFileResponse> { |
| authentication_and_forward_to_management!(self, request, update_input_file) |
| } |
| |
| async fn register_output_file( |
| &self, |
| request: Request<RegisterOutputFileRequest>, |
| ) -> TeaclaveServiceResponseResult<RegisterOutputFileResponse> { |
| authentication_and_forward_to_management!(self, request, register_output_file) |
| } |
| |
| async fn update_output_file( |
| &self, |
| request: Request<UpdateOutputFileRequest>, |
| ) -> TeaclaveServiceResponseResult<UpdateOutputFileResponse> { |
| authentication_and_forward_to_management!(self, request, update_output_file) |
| } |
| |
| async fn register_fusion_output( |
| &self, |
| request: Request<RegisterFusionOutputRequest>, |
| ) -> TeaclaveServiceResponseResult<RegisterFusionOutputResponse> { |
| authentication_and_forward_to_management!(self, request, register_fusion_output) |
| } |
| |
| async fn register_input_from_output( |
| &self, |
| request: Request<RegisterInputFromOutputRequest>, |
| ) -> TeaclaveServiceResponseResult<RegisterInputFromOutputResponse> { |
| authentication_and_forward_to_management!(self, request, register_input_from_output) |
| } |
| |
| async fn get_output_file( |
| &self, |
| request: Request<GetOutputFileRequest>, |
| ) -> TeaclaveServiceResponseResult<GetOutputFileResponse> { |
| authentication_and_forward_to_management!(self, request, get_output_file) |
| } |
| |
| async fn get_input_file( |
| &self, |
| request: Request<GetInputFileRequest>, |
| ) -> TeaclaveServiceResponseResult<GetInputFileResponse> { |
| authentication_and_forward_to_management!(self, request, get_input_file) |
| } |
| |
| async fn register_function( |
| &self, |
| request: Request<RegisterFunctionRequest>, |
| ) -> TeaclaveServiceResponseResult<RegisterFunctionResponse> { |
| authentication_and_forward_to_management!(self, request, register_function) |
| } |
| |
| async fn update_function( |
| &self, |
| request: Request<UpdateFunctionRequest>, |
| ) -> TeaclaveServiceResponseResult<UpdateFunctionResponse> { |
| authentication_and_forward_to_management!(self, request, update_function) |
| } |
| |
| async fn get_function( |
| &self, |
| request: Request<GetFunctionRequest>, |
| ) -> TeaclaveServiceResponseResult<GetFunctionResponse> { |
| authentication_and_forward_to_management!(self, request, get_function) |
| } |
| |
| async fn get_function_usage_stats( |
| &self, |
| request: Request<GetFunctionUsageStatsRequest>, |
| ) -> TeaclaveServiceResponseResult<GetFunctionUsageStatsResponse> { |
| authentication_and_forward_to_management!(self, request, get_function_usage_stats) |
| } |
| |
| async fn delete_function( |
| &self, |
| request: Request<DeleteFunctionRequest>, |
| ) -> TeaclaveServiceResponseResult<()> { |
| authentication_and_forward_to_management!(self, request, delete_function) |
| } |
| |
| async fn disable_function( |
| &self, |
| request: Request<DisableFunctionRequest>, |
| ) -> TeaclaveServiceResponseResult<()> { |
| authentication_and_forward_to_management!(self, request, disable_function) |
| } |
| |
| async fn list_functions( |
| &self, |
| request: Request<ListFunctionsRequest>, |
| ) -> TeaclaveServiceResponseResult<ListFunctionsResponse> { |
| authentication_and_forward_to_management!(self, request, list_functions) |
| } |
| |
| async fn create_task( |
| &self, |
| request: Request<CreateTaskRequest>, |
| ) -> TeaclaveServiceResponseResult<CreateTaskResponse> { |
| authentication_and_forward_to_management!(self, request, create_task) |
| } |
| |
| async fn get_task( |
| &self, |
| request: Request<GetTaskRequest>, |
| ) -> TeaclaveServiceResponseResult<GetTaskResponse> { |
| authentication_and_forward_to_management!(self, request, get_task) |
| } |
| |
| async fn assign_data( |
| &self, |
| request: Request<AssignDataRequest>, |
| ) -> TeaclaveServiceResponseResult<()> { |
| authentication_and_forward_to_management!(self, request, assign_data) |
| } |
| |
| async fn approve_task( |
| &self, |
| request: Request<ApproveTaskRequest>, |
| ) -> TeaclaveServiceResponseResult<()> { |
| authentication_and_forward_to_management!(self, request, approve_task) |
| } |
| |
| async fn invoke_task( |
| &self, |
| request: Request<InvokeTaskRequest>, |
| ) -> TeaclaveServiceResponseResult<()> { |
| authentication_and_forward_to_management!(self, request, invoke_task) |
| } |
| |
| async fn cancel_task( |
| &self, |
| request: Request<CancelTaskRequest>, |
| ) -> TeaclaveServiceResponseResult<()> { |
| authentication_and_forward_to_management!(self, request, cancel_task) |
| } |
| |
| async fn query_audit_logs( |
| &self, |
| request: Request<QueryAuditLogsRequest>, |
| ) -> TeaclaveServiceResponseResult<QueryAuditLogsResponse> { |
| authentication_and_forward_to_management!(self, request, query_audit_logs) |
| } |
| } |
| |
| impl TeaclaveFrontendService { |
| async fn authenticate<T>( |
| &self, |
| request: &Request<T>, |
| ) -> Result<UserAuthClaims, FrontendServiceError> { |
| let id = request |
| .metadata() |
| .get("id") |
| .and_then(|x| x.to_str().ok()) |
| .ok_or(AuthenticationError::MissingUserId)?; |
| let token = request |
| .metadata() |
| .get("token") |
| .and_then(|x| x.to_str().ok()) |
| .ok_or(AuthenticationError::MissingToken)?; |
| let credential = Some(UserCredential::new(id, token)); |
| let auth_request = UserAuthenticateRequest { credential }; |
| let claims = self |
| .authentication_client |
| .clone() |
| .lock() |
| .await |
| .user_authenticate(auth_request) |
| .await |
| .map_err(|_| AuthenticationError::IncorrectCredential)? |
| .into_inner() |
| .claims |
| .and_then(|x| x.try_into().ok()) |
| .ok_or(AuthenticationError::IncorrectCredential)?; |
| |
| Ok(claims) |
| } |
| } |