Merge branch 'feat/backup-phase0a-filename' of github.com:bootjp/elastickv into feat/backup-phase0a-keymap-manifest

Author: bootjp

internal/backup/filename.go

@@ -81,33 +81,64 @@ var ErrShaFallbackNeedsKeymap = errors.New("backup: filename uses SHA fallback;
 // inputs.
 //
 // The encoding is deterministic and idempotent given the same input.
+//
+// Two short-circuits ensure the encoder never trips its own invariants:
+//
+//   - If raw is so large that percent-encoding it would always overflow
+//     maxSegmentBytes (3*len(raw) > maxSegmentBytes), we go straight to
+//     shaFallback without allocating the full expansion. Without this an
+//     adversarial caller could force a very large transient allocation
+//     just to discard it.
+//   - If the percent-encoded form happens to match the SHA-fallback shape
+//     (32 hex chars followed by "__"), we promote it to a real
+//     SHA-fallback so DecodeSegment's structural detection cannot
+//     misclassify a legitimate key. Both isShaFallback and shaFallback
+//     are true on the resulting output, so KEYMAP.jsonl carries the
+//     original bytes for exact-byte recovery.
 func EncodeSegment(raw []byte) string {
+	if len(raw)*percentEncodeMaxExpansion > maxSegmentBytes {
+		return shaFallback(raw)
+	}
 	encoded := percentEncode(raw)
-	if len(encoded) <= maxSegmentBytes {
-		return encoded
+	if len(encoded) > maxSegmentBytes || isShaFallback(encoded) {
+		return shaFallback(raw)
 	}
-	return shaFallback(raw)
+	return encoded
 }
 
 // EncodeBinarySegment encodes a DynamoDB B-attribute (binary) segment as
 // "b64.<base64url-no-padding>" so that binary keys never collide with string
 // keys whose hex-encoding happens to look like base64.
 //
-// b64-encoded segments take the SHA fallback if they exceed maxSegmentBytes
-// after the base64 expansion (~4/3 of the raw length).
+// Short-circuits the SHA-fallback for inputs whose base64 expansion (~4/3 of
+// the raw length, plus the 4-byte "b64." prefix) would always overflow
+// maxSegmentBytes. As with EncodeSegment, this avoids an unnecessary large
+// allocation when the result would have been discarded anyway.
 func EncodeBinarySegment(raw []byte) string {
+	if base64.RawURLEncoding.EncodedLen(len(raw))+len(binaryPrefix) > maxSegmentBytes {
+		return shaFallback(raw)
+	}
 	enc := binaryPrefix + base64.RawURLEncoding.EncodeToString(raw)
-	if len(enc) <= maxSegmentBytes {
-		return enc
+	if len(enc) > maxSegmentBytes {
+		return shaFallback(raw)
 	}
-	return shaFallback(raw)
+	return enc
 }
 
 // DecodeSegment is the inverse of EncodeSegment for percent-encoded and
 // binary-prefixed inputs. SHA-fallback inputs return ErrShaFallbackNeedsKeymap
 // so the caller knows to consult KEYMAP.jsonl rather than treat the partial
 // suffix as the original key.
+//
+// As a defensive measure DecodeSegment refuses inputs longer than
+// maxSegmentBytes. EncodeSegment never produces such inputs, so any caller
+// passing one is either reading a corrupted dump or has a bug; either way the
+// percentDecode allocation should not run.
 func DecodeSegment(seg string) ([]byte, error) {
+	if len(seg) > maxSegmentBytes {
+		return nil, errors.Wrapf(ErrInvalidEncodedSegment,
+			"segment length %d exceeds maximum %d", len(seg), maxSegmentBytes)
+	}
 	if isShaFallback(seg) {
 		return nil, errors.WithStack(ErrShaFallbackNeedsKeymap)
 	}
@@ -121,6 +152,10 @@ func DecodeSegment(seg string) ([]byte, error) {
 	return percentDecode(seg)
 }
 
+// percentEncodeMaxExpansion is the worst-case ratio of encoded length to
+// raw length for percentEncode (every byte expands to "%HH").
+const percentEncodeMaxExpansion = 3
+
 // IsShaFallback reports whether seg uses the SHA-prefix-and-truncated-original
 // form. Such segments cannot be reversed without KEYMAP.jsonl.
 func IsShaFallback(seg string) bool {

internal/backup/filename_test.go

@@ -309,6 +309,72 @@ func TestEncodeBinarySegment_FuzzRoundTripIfNotShaFallback(t *testing.T) {
 	})
 }
 
+func TestEncodeSegment_KeyMatchingShaFallbackShapeIsPromotedToFallback(t *testing.T) {
+	t.Parallel()
+	// A user key that is itself made of 32 hex chars + "__" + suffix
+	// would, under naive encoding, return the raw bytes unchanged
+	// (everything is unreserved) — but DecodeSegment's structural
+	// detection would then misclassify it as a SHA-fallback and
+	// return ErrShaFallbackNeedsKeymap. EncodeSegment must promote
+	// such inputs to a real SHA-fallback so the encoded->decoded
+	// invariant holds (decode refuses; KEYMAP carries the original).
+	raw := []byte("0123456789abcdef0123456789abcdef__suffix")
+	enc := EncodeSegment(raw)
+	if !IsShaFallback(enc) {
+		t.Fatalf("expected SHA fallback for collision-shaped input, got %q", enc)
+	}
+	// The fallback's hex prefix must be the SHA of the raw bytes,
+	// NOT the raw bytes' first 32 chars. That way a KEYMAP entry
+	// keyed on `enc` carries the actual original — not a structural
+	// echo.
+	if _, err := DecodeSegment(enc); !errors.Is(err, ErrShaFallbackNeedsKeymap) {
+		t.Fatalf("decode of promoted fallback: err=%v want ErrShaFallbackNeedsKeymap", err)
+	}
+}
+
+func TestEncodeSegment_HugeInputDoesNotMaterialiseFullExpansion(t *testing.T) {
+	t.Parallel()
+	// A 1 MiB input would, if percent-encoded eagerly, allocate 3
+	// MiB before the length check fired. The early short-circuit
+	// must skip that allocation. We can't directly observe the
+	// allocation here without a profile, but we can assert the
+	// output is correct (SHA fallback, length under the ceiling)
+	// and that the call returns promptly enough to be a no-op
+	// guard in profile-runs.
+	raw := make([]byte, 1<<20) // 1 MiB
+	for i := range raw {
+		raw[i] = byte(i)
+	}
+	enc := EncodeSegment(raw)
+	if !IsShaFallback(enc) {
+		t.Fatalf("expected SHA fallback for huge input")
+	}
+	if len(enc) > maxSegmentBytes {
+		t.Fatalf("encoded len %d > max %d", len(enc), maxSegmentBytes)
+	}
+}
+
+func TestDecodeSegment_RejectsOversizedInput(t *testing.T) {
+	t.Parallel()
+	too := strings.Repeat("a", maxSegmentBytes+1)
+	_, err := DecodeSegment(too)
+	if !errors.Is(err, ErrInvalidEncodedSegment) {
+		t.Fatalf("err=%v want ErrInvalidEncodedSegment for oversized input", err)
+	}
+}
+
+func TestEncodeBinarySegment_HugeInputTakesShaFallbackWithoutEncoding(t *testing.T) {
+	t.Parallel()
+	raw := make([]byte, 1<<20) // 1 MiB
+	enc := EncodeBinarySegment(raw)
+	if !IsShaFallback(enc) {
+		t.Fatalf("expected SHA fallback for huge binary input, got %q", enc[:min(40, len(enc))])
+	}
+	if len(enc) > maxSegmentBytes {
+		t.Fatalf("encoded len %d > max %d", len(enc), maxSegmentBytes)
+	}
+}
+
 func TestEncodeSegment_ShaFallbackEmbedsRecognisableSuffix(t *testing.T) {
 	t.Parallel()
 	// The truncated suffix in the SHA-fallback rendering must be derivable