From 714713695ea1d6e38448246baa341933a9129423 Mon Sep 17 00:00:00 2001
From: Manu NALEPA <enalepa@offchainlabs.com>
Date: Tue, 4 Jun 2024 10:38:59 +0200
Subject: [PATCH] PeerDAS: Withhold data on purpose. (#14076)

* Introduce hidden flag `data-columns-withhold-count`.

* Address Nishant's comment.
---
 .../rpc/prysm/v1alpha1/validator/proposer.go  | 22 +++++++++++++++----
 config/features/config.go                     |  8 +++++++
 config/features/flags.go                      |  8 +++++++
 3 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/beacon-chain/rpc/prysm/v1alpha1/validator/proposer.go b/beacon-chain/rpc/prysm/v1alpha1/validator/proposer.go
index e6ff8fec7324..7f6dd8a879d4 100644
--- a/beacon-chain/rpc/prysm/v1alpha1/validator/proposer.go
+++ b/beacon-chain/rpc/prysm/v1alpha1/validator/proposer.go
@@ -458,14 +458,28 @@ func (vs *Server) broadcastAndReceiveBlobs(ctx context.Context, sidecars []*ethp
 // broadcastAndReceiveDataColumns handles the broadcasting and reception of data columns sidecars.
 func (vs *Server) broadcastAndReceiveDataColumns(ctx context.Context, sidecars []*ethpb.DataColumnSidecar, root [fieldparams.RootLength]byte) error {
 	eg, _ := errgroup.WithContext(ctx)
+
+	dataColumnsWithholdCount := features.Get().DataColumnsWithholdCount
+
 	for i, sd := range sidecars {
 		// Copy the iteration instance to a local variable to give each go-routine its own copy to play with.
 		// See https://golang.org/doc/faq#closures_and_goroutines for more details.
-		colIdx := i
-		sidecar := sd
+		colIdx, sidecar := i, sd
+
 		eg.Go(func() error {
-			if err := vs.P2P.BroadcastDataColumn(ctx, uint64(colIdx)%params.BeaconConfig().DataColumnSidecarSubnetCount, sidecar); err != nil {
-				return errors.Wrap(err, "broadcast data column")
+			// Compute the subnet index based on the column index.
+			subnet := uint64(colIdx) % params.BeaconConfig().DataColumnSidecarSubnetCount
+
+			if colIdx < dataColumnsWithholdCount {
+				log.WithFields(logrus.Fields{
+					"root":            fmt.Sprintf("%#x", root),
+					"subnet":          subnet,
+					"dataColumnIndex": colIdx,
+				}).Warning("Withholding data column")
+			} else {
+				if err := vs.P2P.BroadcastDataColumn(ctx, subnet, sidecar); err != nil {
+					return errors.Wrap(err, "broadcast data column")
+				}
 			}
 
 			roDataColumn, err := blocks.NewRODataColumnWithRoot(sidecar, root)
diff --git a/config/features/config.go b/config/features/config.go
index b9e407aebb4d..3fdfd1f20ad9 100644
--- a/config/features/config.go
+++ b/config/features/config.go
@@ -84,6 +84,9 @@ type Flags struct {
 	// changed on disk. This feature is for advanced use cases only.
 	KeystoreImportDebounceInterval time.Duration
 
+	// DataColumnsWithholdCount specifies the likelihood of withholding a data column sidecar when proposing a block (percentage)
+	DataColumnsWithholdCount int
+
 	// AggregateIntervals specifies the time durations at which we aggregate attestations preparing for forkchoice.
 	AggregateIntervals [3]time.Duration
 }
@@ -269,6 +272,11 @@ func ConfigureBeaconChain(ctx *cli.Context) error {
 		cfg.EnablePeerDAS = true
 	}
 
+	if ctx.IsSet(DataColumnsWithholdCount.Name) {
+		logEnabled(DataColumnsWithholdCount)
+		cfg.DataColumnsWithholdCount = ctx.Int(DataColumnsWithholdCount.Name)
+	}
+
 	cfg.AggregateIntervals = [3]time.Duration{aggregateFirstInterval.Value, aggregateSecondInterval.Value, aggregateThirdInterval.Value}
 	Init(cfg)
 	return nil
diff --git a/config/features/flags.go b/config/features/flags.go
index 1ed7695a5a01..94f85b3e1bb6 100644
--- a/config/features/flags.go
+++ b/config/features/flags.go
@@ -174,6 +174,13 @@ var (
 		Name:  "peer-das",
 		Usage: "Enables Prysm to run with the experimental peer data availability sampling scheme.",
 	}
+	// DataColumnsWithholdCount is a flag for withholding data columns when proposing a block.
+	DataColumnsWithholdCount = &cli.IntFlag{
+		Name:   "data-columns-withhold-count",
+		Usage:  "Number of columns to withhold when proposing a block. DO NOT USE IN PRODUCTION.",
+		Value:  0,
+		Hidden: true,
+	}
 )
 
 // devModeFlags holds list of flags that are set when development mode is on.
@@ -232,6 +239,7 @@ var BeaconChainFlags = append(deprecatedBeaconFlags, append(deprecatedFlags, []c
 	EnableQUIC,
 	DisableCommitteeAwarePacking,
 	EnablePeerDAS,
+	DataColumnsWithholdCount,
 }...)...)
 
 // E2EBeaconChainFlags contains a list of the beacon chain feature flags to be tested in E2E.