Skip to content

Commit

Permalink
work
Browse files Browse the repository at this point in the history
  • Loading branch information
@rok
rok committed Jan 7, 2025
1 parent 1123b1a commit a4105d5
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 31 deletions.
6 changes: 5 additions & 1 deletion parquet/src/encryption/ciphers.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ pub fn create_page_aad(file_aad: &[u8], module_type: ModuleType, row_group_ordin
create_module_aad(file_aad, module_type, row_group_ordinal, column_ordinal, page_ordinal)
}

fn create_module_aad(file_aad: &[u8], module_type: ModuleType, row_group_ordinal: usize,
pub fn create_module_aad(file_aad: &[u8], module_type: ModuleType, row_group_ordinal: usize,
column_ordinal: usize, page_ordinal: Option<usize>) -> Result<Vec<u8>> {

let module_buf = [module_type as u8];
Expand Down Expand Up @@ -325,6 +325,10 @@ impl FileDecryptor {
self.footer_decryptor.unwrap()
}

pub(crate) fn column_decryptor(&self) -> RingGcmBlockDecryptor {
RingGcmBlockDecryptor::new(self.decryption_properties.footer_key.as_ref().unwrap())
}

pub(crate) fn get_column_decryptor(&self, column_name: &[u8]) -> FileDecryptor {
if self.decryption_properties.column_keys.is_none() {
return self.clone();
Expand Down
31 changes: 18 additions & 13 deletions parquet/src/file/metadata/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,9 @@ use crate::basic::{ColumnOrder, Compression, Encoding, Type};
use crate::data_type::AsBytes;
#[cfg(feature = "encryption")]
use crate::encryption::ciphers::FileDecryptor;
use crate::encryption::ciphers::{create_footer_aad, create_page_aad, ModuleType};
use crate::encryption::ciphers::{
create_footer_aad, create_module_aad, create_page_aad, ModuleType,
};
use crate::encryption::ciphers::{
BlockDecryptor, DecryptionPropertiesBuilder, FileDecryptionProperties,
};
Expand Down Expand Up @@ -671,29 +673,32 @@ impl RowGroupMetaData {
cc = ColumnChunkMetaData::from_thrift(d.clone(), c)?;
} else {
let column_name = crypto_metadata.path_in_schema.join(".");
let column_decryptor = decryptor
.unwrap()
.get_column_decryptor(column_name.as_bytes());
let file_decryptor = column_decryptor.footer_decryptor().unwrap();

let aad_file_unique = decryptor.unwrap().aad_file_unique();
let aad_prefix = decryptor
.unwrap()
.decryption_properties()
.aad_prefix()
.unwrap();
let aad: Vec<u8> = [aad_prefix.clone(), aad_file_unique.clone()].concat();
let column_aad = create_page_aad(
aad.as_slice(),

let column_decryptor = decryptor
.unwrap()
.get_column_decryptor(column_name.as_bytes())
.footer_decryptor()
.unwrap();

let column_aad = create_module_aad(
[aad_prefix.as_slice(), aad_file_unique.as_slice()]
.concat()
.as_slice(),
ModuleType::ColumnMetaData,
rg.ordinal.unwrap() as usize,
i,
i as usize,
None,
)?;

let mut buf = c.encrypted_column_metadata.unwrap();
let mut decrypted_cc_buf =
file_decryptor.decrypt(buf.as_slice().as_ref(), column_aad.as_ref())?;
let buf = c.encrypted_column_metadata.unwrap();
let decrypted_cc_buf =
column_decryptor.decrypt(buf.as_slice().as_ref(), column_aad.as_ref())?;

let mut prot = TCompactSliceInputProtocol::new(decrypted_cc_buf.as_slice());
let c = ColumnChunk::read_from_in_protocol(&mut prot)?;
Expand Down
37 changes: 20 additions & 17 deletions parquet/src/file/serialized_reader.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,8 @@ use crate::basic::{Encoding, Type};
use crate::bloom_filter::Sbbf;
use crate::column::page::{Page, PageMetadata, PageReader};
use crate::compression::{create_codec, Codec};
use crate::encryption::ciphers::RingGcmBlockDecryptor;
#[cfg(feature = "encryption")]
use crate::encryption::ciphers::{create_page_aad, BlockDecryptor, CryptoContext, ModuleType};
use crate::encryption::ciphers::{create_module_aad, BlockDecryptor, CryptoContext, ModuleType};
use crate::errors::{ParquetError, Result};
use crate::file::page_index::offset_index::OffsetIndexMetaData;
use crate::file::{
Expand Down Expand Up @@ -346,16 +345,13 @@ pub(crate) fn read_page_header<T: Read>(
input: &mut T,
#[cfg(feature = "encryption")] crypto_context: Option<Arc<CryptoContext>>,
) -> Result<PageHeader> {
// todo: if column is not encrypted skip decryption

#[cfg(feature = "encryption")]
if let Some(crypto_context) = crypto_context {
// crypto_context.data_decryptor().get_column_decryptor()
let decryptor = &crypto_context.data_decryptor();
// todo: get column decryptor
// let file_decryptor = decryptor.ge(crypto_context.column_ordinal);
// if !decryptor.decryption_properties().has_footer_key() {
// return Err(general_err!("Missing footer decryptor"));
// }
let file_decryptor = decryptor.footer_decryptor();

let file_decryptor = decryptor.column_decryptor();
let aad_file_unique = decryptor.aad_file_unique();
let aad_prefix = decryptor.aad_prefix();

Expand All @@ -364,7 +360,7 @@ pub(crate) fn read_page_header<T: Read>(
} else {
ModuleType::DataPageHeader
};
let aad = create_page_aad(
let aad = create_module_aad(
[aad_prefix.as_slice(), aad_file_unique.as_slice()]
.concat()
.as_slice(),
Expand All @@ -374,12 +370,19 @@ pub(crate) fn read_page_header<T: Read>(
crypto_context.page_ordinal,
)?;

let mut len_bytes = [0; 4];
input.read_exact(&mut len_bytes)?;
let ciphertext_len = u32::from_le_bytes(len_bytes) as usize;
let mut ciphertext = vec![0; 4 + ciphertext_len];
input.read_exact(&mut ciphertext[4..])?;
let buf = file_decryptor.unwrap().decrypt(&ciphertext, aad.as_ref())?;
// let mut len_bytes = [0; 4];
// input.read_exact(&mut len_bytes)?;
// let ciphertext_len = u32::from_le_bytes(len_bytes) as usize;
// let mut ciphertext = vec![0; 4 + ciphertext_len];
// input.read_exact(&mut ciphertext[4..])?;
// let mut ciphertext = Vec::new();
// input.read_to_end(&mut ciphertext)?;

let mut ciphertext: Vec<u8> = vec![];
input.read_to_end(&mut ciphertext)?;

// let ciphertext = input.read_to_end();
let buf = file_decryptor.decrypt(&ciphertext, aad.as_ref())?;

let mut prot = TCompactSliceInputProtocol::new(buf.as_slice());
let page_header = PageHeader::read_from_in_protocol(&mut prot)?;
Expand Down Expand Up @@ -478,7 +481,7 @@ pub(crate) fn decode_page(
} else {
ModuleType::DataPage
};
let aad = create_page_aad(
let aad = create_module_aad(
decryptor.aad_file_unique().as_slice(),
module_type,
crypto_context.row_group_ordinal,
Expand Down

0 comments on commit a4105d5

Please sign in to comment.