VAULT-8436 remove <-time.After statements in for loops (#18818)

* replace time.After with ticker in loops

* add semgrep rule

* update to use timers

* remove stop

api/lifetime_watcher.go

@@ -366,10 +366,12 @@ func (r *LifetimeWatcher) doRenewWithOptions(tokenMode bool, nonRenewable bool,
 			return nil
 		}
 
+		timer := time.NewTimer(sleepDuration)
 		select {
 		case <-r.stopCh:
+			timer.Stop()
 			return nil
-		case <-time.After(sleepDuration):
+		case <-timer.C:
 			continue
 		}
 	}

builtin/credential/okta/backend.go

@@ -267,10 +267,12 @@ func (b *backend) Login(ctx context.Context, req *logical.Request, username, pas
 					return nil, logical.ErrorResponse("okta auth backend unexpected failure"), nil, nil
 				}
 
+				timer := time.NewTimer(1 * time.Second)
 				select {
-				case <-time.After(1 * time.Second):
+				case <-timer.C:
 					// Continue
 				case <-ctx.Done():
+					timer.Stop()
 					return nil, logical.ErrorResponse("exiting pending mfa challenge"), nil, nil
 				}
 			case "REJECTED":

builtin/credential/okta/cli.go

@@ -61,10 +61,12 @@ func (h *CLIHandler) Auth(c *api.Client, m map[string]string) (*api.Secret, erro
 
 	go func() {
 		for {
+			timer := time.NewTimer(time.Second)
 			select {
 			case <-doneCh:
+				timer.Stop()
 				return
-			case <-time.After(time.Second):
+			case <-timer.C:
 			}
 
 			resp, _ := c.Logical().Read(fmt.Sprintf("auth/%s/verify/%s", mount, nonce))

command/agent/sink/sink.go

@@ -151,10 +151,12 @@ func (ss *SinkServer) Run(ctx context.Context, incoming chan string, sinks []*Si
 			if err := writeSink(st.sink, st.token); err != nil {
 				backoff := 2*time.Second + time.Duration(ss.random.Int63()%int64(time.Second*2)-int64(time.Second))
 				ss.logger.Error("error returned by sink function, retrying", "error", err, "backoff", backoff.String())
+				timer := time.NewTimer(backoff)
 				select {
 				case <-ctx.Done():
+					timer.Stop()
 					return nil
-				case <-time.After(backoff):
+				case <-timer.C:
 					atomic.AddInt32(ss.remaining, 1)
 					sinkCh <- st
 				}

command/server.go

@@ -2319,9 +2319,11 @@ func (c *ServerCommand) storageMigrationActive(backend physical.Backend) bool {
 		}
 		c.logger.Warn("storage migration check error", "error", err.Error())
 
+		timer := time.NewTimer(2 * time.Second)
 		select {
-		case <-time.After(2 * time.Second):
+		case <-timer.C:
 		case <-c.ShutdownCh:
+			timer.Stop()
 			return true
 		}
 	}
@@ -2609,10 +2611,12 @@ func runUnseal(c *ServerCommand, core *vault.Core, ctx context.Context) {
 		}
 		c.logger.Warn("failed to unseal core", "error", err)
 
+		timer := time.NewTimer(5 * time.Second)
 		select {
 		case <-c.ShutdownCh:
+			timer.Stop()
 			return
-		case <-time.After(5 * time.Second):
+		case <-timer.C:
 		}
 	}
 }

physical/zookeeper/zookeeper.go

@@ -642,8 +642,9 @@ func (i *ZooKeeperHALock) Unlock() error {
 					return
 				}
 
+				timer := time.NewTimer(time.Second)
 				select {
-				case <-time.After(time.Second):
+				case <-timer.C:
 					attempts := attempts + 1
 					if attempts >= 10 {
 						i.logger.Error("release lock max attempts reached. Lock may not be released", "error", err)
@@ -651,6 +652,7 @@ func (i *ZooKeeperHALock) Unlock() error {
 					}
 					continue
 				case <-i.stopCh:
+					timer.Stop()
 					return
 				}
 			}

tools/semgrep/ci/loop-time-after.yml

@@ -0,0 +1,17 @@
+rules:
+  - id: loop-time-after
+    pattern: |
+      for ... {
+        ...
+        select {
+          case ...
+          case <-time.After(...):
+            ...
+          case ...
+        }
+        ...
+      }
+    message: <-time.After() used in for loop, consider using a ticker or a timer instead
+    languages:
+      - go
+    severity: WARNING

vault/ha.go

@@ -663,10 +663,12 @@ func (c *Core) waitForLeadership(newLeaderCh chan func(), manualStepDownCh, stop
 			// Spawn this in a go routine so we can cancel the context and
 			// unblock any inflight requests that are holding the statelock.
 			go func() {
+				timer := time.NewTimer(DefaultMaxRequestDuration)
 				select {
 				case <-activeCtx.Done():
+					timer.Stop()
 					// Attempt to drain any inflight requests
-				case <-time.After(DefaultMaxRequestDuration):
+				case <-timer.C:
 					activeCtxCancel()
 				}
 			}()
@@ -796,8 +798,9 @@ func (c *Core) periodicLeaderRefresh(newLeaderCh chan func(), stopCh chan struct
 
 	clusterAddr := ""
 	for {
+		timer := time.NewTimer(leaderCheckInterval)
 		select {
-		case <-time.After(leaderCheckInterval):
+		case <-timer.C:
 			count := atomic.AddInt32(opCount, 1)
 			if count > 1 {
 				atomic.AddInt32(opCount, -1)
@@ -830,6 +833,7 @@ func (c *Core) periodicLeaderRefresh(newLeaderCh chan func(), stopCh chan struct
 				atomic.AddInt32(lopCount, -1)
 			}()
 		case <-stopCh:
+			timer.Stop()
 			return
 		}
 	}
@@ -842,8 +846,9 @@ func (c *Core) periodicCheckKeyUpgrades(ctx context.Context, stopCh chan struct{
 
 	opCount := new(int32)
 	for {
+		timer := time.NewTimer(keyRotateCheckInterval)
 		select {
-		case <-time.After(keyRotateCheckInterval):
+		case <-timer.C:
 			count := atomic.AddInt32(opCount, 1)
 			if count > 1 {
 				atomic.AddInt32(opCount, -1)
@@ -899,6 +904,7 @@ func (c *Core) periodicCheckKeyUpgrades(ctx context.Context, stopCh chan struct{
 				return
 			}()
 		case <-stopCh:
+			timer.Stop()
 			return
 		}
 	}
@@ -1030,9 +1036,11 @@ func (c *Core) acquireLock(lock physical.Lock, stopCh <-chan struct{}) <-chan st
 
 		// Retry the acquisition
 		c.logger.Error("failed to acquire lock", "error", err)
+		timer := time.NewTimer(lockRetryInterval)
 		select {
-		case <-time.After(lockRetryInterval):
+		case <-timer.C:
 		case <-stopCh:
+			timer.Stop()
 			return nil
 		}
 	}
@@ -1099,13 +1107,15 @@ func (c *Core) cleanLeaderPrefix(ctx context.Context, uuid string, leaderLostCh
 		return
 	}
 	for len(keys) > 0 {
+		timer := time.NewTimer(leaderPrefixCleanDelay)
 		select {
-		case <-time.After(leaderPrefixCleanDelay):
+		case <-timer.C:
 			if keys[0] != uuid {
 				c.barrier.Delete(ctx, coreLeaderPrefix+keys[0])
 			}
 			keys = keys[1:]
 		case <-leaderLostCh:
+			timer.Stop()
 			return
 		}
 	}

vault/login_mfa.go

@@ -1995,11 +1995,13 @@ func (c *Core) validateDuo(ctx context.Context, mfaFactors *MFAFactor, mConfig *
 		case "allow":
 			return nil
 		}
+		timer := time.NewTimer(time.Second)
 
 		select {
 		case <-ctx.Done():
+			timer.Stop()
 			return fmt.Errorf("duo push verification operation canceled")
-		case <-time.After(time.Second):
+		case <-timer.C:
 		}
 	}
 }
@@ -2124,11 +2126,13 @@ func (c *Core) validateOkta(ctx context.Context, mConfig *mfa.Config, username s
 		default:
 			return fmt.Errorf("unknown status code")
 		}
+		timer := time.NewTimer(time.Second)
 
 		select {
 		case <-ctx.Done():
+			timer.Stop()
 			return fmt.Errorf("push verification operation canceled")
-		case <-time.After(time.Second):
+		case <-timer.C:
 		}
 	}
 }

vault/raft.go

@@ -433,8 +433,9 @@ func (c *Core) raftTLSRotateDirect(ctx context.Context, logger hclog.Logger, sto
 				backoff = false
 			}
 
+			timer := time.NewTimer(time.Until(nextRotationTime))
 			select {
-			case <-time.After(time.Until(nextRotationTime)):
+			case <-timer.C:
 				// It's time to rotate the keys
 				next, err := rotateKeyring()
 				if err != nil {
@@ -446,6 +447,7 @@ func (c *Core) raftTLSRotateDirect(ctx context.Context, logger hclog.Logger, sto
 				nextRotationTime = next
 
 			case <-stopCh:
+				timer.Stop()
 				return
 			}
 		}