use std::sync::Arc;
use async_trait::async_trait;
use object_store::DynObjectStore;
use unitycatalog_common::models::credentials::v1::GetCredentialRequest;
use unitycatalog_common::models::tables::v1::{GetTableRequest, Table};
use unitycatalog_common::models::volumes::v1::{GetVolumeRequest, Volume};
use unitycatalog_object_store::store_from_vended_credential;
use unitycatalog_storage_proxy::backend::reject_key_traversal;
use unitycatalog_storage_proxy::{
ProxyCapabilities, ProxyError, ProxyReq, ProxyResult, Securable, StorageProxyBackend,
};
use crate::Error;
use crate::api::RequestContext;
use crate::api::credentials::CredentialHandlerExt;
use crate::api::tables::TableHandler;
use crate::api::volumes::VolumeHandler;
use crate::policy::{Permission, Policy};
use crate::services::ServerHandler;
use crate::services::credential_vending::{VendOperation, vend_credential};
use crate::services::location::StorageLocationUrl;
use crate::services::object_store::find_external_location_for_url;
fn to_proxy_err(e: Error) -> ProxyError {
match e {
Error::NotFound => ProxyError::NotFound("resource not found".into()),
Error::NotAllowed => ProxyError::PermissionDenied("not authorized".into()),
other => ProxyError::Internal(other.to_string()),
}
}
fn to_proxy_err_common(e: unitycatalog_common::Error) -> ProxyError {
match e {
unitycatalog_common::Error::NotFound => ProxyError::NotFound("resource not found".into()),
unitycatalog_common::Error::InvalidArgument(m) => ProxyError::InvalidArgument(m),
unitycatalog_common::Error::InvalidUrl(e) => {
ProxyError::InvalidArgument(format!("invalid storage URL: {e}"))
}
other => ProxyError::Internal(other.to_string()),
}
}
fn required_permission(write: bool) -> Permission {
if write {
Permission::Write
} else {
Permission::Read
}
}
fn vend_operation(write: bool) -> VendOperation {
if write {
VendOperation::ReadWrite
} else {
VendOperation::Read
}
}
async fn authorize_and_locate(
handler: &ServerHandler<RequestContext>,
securable: &Securable,
write: bool,
cx: &RequestContext,
) -> ProxyResult<String> {
let permission = required_permission(write);
match securable {
Securable::Table { full_name } => {
let table: Table = TableHandler::get_table(
handler,
GetTableRequest {
full_name: full_name.clone(),
..Default::default()
},
cx.clone(),
)
.await
.map_err(to_proxy_err)?;
if write {
let id = table.table_id.as_deref().ok_or_else(|| {
ProxyError::Internal("table has no id to authorize write".into())
})?;
authorize_table_write(handler, id, cx).await?;
}
table
.storage_location
.filter(|s| !s.is_empty())
.ok_or_else(|| ProxyError::NotFound("table has no storage location".into()))
}
Securable::Volume { full_name } => {
let volume: Volume = VolumeHandler::get_volume(
handler,
GetVolumeRequest {
name: full_name.clone(),
..Default::default()
},
cx.clone(),
)
.await
.map_err(to_proxy_err)?;
if write {
authorize_volume_write(handler, &volume.volume_id, cx).await?;
}
Ok(volume.storage_location)
}
Securable::Path { url } => {
let storage_url =
StorageLocationUrl::parse(url.as_str()).map_err(to_proxy_err_common)?;
let ext_loc = find_external_location_for_url(&storage_url, handler)
.await
.map_err(to_proxy_err)?;
handler
.authorize_checked(&(&ext_loc).into(), &permission, cx)
.await
.map_err(to_proxy_err)?;
Ok(url.to_string())
}
}
}
async fn authorize_table_write(
handler: &ServerHandler<RequestContext>,
table_id: &str,
cx: &RequestContext,
) -> ProxyResult<()> {
use unitycatalog_common::models::{ResourceIdent, ResourceRef};
let uuid = uuid::Uuid::parse_str(table_id)
.map_err(|_| ProxyError::Internal("table id is not a valid UUID".into()))?;
let ident = ResourceIdent::Table(ResourceRef::Uuid(uuid));
handler
.authorize_checked(&ident, &Permission::Write, cx)
.await
.map_err(to_proxy_err)
}
async fn authorize_volume_write(
handler: &ServerHandler<RequestContext>,
volume_id: &str,
cx: &RequestContext,
) -> ProxyResult<()> {
use unitycatalog_common::models::{ResourceIdent, ResourceRef};
let uuid = uuid::Uuid::parse_str(volume_id)
.map_err(|_| ProxyError::Internal("volume id is not a valid UUID".into()))?;
let ident = ResourceIdent::Volume(ResourceRef::Uuid(uuid));
handler
.authorize_checked(&ident, &Permission::Write, cx)
.await
.map_err(to_proxy_err)
}
async fn open_store(
handler: &ServerHandler<RequestContext>,
location: &str,
write: bool,
) -> ProxyResult<Arc<DynObjectStore>> {
let storage_url = StorageLocationUrl::parse(location).map_err(to_proxy_err_common)?;
let ext_loc = find_external_location_for_url(&storage_url, handler)
.await
.map_err(to_proxy_err)?;
let credential = handler
.get_credential_internal(GetCredentialRequest {
name: ext_loc.credential_name.clone(),
..Default::default()
})
.await
.map_err(to_proxy_err)?;
let vended = vend_credential(&credential, location, vend_operation(write))
.await
.map_err(to_proxy_err)?;
let store = store_from_vended_credential(&vended, None, None)
.map_err(|e| ProxyError::Internal(format!("build store: {e}")))?;
Ok(store.as_dyn())
}
#[async_trait]
impl StorageProxyBackend<RequestContext> for ServerHandler<RequestContext> {
fn capabilities(&self) -> ProxyCapabilities {
ProxyCapabilities { enabled: true }
}
async fn authorize(&self, req: &ProxyReq, cx: &RequestContext) -> ProxyResult<()> {
reject_key_traversal(&req.key)?;
authorize_and_locate(self, &req.securable, req.verb.is_write(), cx).await?;
Ok(())
}
async fn open(&self, req: &ProxyReq, cx: &RequestContext) -> ProxyResult<Arc<DynObjectStore>> {
let write = req.verb.is_write();
let location = authorize_and_locate(self, &req.securable, write, cx).await?;
open_store(self, &location, write).await
}
}