From 271fb6f5737a8b297c6a6aefee1ec7ff2a71b1c5 Mon Sep 17 00:00:00 2001
From: David <david@taiko.xyz>
Date: Mon, 8 Jan 2024 01:34:34 +0800
Subject: [PATCH] fix(prover): fix `--prover.allowance` flag (#490)

---
 prover/prover.go      | 31 ++++++++++++++++++++-----------
 prover/prover_test.go |  4 ++--
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/prover/prover.go b/prover/prover.go
index c34b21ce4..e1a95de93 100644
--- a/prover/prover.go
+++ b/prover/prover.go
@@ -313,7 +313,7 @@ func InitFromConfig(ctx context.Context, p *Prover, cfg *Config) (err error) {
 // setApprovalAmount will set the allowance on the TaikoToken contract for the
 // configured proverAddress as owner and the TaikoL1 contract as spender,
 // if flag is provided for allowance.
-func (p *Prover) setApprovalAmount(ctx context.Context) error {
+func (p *Prover) setApprovalAmount(ctx context.Context, contract common.Address) error {
 	if p.cfg.Allowance == nil || p.cfg.Allowance.Cmp(common.Big0) != 1 {
 		log.Info("Skipping setting approval, `--prover.allowance` flag not set")
 		return nil
@@ -322,19 +322,20 @@ func (p *Prover) setApprovalAmount(ctx context.Context) error {
 	allowance, err := p.rpc.TaikoToken.Allowance(
 		&bind.CallOpts{Context: ctx},
 		p.proverAddress,
-		p.cfg.AssignmentHookAddress,
+		contract,
 	)
 	if err != nil {
 		return err
 	}
 
-	log.Info("Existing allowance for AssignmentHook contract", "allowance", allowance.String())
+	log.Info("Existing allowance for the contract", "allowance", allowance.String(), "contract", contract)
 
 	if allowance.Cmp(p.cfg.Allowance) >= 0 {
 		log.Info(
 			"Skipping setting allowance, allowance already greater or equal",
 			"allowance", allowance.String(),
 			"approvalAmount", p.cfg.Allowance.String(),
+			"contract", contract,
 		)
 		return nil
 	}
@@ -348,11 +349,11 @@ func (p *Prover) setApprovalAmount(ctx context.Context) error {
 	}
 	opts.Context = ctx
 
-	log.Info("Approving AssignmentHook for taiko token", "allowance", p.cfg.Allowance.String())
+	log.Info("Approving the contract for taiko token", "allowance", p.cfg.Allowance.String(), "contract", contract)
 
 	tx, err := p.rpc.TaikoToken.Approve(
 		opts,
-		p.cfg.AssignmentHookAddress,
+		contract,
 		p.cfg.Allowance,
 	)
 	if err != nil {
@@ -364,17 +365,21 @@ func (p *Prover) setApprovalAmount(ctx context.Context) error {
 		return err
 	}
 
-	log.Info("Approved AssignmentHook for taiko token", "txHash", receipt.TxHash.Hex())
+	log.Info(
+		"Approved the contract for taiko token",
+		"txHash", receipt.TxHash.Hex(),
+		"contract", contract,
+	)
 
 	if allowance, err = p.rpc.TaikoToken.Allowance(
 		&bind.CallOpts{Context: ctx},
 		p.proverAddress,
-		p.cfg.AssignmentHookAddress,
+		contract,
 	); err != nil {
 		return err
 	}
 
-	log.Info("New allowance for AssignmentHook contract", "allowance", allowance.String())
+	log.Info("New allowance for the contract", "allowance", allowance.String(), "contract", contract)
 
 	return nil
 }
@@ -384,10 +389,14 @@ func (p *Prover) Start() error {
 	p.wg.Add(1)
 	p.initSubscription()
 
+	if err := p.setApprovalAmount(p.ctx, p.cfg.TaikoL1Address); err != nil {
+		log.Crit("Failed to set approval amount", "contract", p.cfg.TaikoL1Address, "error", err)
+	}
+	if err := p.setApprovalAmount(p.ctx, p.cfg.AssignmentHookAddress); err != nil {
+		log.Crit("Failed to set approval amount", "contract", p.cfg.AssignmentHookAddress, "error", err)
+	}
+
 	go func() {
-		if err := p.setApprovalAmount(p.ctx); err != nil {
-			log.Crit("Failed to set approval amount", "error", err)
-		}
 		if err := p.srv.Start(fmt.Sprintf(":%v", p.cfg.HTTPServerPort)); !errors.Is(err, http.ErrServerClosed) {
 			log.Crit("Failed to start http server", "error", err)
 		}
diff --git a/prover/prover_test.go b/prover/prover_test.go
index 0a7cc5634..58b10407d 100644
--- a/prover/prover_test.go
+++ b/prover/prover_test.go
@@ -410,7 +410,7 @@ func (s *ProverTestSuite) TestSetApprovalAmount() {
 
 	s.p.cfg.Allowance = amt
 
-	s.Nil(s.p.setApprovalAmount(context.Background()))
+	s.Nil(s.p.setApprovalAmount(context.Background(), s.p.cfg.AssignmentHookAddress))
 
 	allowance, err = s.p.rpc.TaikoToken.Allowance(nil, s.p.proverAddress, s.p.cfg.AssignmentHookAddress)
 	s.Nil(err)
@@ -481,7 +481,7 @@ func (s *ProverTestSuite) TestSetApprovalAlreadySetHigher() {
 	amt := common.Big1
 	s.p.cfg.Allowance = amt
 
-	s.Nil(s.p.setApprovalAmount(context.Background()))
+	s.Nil(s.p.setApprovalAmount(context.Background(), s.p.cfg.TaikoL1Address))
 
 	allowance, err := s.p.rpc.TaikoToken.Allowance(&bind.CallOpts{}, s.p.proverAddress, s.p.cfg.TaikoL1Address)
 	s.Nil(err)