1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing, software
12// distributed under the License is distributed on an "AS IS" BASIS,
13// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14// See the License for the specific language governing permissions and
15// limitations under the License.
16
17package encryption
18
19import (
20	"github.com/apache/arrow/go/v6/arrow/memory"
21	"github.com/apache/arrow/go/v6/parquet"
22)
23
24// FileDecryptor is an interface used by the filereader for decrypting an
25// entire parquet file as we go, usually constructed from the DecryptionProperties
26type FileDecryptor interface {
27	// Returns the key for decrypting the footer if provided
28	GetFooterKey() string
29	// Provides the file level AAD security bytes
30	FileAad() string
31	// return which algorithm this decryptor was constructed for
32	Algorithm() parquet.Cipher
33	// return the FileDecryptionProperties that were used for this decryptor
34	Properties() *parquet.FileDecryptionProperties
35	// Clear out the decryption keys, this is automatically called after every
36	// successfully decrypted file to ensure that keys aren't kept around.
37	WipeOutDecryptionKeys()
38	// GetFooterDecryptor returns a Decryptor interface for use to decrypt the footer
39	// of a parquet file.
40	GetFooterDecryptor() Decryptor
41	// GetFooterDecryptorForColumnMeta returns a Decryptor interface for Column Metadata
42	// in the file footer using the AAD bytes provided.
43	GetFooterDecryptorForColumnMeta(aad string) Decryptor
44	// GetFooterDecryptorForColumnData returns the decryptor that can be used for decrypting
45	// actual column data footer bytes, not column metadata.
46	GetFooterDecryptorForColumnData(aad string) Decryptor
47	// GetColumnMetaDecryptor returns a decryptor for the requested column path, key and AAD bytes
48	// but only for decrypting the row group level metadata
49	GetColumnMetaDecryptor(columnPath, columnKeyMetadata, aad string) Decryptor
50	// GetColumnDataDecryptor returns a decryptor for the requested column path, key, and AAD bytes
51	// but only for the rowgroup column data.
52	GetColumnDataDecryptor(columnPath, columnKeyMetadata, aad string) Decryptor
53}
54
55type fileDecryptor struct {
56	// the properties contains the key retriever for us to get keys
57	// from the key metadata
58	props *parquet.FileDecryptionProperties
59	// concatenation of aad_prefix (if exists) and aad_file_unique
60	fileAad                 string
61	columnDataMap           map[string]Decryptor
62	columnMetaDataMap       map[string]Decryptor
63	footerMetadataDecryptor Decryptor
64	footerDataDecryptor     Decryptor
65	alg                     parquet.Cipher
66	footerKeyMetadata       string
67	metaDecryptor           *aesDecryptor
68	dataDecryptor           *aesDecryptor
69	mem                     memory.Allocator
70}
71
72// NewFileDecryptor constructs a decryptor from the provided configuration of properties, cipher and key metadata. Using the provided memory allocator or
73// the default allocator if one isn't provided.
74func NewFileDecryptor(props *parquet.FileDecryptionProperties, fileAad string, alg parquet.Cipher, keymetadata string, mem memory.Allocator) FileDecryptor {
75	if mem == nil {
76		mem = memory.DefaultAllocator
77	}
78	return &fileDecryptor{
79		fileAad:           fileAad,
80		props:             props,
81		alg:               alg,
82		footerKeyMetadata: keymetadata,
83		mem:               mem,
84		columnDataMap:     make(map[string]Decryptor),
85		columnMetaDataMap: make(map[string]Decryptor),
86	}
87}
88
89func (d *fileDecryptor) FileAad() string                               { return d.fileAad }
90func (d *fileDecryptor) Properties() *parquet.FileDecryptionProperties { return d.props }
91func (d *fileDecryptor) Algorithm() parquet.Cipher                     { return d.alg }
92func (d *fileDecryptor) GetFooterKey() string {
93	footerKey := d.props.FooterKey()
94	if footerKey == "" {
95		if d.footerKeyMetadata == "" {
96			panic("no footer key or key metadata")
97		}
98		if d.props.KeyRetriever == nil {
99			panic("no footer key or key retriever")
100		}
101		footerKey = d.props.KeyRetriever.GetKey([]byte(d.footerKeyMetadata))
102	}
103	if footerKey == "" {
104		panic("invalid footer encryption key. Could not parse footer metadata")
105	}
106	return footerKey
107}
108
109func (d *fileDecryptor) GetFooterDecryptor() Decryptor {
110	aad := CreateFooterAad(d.fileAad)
111	return d.getFooterDecryptor(aad, true)
112}
113
114func (d *fileDecryptor) GetFooterDecryptorForColumnMeta(aad string) Decryptor {
115	return d.getFooterDecryptor(aad, true)
116}
117
118func (d *fileDecryptor) GetFooterDecryptorForColumnData(aad string) Decryptor {
119	return d.getFooterDecryptor(aad, false)
120}
121
122func (d *fileDecryptor) GetColumnMetaDecryptor(columnPath, columnKeyMetadata, aad string) Decryptor {
123	return d.getColumnDecryptor(columnPath, columnKeyMetadata, aad, true)
124}
125
126func (d *fileDecryptor) GetColumnDataDecryptor(columnPath, columnKeyMetadata, aad string) Decryptor {
127	return d.getColumnDecryptor(columnPath, columnKeyMetadata, aad, false)
128}
129
130func (d *fileDecryptor) WipeOutDecryptionKeys() {
131	d.props.WipeOutDecryptionKeys()
132}
133
134func (d *fileDecryptor) getFooterDecryptor(aad string, metadata bool) Decryptor {
135	if metadata {
136		if d.footerMetadataDecryptor != nil {
137			return d.footerMetadataDecryptor
138		}
139	} else {
140		if d.footerDataDecryptor != nil {
141			return d.footerDataDecryptor
142		}
143	}
144
145	footerKey := d.GetFooterKey()
146
147	// Create both data and metadata decryptors to avoid redundant retrieval of key
148	// from the key_retriever.
149	aesMetaDecrypt := d.getMetaAesDecryptor()
150	aesDataDecrypt := d.getDataAesDecryptor()
151
152	d.footerMetadataDecryptor = &decryptor{
153		decryptor: aesMetaDecrypt,
154		key:       []byte(footerKey),
155		fileAad:   []byte(d.fileAad),
156		aad:       []byte(aad),
157		mem:       d.mem,
158	}
159	d.footerDataDecryptor = &decryptor{
160		decryptor: aesDataDecrypt,
161		key:       []byte(footerKey),
162		fileAad:   []byte(d.fileAad),
163		aad:       []byte(aad),
164		mem:       d.mem,
165	}
166
167	if metadata {
168		return d.footerMetadataDecryptor
169	}
170	return d.footerDataDecryptor
171}
172
173func (d *fileDecryptor) getColumnDecryptor(columnPath, columnMeta, aad string, metadata bool) Decryptor {
174	if metadata {
175		if res, ok := d.columnMetaDataMap[columnPath]; ok {
176			res.UpdateAad(aad)
177			return res
178		}
179	} else {
180		if res, ok := d.columnDataMap[columnPath]; ok {
181			res.UpdateAad(aad)
182			return res
183		}
184	}
185
186	columnKey := d.props.ColumnKey(columnPath)
187	// No explicit column key given via API. Retrieve via key metadata.
188	if columnKey == "" && columnMeta != "" && d.props.KeyRetriever != nil {
189		columnKey = d.props.KeyRetriever.GetKey([]byte(columnMeta))
190	}
191	if columnKey == "" {
192		panic("hidden column exception, path=" + columnPath)
193	}
194
195	aesDataDecrypt := d.getDataAesDecryptor()
196	aesMetaDecrypt := d.getMetaAesDecryptor()
197
198	d.columnDataMap[columnPath] = &decryptor{
199		decryptor: aesDataDecrypt,
200		key:       []byte(columnKey),
201		fileAad:   []byte(d.fileAad),
202		aad:       []byte(aad),
203		mem:       d.mem,
204	}
205	d.columnMetaDataMap[columnPath] = &decryptor{
206		decryptor: aesMetaDecrypt,
207		key:       []byte(columnKey),
208		fileAad:   []byte(d.fileAad),
209		aad:       []byte(aad),
210		mem:       d.mem,
211	}
212
213	if metadata {
214		return d.columnMetaDataMap[columnPath]
215	}
216	return d.columnDataMap[columnPath]
217}
218
219func (d *fileDecryptor) getMetaAesDecryptor() *aesDecryptor {
220	if d.metaDecryptor == nil {
221		d.metaDecryptor = newAesDecryptor(d.alg, true)
222	}
223	return d.metaDecryptor
224}
225
226func (d *fileDecryptor) getDataAesDecryptor() *aesDecryptor {
227	if d.dataDecryptor == nil {
228		d.dataDecryptor = newAesDecryptor(d.alg, false)
229	}
230	return d.dataDecryptor
231}
232
233// Decryptor is the basic interface for any decryptor generated from a FileDecryptor
234type Decryptor interface {
235	// returns the File Level AAD bytes
236	FileAad() string
237	// returns the current allocator that was used for any extra allocations of buffers
238	Allocator() memory.Allocator
239	// returns the CiphertextSizeDelta from the decryptor
240	CiphertextSizeDelta() int
241	// Decrypt just returns the decrypted plaintext from the src ciphertext
242	Decrypt(src []byte) []byte
243	// set the AAD bytes of the decryptor to the provided string
244	UpdateAad(string)
245}
246
247type decryptor struct {
248	decryptor *aesDecryptor
249	key       []byte
250	fileAad   []byte
251	aad       []byte
252	mem       memory.Allocator
253}
254
255func (d *decryptor) Allocator() memory.Allocator { return d.mem }
256func (d *decryptor) FileAad() string             { return string(d.fileAad) }
257func (d *decryptor) UpdateAad(aad string)        { d.aad = []byte(aad) }
258func (d *decryptor) CiphertextSizeDelta() int    { return d.decryptor.CiphertextSizeDelta() }
259func (d *decryptor) Decrypt(src []byte) []byte {
260	return d.decryptor.Decrypt(src, d.key, d.aad)
261}
262