RexXiong commented on PR #3689:
URL: https://github.com/apache/celeborn/pull/3689#issuecomment-4478748117
## PR Review: [CELEBORN-2329][CIP22] Encryption at Rest Spark Impl
**Overall: Good architecture, some performance and safety concerns.**
The encrypt-at-client / decrypt-at-client approach is clean — workers store
ciphertext, no key distribution needed to workers. Compress-then-encrypt
ordering is correct.
---
### Critical Issues
**1. Unbounded allocation in `decrypt` — OOM vector**
```java
int decryptedLength = dis.readInt();
byte[] decrypted = new byte[decryptedLength];
```
If the encrypted data is corrupted (bit flip, wrong key, truncation),
`decryptedLength` could be any value up to `Integer.MAX_VALUE`. The check `< 0`
catches negative values but not unreasonably large ones (e.g., 2GB). Should add
an upper bound check:
```java
if (decryptedLength < 0 || decryptedLength > MAX_EXPECTED_BATCH_SIZE) {
throw new IOException("Invalid decrypted length: " + decryptedLength);
}
```
---
### Performance Concerns
**2. Heavy allocation on hot path — every batch allocates multiple objects**
`SparkCryptoHandler.encrypt()` and `decrypt()` each create:
- ByteArrayOutputStream/InputStream
- DataOutputStream/InputStream
- CryptoStream (which internally allocates IV, cipher state)
- Output byte[] array
For a typical shuffle with thousands of batches per second per executor,
this generates significant GC pressure. Consider:
- Reusing crypto streams via `ThreadLocal` (reset with new IV per call)
- Using a buffer pool for the output arrays
- Or at minimum, pre-sizing `ByteArrayOutputStream` with `length + overhead`
to avoid internal resizing
**3. Extra copy in write path**
```java
byte[] encrypted = cryptoHandler.get().encrypt(data, offset, length);
data = encrypted;
// then immediately copied into: new byte[BATCH_HEADER_SIZE + length]
```
The encrypted data is allocated, then immediately copied into the batch
body. Two allocations + one copy for every batch. Could be optimized by having
encrypt write directly into the batch body buffer at an offset.
---
### Moderate Issues
**4. `commons-crypto` dependency added to generic client module**
Non-Spark engines (Flink, MR) using the Celeborn client now pull in
`commons-crypto` even if they never use encryption. Consider making this
dependency `optional` in the client pom and keeping it non-optional only in the
spark-shaded modules.
**5. `getCryptoHandler` called multiple times per reader/writer creation**
```java
SparkCommonUtils.getCryptoHandler(conf) // called in getWriter
SparkCommonUtils.getCryptoHandler(conf) // called in getReader (multiple
times for columnar path)
```
Each call accesses `SparkEnv.get()` and
`securityManager().getIOEncryptionKey()`. While cheap, it creates a new
`SparkCryptoHandler` instance each time. Since the key doesn't change within an
app, this should be cached (e.g., at the `SparkShuffleManager` level).
---
### Minor / Nits
**6. No AES-GCM / integrity protection**
Spark's `CryptoStreamUtils` uses AES-CTR by default which provides
confidentiality but not integrity. A compromised worker could tamper with
ciphertext without detection. This is a pre-existing Spark limitation (not
introduced by this PR) but worth documenting in the CIP as a known gap.
**7. Test coverage is good** — roundtrip, wrong key, large data, empty data,
offset. Could add a test for corrupted ciphertext (flip a byte in encrypted
output, verify decrypt fails or produces wrong output).
**8. Debug logging on every batch**
```java
logger.debug("Encrypted shuffle data for shuffle {} ...", shuffleId, ...);
logger.debug("Decrypted shuffle data for shuffle {} ...", shuffleId, ...);
```
Even with debug disabled, the varargs boxing and string formatting
preparation happens. On a hot path with millions of batches, this adds up.
Consider guarding with `if (logger.isDebugEnabled())`.
---
*Reviewed with Claude Code*
--
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]
