From 0323dd01b2c2e75b399f83fcac2d8f6fe6cc28da Mon Sep 17 00:00:00 2001
From: Andrew Dunham <andrew@du.nham.ca>
Date: Wed, 26 Jun 2024 11:29:53 -0400
Subject: [PATCH] ci: enable checklocks workflow for specific packages

This turns the checklocks workflow into a real check, and adds
annotations to a few basic packages as a starting point.

Updates #12625

Signed-off-by: Andrew Dunham <andrew@du.nham.ca>
Change-Id: I2b0185bae05a843b5257980fc6bde732b1bdd93f
---
 .github/workflows/checklocks.yml | 10 ++++++++--
 envknob/envknob.go               | 23 +++++++++++++++++------
 ipn/store/mem/store_mem.go       |  3 ++-
 net/stun/stuntest/stuntest.go    |  4 +++-
 net/wsconn/wsconn.go             |  6 +++++-
 proxymap/proxymap.go             |  7 ++++++-
 6 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/.github/workflows/checklocks.yml b/.github/workflows/checklocks.yml
index 17d853f84..9d69783ca 100644
--- a/.github/workflows/checklocks.yml
+++ b/.github/workflows/checklocks.yml
@@ -24,5 +24,11 @@ jobs:
         run: ./tool/go build -o /tmp/checklocks gvisor.dev/gvisor/tools/checklocks/cmd/checklocks
 
       - name: Run checklocks vet
-        # TODO: remove || true once we have applied checklocks annotations everywhere.
-        run: ./tool/go vet -vettool=/tmp/checklocks ./... || true
+        # TODO(#12625): add more packages as we add annotations
+        run: |-
+          ./tool/go vet -vettool=/tmp/checklocks \
+            ./envknob           \
+            ./ipn/store/mem     \
+            ./net/stun/stuntest \
+            ./net/wsconn        \
+            ./proxymap
diff --git a/envknob/envknob.go b/envknob/envknob.go
index e1ffaa8e8..8873f0077 100644
--- a/envknob/envknob.go
+++ b/envknob/envknob.go
@@ -36,13 +36,19 @@
 )
 
 var (
-	mu          sync.Mutex
-	set         = map[string]string{}
-	regStr      = map[string]*string{}
-	regBool     = map[string]*bool{}
-	regOptBool  = map[string]*opt.Bool{}
+	mu sync.Mutex
+	// +checklocks:mu
+	set = map[string]string{}
+	// +checklocks:mu
+	regStr = map[string]*string{}
+	// +checklocks:mu
+	regBool = map[string]*bool{}
+	// +checklocks:mu
+	regOptBool = map[string]*opt.Bool{}
+	// +checklocks:mu
 	regDuration = map[string]*time.Duration{}
-	regInt      = map[string]*int{}
+	// +checklocks:mu
+	regInt = map[string]*int{}
 )
 
 func noteEnv(k, v string) {
@@ -51,6 +57,7 @@ func noteEnv(k, v string) {
 	noteEnvLocked(k, v)
 }
 
+// +checklocks:mu
 func noteEnvLocked(k, v string) {
 	if v != "" {
 		set[k] = v
@@ -202,6 +209,7 @@ func RegisterInt(envVar string) func() int {
 	return func() int { return *p }
 }
 
+// +checklocks:mu
 func setBoolLocked(p *bool, envVar, val string) {
 	noteEnvLocked(envVar, val)
 	if val == "" {
@@ -215,6 +223,7 @@ func setBoolLocked(p *bool, envVar, val string) {
 	}
 }
 
+// +checklocks:mu
 func setOptBoolLocked(p *opt.Bool, envVar, val string) {
 	noteEnvLocked(envVar, val)
 	if val == "" {
@@ -228,6 +237,7 @@ func setOptBoolLocked(p *opt.Bool, envVar, val string) {
 	p.Set(b)
 }
 
+// +checklocks:mu
 func setDurationLocked(p *time.Duration, envVar, val string) {
 	noteEnvLocked(envVar, val)
 	if val == "" {
@@ -241,6 +251,7 @@ func setDurationLocked(p *time.Duration, envVar, val string) {
 	}
 }
 
+// +checklocks:mu
 func setIntLocked(p *int, envVar, val string) {
 	noteEnvLocked(envVar, val)
 	if val == "" {
diff --git a/ipn/store/mem/store_mem.go b/ipn/store/mem/store_mem.go
index 8835374b0..f3a308ae5 100644
--- a/ipn/store/mem/store_mem.go
+++ b/ipn/store/mem/store_mem.go
@@ -20,7 +20,8 @@ func New(logger.Logf, string) (ipn.StateStore, error) {
 
 // Store is an ipn.StateStore that keeps state in memory only.
 type Store struct {
-	mu    sync.Mutex
+	mu sync.Mutex
+	// +checklocks:mu
 	cache map[ipn.StateKey][]byte
 }
 
diff --git a/net/stun/stuntest/stuntest.go b/net/stun/stuntest/stuntest.go
index 321757807..096841600 100644
--- a/net/stun/stuntest/stuntest.go
+++ b/net/stun/stuntest/stuntest.go
@@ -21,8 +21,10 @@
 )
 
 type stunStats struct {
-	mu       sync.Mutex
+	mu sync.Mutex
+	// +checklocks:mu
 	readIPv4 int
+	// +checklocks:mu
 	readIPv6 int
 }
 
diff --git a/net/wsconn/wsconn.go b/net/wsconn/wsconn.go
index 2d708ac53..2e952f14f 100644
--- a/net/wsconn/wsconn.go
+++ b/net/wsconn/wsconn.go
@@ -105,7 +105,11 @@ type netConn struct {
 	afterReadDeadline atomic.Bool
 
 	readMu sync.Mutex
-	eofed  bool
+	// eofed is true if the reader should return io.EOF from the Read call.
+	//
+	// +checklocks:readMu
+	eofed bool
+	// +checklocks:readMu
 	reader io.Reader
 }
 
diff --git a/proxymap/proxymap.go b/proxymap/proxymap.go
index 5df93d61c..8a7f1f95e 100644
--- a/proxymap/proxymap.go
+++ b/proxymap/proxymap.go
@@ -19,7 +19,12 @@
 // given localhost:port corresponds to.
 type Mapper struct {
 	mu sync.Mutex
-	m  map[string]map[netip.AddrPort]netip.Addr // proto ("tcp", "udp") => ephemeral => tailscale IP
+
+	// m holds the mapping from localhost IP:ports to Tailscale IPs. It is
+	// keyed first by the protocol ("tcp" or "udp"), then by the IP:port.
+	//
+	// +checklocks:mu
+	m map[string]map[netip.AddrPort]netip.Addr
 }
 
 // RegisterIPPortIdentity registers a given node (identified by its