From 7124367b820889d0ac17eaf6fef2f595b6707ab9 Mon Sep 17 00:00:00 2001
From: Kevin Krakauer <krakauer@google.com>
Date: Wed, 27 Apr 2022 11:22:20 -0700
Subject: [PATCH] prohibit direct use of sync/atomic (u)int32 functions

Also adds the "// +checkalignedignore" escape hatch for packages to opt out of
checking.

See cl/439349432 for justification (https://github.com/google/gvisor/pull/7376).

PiperOrigin-RevId: 444918125
---
 pkg/atomicbitops/atomicbitops.go   |  2 ++
 pkg/sync/sync.go                   |  2 ++
 tools/checkaligned/checkaligned.go | 36 ++++++++++++++++++++++++------
 3 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/pkg/atomicbitops/atomicbitops.go b/pkg/atomicbitops/atomicbitops.go
index 696f0108e..36620b3cb 100644
--- a/pkg/atomicbitops/atomicbitops.go
+++ b/pkg/atomicbitops/atomicbitops.go
@@ -19,6 +19,8 @@
 //
 // All read-modify-write operations implemented by this package have
 // acquire-release memory ordering (like sync/atomic).
+//
+// +checkalignedignore
 package atomicbitops
 
 // AndUint32 atomically applies bitwise AND operation to *addr with val.
diff --git a/pkg/sync/sync.go b/pkg/sync/sync.go
index b16cf5333..a9bf146db 100644
--- a/pkg/sync/sync.go
+++ b/pkg/sync/sync.go
@@ -4,4 +4,6 @@
 // license that can be found in the LICENSE file.
 
 // Package sync provides synchronization primitives.
+//
+// +checkalignedignore
 package sync
diff --git a/tools/checkaligned/checkaligned.go b/tools/checkaligned/checkaligned.go
index 0c4977af8..b8299d09f 100644
--- a/tools/checkaligned/checkaligned.go
+++ b/tools/checkaligned/checkaligned.go
@@ -12,13 +12,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package checkaligned ensures that atomic (u)int64 operations happen
+// Package checkaligned ensures that atomic (u)int operations happen
 // exclusively via the atomicbitops package.
+//
+// We support a "// +checkalignedignore" escape hatch in the package comment
+// that disables checking throughout the package.
 package checkaligned
 
 import (
 	"fmt"
 	"go/ast"
+	"strings"
 
 	"golang.org/x/tools/go/analysis"
 )
@@ -26,14 +30,14 @@ import (
 // Analyzer defines the entrypoint.
 var Analyzer = &analysis.Analyzer{
 	Name: "checkaligned",
-	Doc:  "prohibits direct use of 64 bit atomic operations",
+	Doc:  "prohibits direct use of atomic int operations",
 	Run:  run,
 }
 
 // blocklist lists prohibited identifiers in the atomic package.
 //
-// TODO(b/228378998): We should do this for 32 bit values too. Can also further
-// genericize this to ban other things we don't like (e.g. os.File).
+// TODO(b/228378998): We can further genericize this to ban other things we
+// don't like (e.g. os.File).
 var blocklist = []string{
 	"AddInt64",
 	"AddUint64",
@@ -45,12 +49,30 @@ var blocklist = []string{
 	"StoreUint64",
 	"SwapInt64",
 	"SwapUint64",
+
+	"AddInt32",
+	"AddUint32",
+	"CompareAndSwapInt32",
+	"CompareAndSwapUint32",
+	"LoadInt32",
+	"LoadUint32",
+	"StoreInt32",
+	"StoreUint32",
+	"SwapInt32",
+	"SwapUint32",
 }
 
 func run(pass *analysis.Pass) (interface{}, error) {
-	// atomicbitops uses 64 bit values safely.
-	if pass.Pkg.Name() == "atomicbitops" {
-		return nil, nil
+	// Check for the "// +checkalignedignore" escape hatch.
+	for _, file := range pass.Files {
+		if file.Doc == nil {
+			continue
+		}
+		for _, comment := range file.Doc.List {
+			if len(comment.Text) > 2 && strings.HasPrefix(comment.Text[2:], " +checkalignedignore") {
+				return nil, nil
+			}
+		}
 	}
 
 	for _, file := range pass.Files {