tanmayrauth commented on code in PR #1620: URL: https://github.com/apache/iceberg-go/pull/1620#discussion_r3693021679
########## encryption/standard_manager.go: ########## @@ -0,0 +1,488 @@ +// 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. + +package encryption + +import ( + "context" + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + + icebergio "github.com/apache/iceberg-go/io" +) + +// Defaults for [StandardEncryptionManager]. +const ( + // StandardDefaultDEKLength is the default length, in bytes, of the + // per-file data encryption key (DEK) generated for AES-256-GCM. + StandardDefaultDEKLength = 32 + + // StandardDefaultBlockSize is the default plaintext block size, in + // bytes, used to split a file into independently authenticated AES-GCM + // blocks. Blocks allow random access (Seek/ReadAt) without buffering or + // decrypting the whole file. + StandardDefaultBlockSize = 64 * 1024 +) + +// Sentinel errors returned by [StandardEncryptionManager]. +var ( + // ErrKeyIDRequired is returned by + // [StandardEncryptionManager.NewEncryptedOutputFile] when keyID is empty. + // StandardEncryptionManager always encrypts, so it requires a KEK to wrap + // the generated DEK; use [PlaintextEncryptionManager] for unencrypted + // tables instead of passing an empty keyID here. + ErrKeyIDRequired = errors.New("encryption: StandardEncryptionManager requires a non-empty keyID") + + // ErrKeyMetadataRequired is returned by + // [StandardEncryptionManager.NewDecryptedInputFile] when keyMetadata is + // empty. StandardEncryptionManager always decrypts, so it requires the + // per-file key metadata produced by [StandardEncryptionManager.NewEncryptedOutputFile]. + ErrKeyMetadataRequired = errors.New("encryption: StandardEncryptionManager requires non-empty key metadata") + + // ErrUnsupportedKeyMetadataVersion is returned when key metadata was + // produced by a newer, incompatible encoding version. + ErrUnsupportedKeyMetadataVersion = errors.New("encryption: unsupported key metadata version") +) + +// standardKeyMetadataVersion is the current encoding version written by +// [StandardEncryptionManager]. It is bumped whenever the on-disk layout of +// standardKeyMetadata or the block ciphertext format changes incompatibly. +const standardKeyMetadataVersion = 1 + +// standardKeyMetadata is the JSON-encoded structure stored as the opaque +// [EncryptionKeyMetadata] for files produced by [StandardEncryptionManager]. +type standardKeyMetadata struct { + Version int `json:"v"` + KeyID string `json:"key-id"` + WrappedKey []byte `json:"wrapped-key"` + NoncePrefix []byte `json:"nonce-prefix"` + BlockSize int `json:"block-size"` + PlaintextLength int64 `json:"plaintext-length"` +} + +// StandardEncryptionManager is a generic, format-agnostic [EncryptionManager] +// that provides envelope encryption for arbitrary files (e.g. manifests, +// manifest lists, Puffin statistics) using a [KeyManagementClient] to wrap +// and unwrap a fresh AES-256-GCM data encryption key (DEK) per file. +// +// Each file is split into fixed-size plaintext blocks, and each block is +// sealed independently with AES-GCM using a unique nonce (a per-file random +// prefix combined with the block index). This bounds memory usage and +// supports random access (Seek/ReadAt) on the decrypted file without +// buffering or decrypting more than the requested blocks. +// +// StandardEncryptionManager always encrypts and always decrypts: it fails +// closed, returning [ErrKeyIDRequired] or [ErrKeyMetadataRequired] rather +// than silently falling back to plaintext. Use [PlaintextEncryptionManager] +// for tables or files that are not encrypted. +type StandardEncryptionManager struct { + kms KeyManagementClient + dekLength int + blockSize int +} + +var _ EncryptionManager = (*StandardEncryptionManager)(nil) + +// StandardManagerOption configures a [StandardEncryptionManager] created by +// [NewStandardEncryptionManager]. +type StandardManagerOption func(*StandardEncryptionManager) + +// WithDEKLength overrides the default data encryption key length (in bytes). +// Valid AES key lengths are 16, 24, or 32 bytes. +func WithDEKLength(length int) StandardManagerOption { + return func(m *StandardEncryptionManager) { m.dekLength = length } +} + +// WithBlockSize overrides the default plaintext block size (in bytes) used +// to split files for independent block-level authentication. +func WithBlockSize(size int) StandardManagerOption { + return func(m *StandardEncryptionManager) { m.blockSize = size } +} + +// NewStandardEncryptionManager creates a [StandardEncryptionManager] backed +// by kms. kms must not be nil. +func NewStandardEncryptionManager(kms KeyManagementClient, opts ...StandardManagerOption) *StandardEncryptionManager { + m := &StandardEncryptionManager{ + kms: kms, + dekLength: StandardDefaultDEKLength, + blockSize: StandardDefaultBlockSize, + } + for _, opt := range opts { + opt(m) + } + + return m +} + +// NewEncryptedOutputFile creates a new AES-GCM block-encrypted output file. +// keyID identifies the KEK used to wrap the freshly generated per-file DEK, +// and must be non-empty; otherwise [ErrKeyIDRequired] is returned. +func (m *StandardEncryptionManager) NewEncryptedOutputFile(ctx context.Context, writer icebergio.FileWriter, keyID string) (EncryptedOutputFile, error) { + if keyID == "" { + return nil, ErrKeyIDRequired + } + + var ( + plainDEK, wrappedDEK []byte + err error + ) + if m.kms.SupportsKeyGeneration() { + plainDEK, wrappedDEK, err = m.kms.GenerateKey(ctx, keyID, m.dekLength) + if err != nil { + return nil, fmt.Errorf("encryption: failed to generate DEK: %w", err) + } + } else { + plainDEK = make([]byte, m.dekLength) + if _, err = rand.Read(plainDEK); err != nil { + return nil, fmt.Errorf("encryption: failed to generate DEK: %w", err) + } + if wrappedDEK, err = m.kms.WrapKey(ctx, keyID, plainDEK); err != nil { + return nil, fmt.Errorf("encryption: failed to wrap DEK: %w", err) + } + } + + aead, err := newStandardAEAD(plainDEK) + if err != nil { + return nil, err + } + + noncePrefix := make([]byte, 4) + if _, err := rand.Read(noncePrefix); err != nil { + return nil, fmt.Errorf("encryption: failed to generate nonce prefix: %w", err) + } + + return &standardOutputFile{ + FileWriter: writer, + aead: aead, + noncePrefix: noncePrefix, + blockSize: m.blockSize, + keyID: keyID, + wrappedKey: wrappedDEK, + }, nil +} + +// NewDecryptedInputFile wraps file for transparent block-level AES-GCM +// decryption. keyMetadata must be the non-empty blob produced by +// [StandardEncryptionManager.NewEncryptedOutputFile]; otherwise +// [ErrKeyMetadataRequired] is returned. +func (m *StandardEncryptionManager) NewDecryptedInputFile(ctx context.Context, file icebergio.File, keyMetadata EncryptionKeyMetadata) (EncryptedInputFile, error) { + if len(keyMetadata) == 0 { + return nil, ErrKeyMetadataRequired + } + + var meta standardKeyMetadata + if err := json.Unmarshal(keyMetadata, &meta); err != nil { + return nil, fmt.Errorf("encryption: failed to decode key metadata: %w", err) + } + if meta.Version != standardKeyMetadataVersion { + return nil, fmt.Errorf("%w: %d", ErrUnsupportedKeyMetadataVersion, meta.Version) + } + + plainDEK, err := m.kms.UnwrapKey(ctx, meta.KeyID, meta.WrappedKey) + if err != nil { + return nil, fmt.Errorf("encryption: failed to unwrap DEK: %w", err) + } + + aead, err := newStandardAEAD(plainDEK) + if err != nil { + return nil, err + } + + return &standardInputFile{ + underlying: file, + aead: aead, + noncePrefix: meta.NoncePrefix, + blockSize: meta.BlockSize, + plaintextLength: meta.PlaintextLength, + keyMetadata: keyMetadata, + }, nil +} + +func newStandardAEAD(key []byte) (cipher.AEAD, error) { + block, err := aes.NewCipher(key) + if err != nil { + return nil, fmt.Errorf("%w: %v", ErrInvalidKeyLength, err) + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, fmt.Errorf("encryption: failed to create GCM: %w", err) + } + + return gcm, nil +} + +// standardBlockNonce derives the AES-GCM nonce for blockIndex: the 4-byte +// per-file random prefix followed by the 8-byte big-endian block index. The +// (key, nonce) pair is unique per block since the DEK is fresh per file and +// no two blocks in the same file share an index. +func standardBlockNonce(prefix []byte, blockIndex uint64) []byte { + nonce := make([]byte, 12) + copy(nonce, prefix) + binary.BigEndian.PutUint64(nonce[4:], blockIndex) + + return nonce +} + +// standardOutputFile is an [EncryptedOutputFile] that seals fixed-size +// plaintext blocks with AES-GCM as they are written. +type standardOutputFile struct { + icebergio.FileWriter + + aead cipher.AEAD + noncePrefix []byte + blockSize int + keyID string + wrappedKey []byte + + buf []byte + blockIndex uint64 + written int64 + closed bool + + keyMetadata EncryptionKeyMetadata +} + +var _ EncryptedOutputFile = (*standardOutputFile)(nil) + +func (f *standardOutputFile) Write(p []byte) (int, error) { + total := len(p) + for len(p) > 0 { + space := f.blockSize - len(f.buf) Review Comment: Nit: NewStandardEncryptionManager doesn't validate blockSize, so a bad WithBlockSize only surfaces at write time — WithBlockSize(-1) panics slice bounds out of range [:-1] on the first Write, and WithBlockSize(0) spins forever here (len(f.buf) == f.blockSize is 0 == 0 every iteration, so p is never consumed). A blockSize > 0 guard in NewEncryptedOutputFile (which already returns an error) rejects it early. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
