hectar-glitches commented on code in PR #1620: URL: https://github.com/apache/iceberg-go/pull/1620#discussion_r4262469935
########## encryption/standard_manager.go: ########## @@ -0,0 +1,862 @@ +// 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" + "errors" + "fmt" + "io" + "io/fs" + "math" + "slices" + "sync" + + icebergio "github.com/apache/iceberg-go/io" +) + +// StandardDefaultDataKeyLength is the default length, in bytes, of the +// per-file data encryption key (DEK). It is 128-bit, matching the default of +// the encryption.data-key-length table property in the Java and Rust +// implementations. +const StandardDefaultDataKeyLength = 16 + +// Sentinel errors returned by [StandardEncryptionManager]. +var ( + // ErrKeyIDRequired is returned by + // [StandardEncryptionManager.NewEncryptedOutputFile] when keyID is empty. + // StandardEncryptionManager always encrypts, so it requires the table's + // encryption key ID; 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 starts + // with a version byte other than the one this package understands. + ErrUnsupportedKeyMetadataVersion = errors.New("encryption: unsupported key metadata version") + + // ErrInvalidStreamHeader is returned when the AES GCM Stream header (the + // "AGS1" magic and little-endian plain block length at the start of an + // encrypted file, per the Iceberg AES GCM Stream spec) is missing, + // truncated, or declares a block length other than the fixed 1 MiB. + ErrInvalidStreamHeader = errors.New("encryption: invalid AES GCM Stream header") + + // ErrInvalidKeyMetadata is returned by + // [StandardEncryptionManager.NewDecryptedInputFile] when key metadata is + // not a well-formed StandardKeyMetadata record or fails basic sanity + // checks (e.g. a truncated record, a missing or wrong-sized AAD prefix, or + // a missing or impossible file length). Key metadata is untrusted input on + // a crypto read path, so it is validated rather than trusted blindly. + ErrInvalidKeyMetadata = errors.New("encryption: invalid key metadata") + + // ErrOutputFileClosed is returned by [standardOutputFile.Write] when + // called after a successful Close. It wraps [fs.ErrClosed] so callers can + // test with errors.Is(err, fs.ErrClosed). A Write after a failed flush or + // a failed Close returns that original error instead. + ErrOutputFileClosed = fmt.Errorf("encryption: write to closed StandardEncryptionManager output file: %w", fs.ErrClosed) + + // ErrBlockTruncated is returned by [standardInputFile.ReadAt] when the + // underlying storage returns fewer ciphertext bytes for a block than its + // recorded length requires, indicating the file was truncated at rest. + // This is distinct from [ErrCiphertextTooShort], which [KeyManagementClient] + // implementations use for a too-short wrapped key or KMS-encrypted + // payload; keeping them separate lets a caller tell a malformed KMS blob + // apart from a short block read. + ErrBlockTruncated = errors.New("encryption: block truncated: read fewer ciphertext bytes than expected") +) + +// Constants describing the Iceberg AES GCM Stream ("AGS1") wire format used +// for the ciphertext produced by [StandardEncryptionManager]. See +// https://iceberg.apache.org/gcm-stream-spec/ for the full specification. +const ( + // gcmStreamMagic identifies an AES GCM Stream version 1 file. + gcmStreamMagic = "AGS1" + + // gcmStreamHeaderLength is the length, in bytes, of the magic plus the + // 4-byte little-endian plain block length written at the start of every + // file. + gcmStreamHeaderLength = 8 + + // gcmStreamPlainBlockSize is the plaintext block size. It is fixed: + // Java's Ciphers.PLAIN_BLOCK_SIZE and iceberg-rust's PLAIN_BLOCK_SIZE are + // both 1 MiB, and Java's AesGcmInputStream rejects any other value in the + // header, so a stream with a different block size is unreadable there. + gcmStreamPlainBlockSize = 1024 * 1024 + + // gcmStreamNonceLength is the length, in bytes, of the random AES-GCM + // nonce stored at the start of every cipher block. + gcmStreamNonceLength = 12 + + // gcmStreamTagLength is the length, in bytes, of the AES-GCM + // authentication tag appended to every cipher block's ciphertext. + gcmStreamTagLength = 16 + + // gcmStreamBlockOverhead is the number of ciphertext bytes added to + // each block beyond its plaintext length (nonce + tag). + gcmStreamBlockOverhead = gcmStreamNonceLength + gcmStreamTagLength + + // gcmStreamCipherBlockSize is the on-disk length of a full block. + gcmStreamCipherBlockSize = gcmStreamPlainBlockSize + gcmStreamBlockOverhead + + // gcmStreamMinLength is the length of the smallest valid stream: the + // header followed by one block holding no plaintext. Java and iceberg-rust + // always write block 0, even for an empty file, and iceberg-rust rejects + // anything shorter. + gcmStreamMinLength = gcmStreamHeaderLength + gcmStreamBlockOverhead + + // gcmStreamAADPrefixLength is the length, in bytes, of the random + // per-file AAD prefix. It matches Java's ENCRYPTION_AAD_LENGTH_DEFAULT. + gcmStreamAADPrefixLength = 16 +) + +// standardKeyMetadataVersion is the leading version byte of Iceberg's +// StandardKeyMetadata encoding. +const standardKeyMetadataVersion = 1 + +// standardKeyMetadata is the decoded form of Iceberg's StandardKeyMetadata: +// a version byte followed by an Avro binary record +// +// {encryption_key: bytes, aad_prefix: ["null", bytes], file_length: ["null", long]} +// +// This is the key_metadata encoding Java's StandardEncryptionManager, +// iceberg-rust and pyiceberg read and write. It holds the raw DEK, so it is a +// secret and must be protected by whatever stores it. +type standardKeyMetadata struct { + encryptionKey []byte + + // aadPrefix is combined with each block's little-endian index to form + // the AES GCM Stream additional authenticated data, binding every + // ciphertext block to this file and to its position so that blocks + // cannot be silently reordered, replayed from another file, or spliced + // in from elsewhere in the same file. It is not secret. It is nil when + // the record omits it. + aadPrefix []byte + + // fileLength is the length of the encrypted stream, header included. + // Per the AES GCM Stream spec's "File length" note, a reader must take + // the length from a trusted source rather than the underlying storage's + // reported size, since storage size alone cannot distinguish a genuinely + // short file from one truncated by an attacker who does not also control + // this metadata. hasFileLength is false when the record omits it. + fileLength int64 + hasFileLength bool +} + +// encodeStandardKeyMetadata encodes the version byte and Avro record for a +// file, with both optional fields present. +func encodeStandardKeyMetadata(encryptionKey, aadPrefix []byte, fileLength int64) EncryptionKeyMetadata { + b := []byte{standardKeyMetadataVersion} + b = binary.AppendVarint(b, int64(len(encryptionKey))) + b = append(b, encryptionKey...) + b = binary.AppendVarint(b, 1) // union branch 1 (bytes): aad_prefix is present + b = binary.AppendVarint(b, int64(len(aadPrefix))) + b = append(b, aadPrefix...) + b = binary.AppendVarint(b, 1) // union branch 1 (long): file_length is present + b = binary.AppendVarint(b, fileLength) + + return b +} + +// decodeStandardKeyMetadata decodes the version byte and Avro record written +// by [encodeStandardKeyMetadata] or by Java's StandardKeyMetadata.buffer(). +func decodeStandardKeyMetadata(b []byte) (standardKeyMetadata, error) { + if b[0] != standardKeyMetadataVersion { + return standardKeyMetadata{}, fmt.Errorf("%w: %d", ErrUnsupportedKeyMetadataVersion, b[0]) + } + + r := avroReader{buf: b[1:]} + var k standardKeyMetadata + k.encryptionKey = r.bytes() + if r.present() { + k.aadPrefix = r.bytes() + } + if r.present() { + k.fileLength = r.long() + k.hasFileLength = true + } + if r.err == nil && len(r.buf) != 0 { + r.err = fmt.Errorf("%d unexpected trailing bytes", len(r.buf)) + } + if r.err != nil { + return standardKeyMetadata{}, fmt.Errorf("%w: %w", ErrInvalidKeyMetadata, r.err) + } + + return k, nil +} + +// avroReader decodes the few Avro binary primitives StandardKeyMetadata uses, +// remembering the first error. Avro longs are zig-zag varints, which is what +// encoding/binary's Varint implements. +type avroReader struct { + buf []byte + err error +} + +func (r *avroReader) long() int64 { + if r.err != nil { + return 0 + } + v, n := binary.Varint(r.buf) + if n <= 0 { + r.err = errors.New("truncated or oversized varint") + + return 0 + } + r.buf = r.buf[n:] + + return v +} + +func (r *avroReader) bytes() []byte { + n := r.long() + if r.err != nil { + return nil + } + if n < 0 || n > int64(len(r.buf)) { + r.err = fmt.Errorf("bytes length %d out of range (%d bytes remain)", n, len(r.buf)) + + return nil + } + out := slices.Clone(r.buf[:n]) + r.buf = r.buf[n:] + + return out +} + +// present reads the branch index of a ["null", T] union and reports whether +// the value is present (branch 1) rather than null (branch 0). +func (r *avroReader) present() bool { + branch := r.long() + if r.err != nil { + return false + } + if branch != 0 && branch != 1 { + r.err = fmt.Errorf("unexpected union branch %d", branch) + + return false + } + + return branch == 1 +} + +// plaintextLengthFromStreamLength derives the plaintext length from the +// length of the encrypted stream, the way Java's +// AesGcmInputStream.calculatePlaintextLength does. It accepts the 36-byte +// form of an empty file (header plus one empty block) and rejects lengths no +// writer could have produced. +func plaintextLengthFromStreamLength(streamLength int64) (int64, error) { + if streamLength < gcmStreamMinLength { + return 0, fmt.Errorf("%w: file-length %d is shorter than the %d-byte minimum stream", ErrInvalidKeyMetadata, streamLength, gcmStreamMinLength) + } + + body := streamLength - gcmStreamHeaderLength + plainLength := body / gcmStreamCipherBlockSize * gcmStreamPlainBlockSize + if lastBlock := body % gcmStreamCipherBlockSize; lastBlock != 0 { + if lastBlock < gcmStreamBlockOverhead { + return 0, fmt.Errorf("%w: file-length %d ends in a %d-byte block, shorter than its %d-byte nonce and tag", ErrInvalidKeyMetadata, streamLength, lastBlock, gcmStreamBlockOverhead) + } + plainLength += lastBlock - gcmStreamBlockOverhead + } + + return plainLength, nil +} + +// StandardEncryptionManager is a generic, format-agnostic [EncryptionManager] +// for arbitrary files (e.g. manifests, manifest lists, Puffin statistics). It +// reads and writes the same bytes as Java's StandardEncryptionManager, +// iceberg-rust and pyiceberg: an Iceberg AES GCM Stream ("AGS1") with fixed +// 1 MiB plaintext blocks, and key metadata in Iceberg's StandardKeyMetadata +// encoding. +// +// Each file is encrypted under a fresh random data encryption key (DEK), +// generated locally with no KMS call. Each block carries its own random nonce +// and is authenticated with an AAD that binds it to the file and to its +// position. This bounds memory usage and supports random access +// (Seek/ReadAt) on the decrypted file without buffering or decrypting more +// than the requested blocks. +// +// The key metadata returned by Close contains the raw DEK, so it is a secret: +// persist it only where something protects it, such as a manifest or manifest +// list that is itself encrypted under a key encryption key. Wrapping the DEK +// that way is not done by this type. +// +// 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 { + dataKeyLength int +} + +var _ EncryptionManager = (*StandardEncryptionManager)(nil) + +// StandardManagerOption configures a [StandardEncryptionManager] created by +// [NewStandardEncryptionManager]. +type StandardManagerOption func(*StandardEncryptionManager) + +// WithStandardDataKeyLength sets the length, in bytes, of the per-file data +// encryption key. It corresponds to the encryption.data-key-length table +// property. Valid AES key lengths are 16, 24, or 32 bytes; the default is +// [StandardDefaultDataKeyLength]. +func WithStandardDataKeyLength(length int) StandardManagerOption { + return func(m *StandardEncryptionManager) { m.dataKeyLength = length } +} + +// NewStandardEncryptionManager creates a [StandardEncryptionManager]. +func NewStandardEncryptionManager(opts ...StandardManagerOption) *StandardEncryptionManager { + m := &StandardEncryptionManager{dataKeyLength: StandardDefaultDataKeyLength} + for _, opt := range opts { + opt(m) + } + + return m +} + +// NewEncryptedOutputFile creates a new AES-GCM block-encrypted output file. +// keyID is the table's encryption key ID and must be non-empty; otherwise +// [ErrKeyIDRequired] is returned. It is not recorded in the file's key +// metadata. +func (m *StandardEncryptionManager) NewEncryptedOutputFile(_ context.Context, writer icebergio.FileWriter, keyID string) (EncryptedOutputFile, error) { + if keyID == "" { + return nil, ErrKeyIDRequired + } + switch m.dataKeyLength { + case 16, 24, 32: + default: + return nil, fmt.Errorf("%w: DEK length must be 16, 24, or 32 bytes; got %d", ErrInvalidKeyLength, m.dataKeyLength) + } + + // A fresh DEK per file keeps the number of random-nonce encryptions under + // any one key within the 2^32 limit NIST SP 800-38D sets for random + // 96-bit nonces: a file holds at most 2^32-1 blocks. + dek := make([]byte, m.dataKeyLength) + if _, err := io.ReadFull(rand.Reader, dek); err != nil { + return nil, fmt.Errorf("encryption: failed to generate DEK: %w", err) + } + + aead, err := newStandardAEAD(dek) + if err != nil { + return nil, err + } + + aadPrefix := make([]byte, gcmStreamAADPrefixLength) + if _, err := io.ReadFull(rand.Reader, aadPrefix); err != nil { + return nil, fmt.Errorf("encryption: failed to generate AAD prefix: %w", err) + } + + // Write the AES GCM Stream header (magic + little-endian block length) + // up front, before any ciphertext blocks, per the format spec. + header := make([]byte, gcmStreamHeaderLength) + copy(header, gcmStreamMagic) + binary.LittleEndian.PutUint32(header[len(gcmStreamMagic):], gcmStreamPlainBlockSize) + if _, err := writer.Write(header); err != nil { + return nil, fmt.Errorf("encryption: failed to write stream header: %w", err) + } + + return &standardOutputFile{ + FileWriter: writer, + aead: aead, + dek: dek, + aadPrefix: aadPrefix, + encryptedLength: gcmStreamHeaderLength, + }, nil +} + +// NewDecryptedInputFile wraps file for transparent block-level AES-GCM +// decryption. keyMetadata must be the non-empty StandardKeyMetadata blob +// produced by [StandardEncryptionManager.NewEncryptedOutputFile] or by another +// Iceberg implementation; otherwise [ErrKeyMetadataRequired] is returned. +// +// The plaintext length is derived from the file length recorded in key +// metadata, never from the underlying file's size. Key metadata without a file +// length is rejected, because there is no other trusted source for it here. +// Key metadata without an AAD prefix, which the encoding permits, is rejected +// too: the prefix is what binds a block to its file. +func (m *StandardEncryptionManager) NewDecryptedInputFile(_ context.Context, file icebergio.File, keyMetadata EncryptionKeyMetadata) (EncryptedInputFile, error) { + if len(keyMetadata) == 0 { + return nil, ErrKeyMetadataRequired + } + + meta, err := decodeStandardKeyMetadata(keyMetadata) + if err != nil { + return nil, err + } + // Java, iceberg-rust and this writer all emit 16-byte prefixes. Bounding + // it also keeps the per-block AAD allocation in readBlock small. + if len(meta.aadPrefix) != gcmStreamAADPrefixLength { + return nil, fmt.Errorf("%w: aad-prefix must be exactly %d bytes, got %d", ErrInvalidKeyMetadata, gcmStreamAADPrefixLength, len(meta.aadPrefix)) + } + if !meta.hasFileLength { Review Comment: NewDecryptedInputFileWithLength (standard_manager.go:412) uses file_length when it’s present, otherwise the caller’s length, and fails only when neither exists. -- 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]
