feat: streaming framed part encryption (O(frame) UploadPart memory) (#79)

Author: ServerSideHannes

e2e/postgres/templates/postgres-cluster.yaml

@@ -36,5 +36,11 @@ spec:
       wal:
         compression: gzip
       data:
-        compression: gzip
+        # Compression intentionally OFF: pgbench data is highly compressible, so
+        # gzip would shrink the backup below one 512MB chunk and we'd never get
+        # concurrent large uploads. Uncompressed, the dataset spans several 512MB
+        # chunks which (with jobs=4) reproduces the multipart-upload OOM.
+        jobs: 4
+        additionalCommandArgs:
+          - "--min-chunk-size=512MB"
     retentionPolicy: "7d"

e2e/postgres/test.sh

@@ -20,7 +20,7 @@ fi
 
 NAMESPACE="postgres-test"
 export CLUSTER_NAME="pg-cluster"
-DATA_SIZE_GB=2
+DATA_SIZE_GB=3
 SCALE_FACTOR=$((DATA_SIZE_GB * 70))  # pgbench scale: ~15MB per scale factor
 
 # Colors for output
@@ -33,6 +33,44 @@ log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
 log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
 log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
 
+S3_ENDPOINT="http://s3proxy-python-frontproxy.s3proxy:80"
+S3_WAL_CHECK_POD="s3-wal-check"
+
+ensure_s3_wal_checker() {
+    if kubectl get pod -n "$NAMESPACE" "$S3_WAL_CHECK_POD" >/dev/null 2>&1; then
+        kubectl wait -n "$NAMESPACE" --for=condition=Ready "pod/${S3_WAL_CHECK_POD}" --timeout=120s
+        return
+    fi
+    kubectl run "$S3_WAL_CHECK_POD" --restart=Never -n "$NAMESPACE" \
+        --image=amazon/aws-cli:2.15.0 \
+        --overrides='{"spec":{"containers":[{"name":"'"$S3_WAL_CHECK_POD"'","image":"amazon/aws-cli:2.15.0","command":["sleep","3600"],"envFrom":[{"secretRef":{"name":"s3-credentials"}}]}]}}'
+    kubectl wait -n "$NAMESPACE" --for=condition=Ready "pod/${S3_WAL_CHECK_POD}" --timeout=120s
+}
+
+wait_for_end_wal_in_s3() {
+    local end_wal="$1"
+    local timeline_prefix="${end_wal:0:16}"
+    local wal_object="pg-cluster/wals/${timeline_prefix}/${end_wal}.gz"
+
+    log_info "Waiting for backup end WAL in S3: ${wal_object}"
+    ensure_s3_wal_checker
+
+    local deadline=$((SECONDS + 600))
+    while [ "$SECONDS" -lt "$deadline" ]; do
+        if kubectl exec -n "$NAMESPACE" "$S3_WAL_CHECK_POD" -- sh -c "
+            export AWS_ACCESS_KEY_ID=\$ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY=\$ACCESS_SECRET_KEY
+            aws --endpoint-url ${S3_ENDPOINT} s3 ls s3://postgres-backups/${wal_object} >/dev/null 2>&1
+        "; then
+            log_info "✓ End WAL archived: ${end_wal}.gz"
+            return 0
+        fi
+        sleep 5
+    done
+
+    log_error "Timeout waiting for end WAL ${end_wal}.gz in S3"
+    return 1
+}
+
 cleanup() {
     log_info "Cleaning up..."
     kubectl delete namespace "$NAMESPACE" --ignore-not-found --wait=false || true
@@ -218,41 +256,63 @@ log_info "Backup completed!"
 kubectl get backup -n "$NAMESPACE" ${CLUSTER_NAME}-backup-1 -o yaml | grep -A5 "status:"
 
 # ============================================================================
-# STEP 5: Verify encryption + Delete cluster + Create new cluster (ALL PARALLEL)
+# STEP 4b: Assert the s3proxy pods survived (the OOM we are gating against)
+# ============================================================================
+log_info "=== Checking s3proxy pods were not OOM-killed during backup ==="
+oom_found=0
+for p in $(kubectl get pods -n s3proxy -l app.kubernetes.io/name=s3proxy-python,app.kubernetes.io/component=server -o name); do
+    rc=$(kubectl get -n s3proxy "$p" -o jsonpath='{.status.containerStatuses[0].restartCount}' 2>/dev/null || echo 0)
+    reason=$(kubectl get -n s3proxy "$p" -o jsonpath='{.status.containerStatuses[0].lastState.terminated.reason}' 2>/dev/null || echo "")
+    log_info "  $p restarts=${rc:-0} lastTerminated=${reason:-none}"
+    if [ "${reason}" = "OOMKilled" ] || [ "${rc:-0}" -gt 0 ]; then oom_found=1; fi
+done
+if [ "$oom_found" -eq 1 ]; then
+    log_error "s3proxy was OOM-killed/restarted during backup"
+    exit 1
+fi
+log_info "✓ s3proxy survived the backup"
+
+# ============================================================================
+# STEP 4c: Wait for backup end WAL to reach S3 (restore needs it)
+# CNPG may mark backup completed before the archiver uploads the final segment.
 # ============================================================================
-log_info "=== Step 5: Parallel - verify encryption, delete old, create new ==="
+END_WAL=$(kubectl get backup -n "$NAMESPACE" "${CLUSTER_NAME}-backup-1" -o jsonpath='{.status.endWal}')
+if [ -z "$END_WAL" ]; then
+    log_error "Backup status.endWal is empty"
+    exit 1
+fi
+log_info "=== Step 4c: Waiting for end WAL ${END_WAL} in S3 ==="
+wait_for_end_wal_in_s3 "$END_WAL"
+
+# ============================================================================
+# STEP 5: Verify encryption, restore, then delete source cluster
+# Keep the source cluster alive until end WAL is archived and restore succeeds.
+# ============================================================================
+log_info "=== Step 5: Verify encryption and restore from backup ==="
 
-# 1. Start encryption verification in background
 verify_encryption "postgres-backups" "" "$NAMESPACE" ".gz|.tar|.backup|.data" &
 VERIFY_PID=$!
 
-# 2. Delete old cluster in background
-(
-    kubectl delete cluster -n "$NAMESPACE" ${CLUSTER_NAME} --wait
-    kubectl wait --namespace "$NAMESPACE" \
-        --for=delete pod -l cnpg.io/cluster=${CLUSTER_NAME} \
-        --timeout=300s || true
-    log_info "✓ Old cluster deleted"
-) &
-DELETE_PID=$!
-
-# 3. Create new cluster immediately (different name, can coexist)
-log_info "Creating restored cluster (parallel with deletion)..."
+log_info "Creating restored cluster..."
 envsubst < "${SCRIPT_DIR}/templates/postgres-cluster-restore.yaml" | kubectl apply -n "$NAMESPACE" -f -
 
-# Wait for all parallel operations
 wait $VERIFY_PID || { log_error "Encryption verification failed"; exit 1; }
 log_info "✓ Encryption verified"
 
-wait $DELETE_PID || { log_error "Old cluster deletion failed"; exit 1; }
-
 log_info "Waiting for restored cluster to be ready..."
 kubectl wait --namespace "$NAMESPACE" \
     --for=condition=Ready cluster/${CLUSTER_NAME}-restored \
     --timeout=1800s
 
 log_info "Restored cluster is ready!"
 
+log_info "Deleting source cluster (no longer needed)..."
+kubectl delete cluster -n "$NAMESPACE" "${CLUSTER_NAME}" --wait
+kubectl wait --namespace "$NAMESPACE" \
+    --for=delete pod -l "cnpg.io/cluster=${CLUSTER_NAME}" \
+    --timeout=300s || true
+log_info "✓ Old cluster deleted"
+
 # ============================================================================
 # STEP 6: Validate restored data
 # ============================================================================

s3proxy/client/s3.py

@@ -195,7 +195,8 @@ async def upload_part(
             Body=body,
         )
         duration = time.monotonic() - start
-        size_mb = len(body) / 1024 / 1024
+        size = len(body)
+        size_mb = size / 1024 / 1024
         logger.debug(
             "S3 upload_part completed",
             bucket=bucket,

s3proxy/crypto.py

@@ -37,6 +37,28 @@
 # Sweet spot: 8MB balances part count vs per-part overhead
 MAX_BUFFER_SIZE = 8 * 1024 * 1024  # 8 MB per internal part
 
+# Framed internal-part format.
+# S3's 10,000-part limit forces large internal parts for big objects, but we do
+# not want to hold a whole part (plaintext + ciphertext) in memory. So a part is
+# encrypted as a sequence of independent AES-GCM *frames*, each sealing up to
+# FRAME_PLAINTEXT_SIZE bytes:
+#
+#   part        = frame[0] || frame[1] || ...
+#   frame[i]    = nonce(12) || ciphertext || tag(16)
+#
+# Both writer and reader stream one frame at a time, so peak memory is O(frame),
+# independent of part size. The reader recovers the frame count purely from the
+# stored sizes:
+#
+#   num_frames = (ciphertext_size - plaintext_size) / ENCRYPTION_OVERHEAD
+#
+# A legacy single-seal part is exactly the num_frames == 1 case, so it is read by
+# the same code path with no migration of stored data.
+#
+# FROZEN: never change FRAME_PLAINTEXT_SIZE. Parts already written were framed at
+# this boundary and the reader splits ciphertext on (FRAME_PLAINTEXT_SIZE + overhead).
+FRAME_PLAINTEXT_SIZE = 8 * 1024 * 1024  # 8 MB plaintext per frame
+
 
 def calculate_optimal_part_size(content_length: int) -> int:
     """Calculate optimal part size to avoid creating parts < 5MB that aren't the final part."""
@@ -140,6 +162,78 @@ def derive_part_nonce(upload_id: str, part_number: int) -> bytes:
     return hashlib.sha256(data).digest()[:NONCE_SIZE]
 
 
+def derive_frame_nonce(upload_id: str, part_number: int, frame_index: int) -> bytes:
+    """Deterministic, unique nonce for one frame of an internal part.
+
+    Within an upload the DEK is fixed, so the (part_number, frame_index) pair must
+    be globally unique to never reuse an AES-GCM nonce. Internal part numbers are
+    unique per upload and frame indexes are unique within a part, so this holds.
+    """
+    data = f"{upload_id}:{part_number}:{frame_index}".encode()
+    return hashlib.sha256(data).digest()[:NONCE_SIZE]
+
+
+def frame_count(plaintext_size: int) -> int:
+    """Number of frames a plaintext of this size is encrypted into."""
+    if plaintext_size <= FRAME_PLAINTEXT_SIZE:
+        return 1
+    return (plaintext_size + FRAME_PLAINTEXT_SIZE - 1) // FRAME_PLAINTEXT_SIZE
+
+
+def framed_ciphertext_size(plaintext_size: int) -> int:
+    """Total stored size of a framed part: plaintext + per-frame GCM overhead."""
+    return plaintext_size + frame_count(plaintext_size) * ENCRYPTION_OVERHEAD
+
+
+def encrypt_frame(
+    plaintext: bytes, dek: bytes, upload_id: str, part_number: int, frame_index: int
+) -> bytes:
+    """Encrypt a single frame (nonce || ciphertext || tag) with its derived nonce."""
+    return encrypt(plaintext, dek, derive_frame_nonce(upload_id, part_number, frame_index))
+
+
+def ciphertext_frame_byte_sizes(plaintext_size: int, ciphertext_size: int) -> list[int]:
+    """Ciphertext byte length of each frame in a (possibly framed) internal part."""
+    return _ciphertext_frame_sizes(plaintext_size, ciphertext_size - plaintext_size)
+
+
+def _ciphertext_frame_sizes(plaintext_size: int, stored_overhead: int) -> list[int]:
+    """Ciphertext byte length of each frame, derived from the stored sizes.
+
+    `stored_overhead = ciphertext_size - plaintext_size` tells us the real frame
+    count for the part as it was written, which is authoritative even if
+    FRAME_PLAINTEXT_SIZE were ever reinterpreted: a legacy single-seal part has
+    overhead == ENCRYPTION_OVERHEAD (one frame) regardless of its plaintext size.
+    """
+    num_frames = stored_overhead // ENCRYPTION_OVERHEAD
+    if num_frames <= 1:
+        return [plaintext_size + ENCRYPTION_OVERHEAD]
+    sizes = []
+    remaining = plaintext_size
+    for _ in range(num_frames):
+        pt = min(FRAME_PLAINTEXT_SIZE, remaining)
+        sizes.append(pt + ENCRYPTION_OVERHEAD)
+        remaining -= pt
+    return sizes
+
+
+def decrypt_framed(ciphertext: bytes, dek: bytes, plaintext_size: int) -> bytes:
+    """Decrypt a (possibly framed) internal part held whole in memory.
+
+    Backward compatible: a legacy single-seal part has only one frame's worth of
+    overhead, so it is decrypted in a single call via decrypt().
+    """
+    sizes = _ciphertext_frame_sizes(plaintext_size, len(ciphertext) - plaintext_size)
+    if len(sizes) == 1:
+        return decrypt(ciphertext, dek)
+    out = bytearray()
+    offset = 0
+    for fsize in sizes:
+        out += decrypt(ciphertext[offset : offset + fsize], dek)
+        offset += fsize
+    return bytes(out)
+
+
 def wrap_key(dek: bytes, kek: bytes) -> bytes:
     """Wrap DEK using AES-KWP (Key Wrap with Padding)."""
     try: