From 118e065de031b63ec07231b63e555943f1292fad Mon Sep 17 00:00:00 2001
From: NC <17676176+ensi321@users.noreply.github.com>
Date: Wed, 23 Oct 2024 00:42:34 -0700
Subject: [PATCH] Rewrite SyncCommitteeWitnessRepository

---
 .../src/chain/lightClient/index.ts            | 11 +++-
 .../src/chain/lightClient/types.ts            |  2 +-
 .../lightclientSyncCommitteeWitness.ts        | 55 ++++++++++++++++++-
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/packages/beacon-node/src/chain/lightClient/index.ts b/packages/beacon-node/src/chain/lightClient/index.ts
index bb7de1f0478a..d3a620008c11 100644
--- a/packages/beacon-node/src/chain/lightClient/index.ts
+++ b/packages/beacon-node/src/chain/lightClient/index.ts
@@ -622,6 +622,16 @@ export class LightClientServer {
     if (!syncCommitteeWitness) {
       throw Error(`syncCommitteeWitness not available at ${toRootHex(attestedData.blockRoot)}`);
     }
+
+    const attestedFork = this.config.getForkName(attestedHeader.beacon.slot);
+    const numWitness = syncCommitteeWitness.witness.length;
+    if (isForkPostElectra(attestedFork) && numWitness !== NUM_WITNESS_ELECTRA) {
+      throw Error(`Expected ${NUM_WITNESS_ELECTRA} witnesses in post-Electra numWiteness=${numWitness}`);
+    }
+    if (!isForkPostElectra(attestedFork) && numWitness !== NUM_WITNESS) {
+      throw Error(`Expected ${NUM_WITNESS} witnesses in pre-Electra numWiteness=${numWitness}`);
+    }
+
     const nextSyncCommittee = await this.db.syncCommittee.get(syncCommitteeWitness.nextSyncCommitteeRoot);
     if (!nextSyncCommittee) {
       throw Error("nextSyncCommittee not available");
@@ -642,7 +652,6 @@ export class LightClientServer {
       finalityBranch = attestedData.finalityBranch;
       finalizedHeader = finalizedHeaderAttested;
       // Fork of LightClientUpdate is based off on attested header's fork
-      const attestedFork = this.config.getForkName(attestedHeader.beacon.slot);
       if (this.config.getForkName(finalizedHeader.beacon.slot) !== attestedFork) {
         finalizedHeader = upgradeLightClientHeader(this.config, attestedFork, finalizedHeader);
       }
diff --git a/packages/beacon-node/src/chain/lightClient/types.ts b/packages/beacon-node/src/chain/lightClient/types.ts
index b253c05d45fb..00a819f30c15 100644
--- a/packages/beacon-node/src/chain/lightClient/types.ts
+++ b/packages/beacon-node/src/chain/lightClient/types.ts
@@ -26,7 +26,7 @@
  * ```
  */
 export type SyncCommitteeWitness = {
-  /** Vector[Bytes32, 4] */
+  /** Vector[Bytes32, 4] or Vector[Bytes32, 5] depends on the fork */
   witness: Uint8Array[];
   currentSyncCommitteeRoot: Uint8Array;
   nextSyncCommitteeRoot: Uint8Array;
diff --git a/packages/beacon-node/src/db/repositories/lightclientSyncCommitteeWitness.ts b/packages/beacon-node/src/db/repositories/lightclientSyncCommitteeWitness.ts
index 45f91f159997..e323c3e55f61 100644
--- a/packages/beacon-node/src/db/repositories/lightclientSyncCommitteeWitness.ts
+++ b/packages/beacon-node/src/db/repositories/lightclientSyncCommitteeWitness.ts
@@ -5,6 +5,15 @@ import {ssz} from "@lodestar/types";
 import {SyncCommitteeWitness} from "../../chain/lightClient/types.js";
 import {Bucket, getBucketNameByValue} from "../buckets.js";
 
+// We add a 1-byte prefix where 0 means pre-electra and 1 means post-electra
+enum PrefixByte {
+  PRE_ELECTRA = 0,
+  POST_ELECTRA = 1,
+}
+
+export const NUM_WITNESS = 4;
+export const NUM_WITNESS_ELECTRA = 5;
+
 /**
  * Historical sync committees witness by block root
  *
@@ -13,12 +22,56 @@ import {Bucket, getBucketNameByValue} from "../buckets.js";
 export class SyncCommitteeWitnessRepository extends Repository<Uint8Array, SyncCommitteeWitness> {
   constructor(config: ChainForkConfig, db: DatabaseController<Uint8Array, Uint8Array>) {
     const bucket = Bucket.lightClient_syncCommitteeWitness;
+    // Pick some type but won't be used. Witness can be 4 or 5 so need to handle dynamically
     const type = new ContainerType({
-      witness: new VectorCompositeType(ssz.Root, 4),
+      witness: new VectorCompositeType(ssz.Root, NUM_WITNESS),
       currentSyncCommitteeRoot: ssz.Root,
       nextSyncCommitteeRoot: ssz.Root,
     });
 
     super(config, db, bucket, type, getBucketNameByValue(bucket));
   }
+
+  // Overrides for multi-fork
+  encodeValue(value: SyncCommitteeWitness): Uint8Array {
+    const numWitness = value.witness.length;
+
+    if (numWitness !== NUM_WITNESS && numWitness !== NUM_WITNESS_ELECTRA) {
+      throw Error(`Number of witness can only be 4 pre-electra or 5 post-electra numWitness=${numWitness}`);
+    }
+
+    const type = new ContainerType({
+      witness: new VectorCompositeType(ssz.Root, numWitness),
+      currentSyncCommitteeRoot: ssz.Root,
+      nextSyncCommitteeRoot: ssz.Root,
+    });
+
+    const valueBytes = type.serialize(value);
+
+    // We need to differentiate between post-electra and pre-electra witness
+    // such that we can deserialize correctly
+    const isPostElectra = numWitness === NUM_WITNESS_ELECTRA;
+    const prefixByte = new Uint8Array(1);
+    prefixByte[0] = isPostElectra ? PrefixByte.POST_ELECTRA : PrefixByte.PRE_ELECTRA;
+
+    const prefixedData = new Uint8Array(1 + valueBytes.length);
+    prefixedData.set(prefixByte, 0);
+    prefixedData.set(valueBytes, 1);
+
+    return prefixedData;
+  }
+
+  decodeValue(data: Uint8Array): SyncCommitteeWitness {
+    // First byte is written
+    const prefix = data.subarray(0, 1);
+    const isPostElectra = prefix[0] === PrefixByte.POST_ELECTRA;
+
+    const type = new ContainerType({
+      witness: new VectorCompositeType(ssz.Root, isPostElectra ? NUM_WITNESS_ELECTRA : NUM_WITNESS),
+      currentSyncCommitteeRoot: ssz.Root,
+      nextSyncCommitteeRoot: ssz.Root,
+    });
+
+    return type.deserialize(data.subarray(1));
+  }
 }