Add new locks with the correctness validator

All locks are separated into classes. The validator builds a dependency
graph and checks that it doesn't have cycles.

PiperOrigin-RevId: 447812244

pkg/sync/locking/BUILD

@@ -0,0 +1,136 @@
+load("//tools:defs.bzl", "go_library", "go_test")
+load("//tools/go_generics:defs.bzl", "go_template", "go_template_instance")
+load("//pkg/sync/locking:locking.bzl", "declare_mutex", "declare_rwmutex")
+
+package(
+    default_visibility = ["//:sandbox"],
+    licenses = ["notice"],
+)
+
+go_library(
+    name = "locking",
+    srcs = [
+        "atomicptrmap_ancestors_unsafe.go",
+        "atomicptrmap_class_unsafe.go",
+        "atomicptrmap_goroutine_unsafe.go",
+        "atomicptrmap_subclass_unsafe.go",
+        "lockdep.go",
+        "lockdep_norace.go",
+        "locking.go",
+    ],
+    marshal = False,
+    stateify = False,
+    visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/gohacks",
+        "//pkg/goid",
+        "//pkg/log",
+        "//pkg/sync",
+    ],
+)
+
+go_template_instance(
+    name = "atomicptrmap_goroutine",
+    out = "atomicptrmap_goroutine_unsafe.go",
+    package = "locking",
+    prefix = "goroutineLocks",
+    template = "//pkg/sync/atomicptrmap:generic_atomicptrmap",
+    types = {
+        "Key": "int64",
+        "Value": "goroutineLocks",
+    },
+)
+
+go_template_instance(
+    name = "atomicptrmap_class",
+    out = "atomicptrmap_class_unsafe.go",
+    imports = {
+        "reflect": "reflect",
+    },
+    package = "locking",
+    prefix = "class",
+    template = "//pkg/sync/atomicptrmap:generic_atomicptrmap",
+    types = {
+        "Key": "*MutexClass",
+        "Value": "reflect.Type",
+    },
+)
+
+go_template_instance(
+    name = "atomicptrmap_subclass",
+    out = "atomicptrmap_subclass_unsafe.go",
+    imports = {
+        "reflect": "reflect",
+    },
+    package = "locking",
+    prefix = "subclass",
+    template = "//pkg/sync/atomicptrmap:generic_atomicptrmap",
+    types = {
+        "Key": "uint32",
+        "Value": "MutexClass",
+    },
+)
+
+go_template_instance(
+    name = "atomicptrmap_ancestors",
+    out = "atomicptrmap_ancestors_unsafe.go",
+    imports = {
+        "reflect": "reflect",
+    },
+    package = "locking",
+    prefix = "ancestors",
+    template = "//pkg/sync/atomicptrmap:generic_atomicptrmap",
+    types = {
+        "Key": "*MutexClass",
+        "Value": "string",
+    },
+)
+
+go_template(
+    name = "generic_mutex",
+    srcs = ["generic_mutex.go"],
+    visibility = ["//:sandbox"],
+)
+
+go_template(
+    name = "generic_rwmutex",
+    srcs = ["generic_rwmutex.go"],
+    visibility = ["//:sandbox"],
+)
+
+declare_mutex(
+    name = "mutex_test",
+    out = "mutex_test.go",
+    package = "locking_test",
+    prefix = "test",
+)
+
+declare_rwmutex(
+    name = "mutex_test2",
+    out = "mutex_test2.go",
+    package = "locking_test",
+    prefix = "test2",
+)
+
+declare_mutex(
+    name = "mutex_test3",
+    out = "mutex_test3.go",
+    package = "locking_test",
+    prefix = "test3",
+)
+
+go_test(
+    name = "locking_test",
+    size = "small",
+    srcs = [
+        "lockdep_nolockdep_test.go",
+        "lockdep_test.go",
+        "mutex_test.go",
+        "mutex_test2.go",
+        "mutex_test3.go",
+    ],
+    deps = [
+        "//pkg/sync",
+        "//pkg/sync/locking",
+    ],
+)

pkg/sync/locking/generic_mutex.go

@@ -0,0 +1,61 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package locking
+
+import (
+	"reflect"
+
+	"gvisor.dev/gvisor/pkg/sync"
+	"gvisor.dev/gvisor/pkg/sync/locking"
+)
+
+// Mutex is sync.Mutex with the correctness validator.
+type Mutex struct {
+	mu sync.Mutex
+}
+
+// Lock locks m.
+// +checklocksignore
+func (m *Mutex) Lock() {
+	locking.AddGLock(genericMarkIndex, 0)
+	m.mu.Lock()
+}
+
+// NestedLock locks m knowing that another lock of the same type is held.
+// +checklocksignore
+func (m *Mutex) NestedLock() {
+	locking.AddGLock(genericMarkIndex, 1)
+	m.mu.Lock()
+}
+
+// Unlock unlocks m.
+// +checklocksignore
+func (m *Mutex) Unlock() {
+	locking.DelGLock(genericMarkIndex, 0)
+	m.mu.Unlock()
+}
+
+// NestedUnlock unlocks m knowing that another lock of the same type is held.
+// +checklocksignore
+func (m *Mutex) NestedUnlock() {
+	locking.DelGLock(genericMarkIndex, 1)
+	m.mu.Unlock()
+}
+
+var genericMarkIndex *locking.MutexClass
+
+func init() {
+	genericMarkIndex = locking.NewMutexClass(reflect.TypeOf(Mutex{}))
+}

pkg/sync/locking/generic_rwmutex.go

@@ -0,0 +1,93 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package locking
+
+import (
+	"reflect"
+
+	"gvisor.dev/gvisor/pkg/sync"
+	"gvisor.dev/gvisor/pkg/sync/locking"
+)
+
+// RWMutex is sync.RWMutex with the correctness validator.
+type RWMutex struct {
+	mu sync.RWMutex
+}
+
+// Lock locks m.
+// +checklocksignore
+func (m *RWMutex) Lock() {
+	locking.AddGLock(genericMarkIndex, 0)
+	m.mu.Lock()
+}
+
+// NestedLock locks m knowing that another lock of the same type is held.
+// +checklocksignore
+func (m *RWMutex) NestedLock() {
+	locking.AddGLock(genericMarkIndex, 1)
+	m.mu.Lock()
+}
+
+// Unlock unlocks m.
+// +checklocksignore
+func (m *RWMutex) Unlock() {
+	m.mu.Unlock()
+	locking.DelGLock(genericMarkIndex, 0)
+}
+
+// NestedUnlock unlocks m knowing that another lock of the same type is held.
+// +checklocksignore
+func (m *RWMutex) NestedUnlock() {
+	m.mu.Unlock()
+	locking.DelGLock(genericMarkIndex, 1)
+}
+
+// RLock locks m for reading.
+// +checklocksignore
+func (m *RWMutex) RLock() {
+	locking.AddGLock(genericMarkIndex, 0)
+	m.mu.RLock()
+}
+
+// RUnlock undoes a single RLock call.
+// +checklocksignore
+func (m *RWMutex) RUnlock() {
+	m.mu.RUnlock()
+	locking.DelGLock(genericMarkIndex, 0)
+}
+
+// RLockBypass locks m for reading without executing the validator.
+// +checklocksignore
+func (m *RWMutex) RLockBypass() {
+	m.mu.RLock()
+}
+
+// RUnlockBypass undoes a single RLockBypass call.
+// +checklocksignore
+func (m *RWMutex) RUnlockBypass() {
+	m.mu.RUnlock()
+}
+
+// DowngradeLock atomically unlocks rw for writing and locks it for reading.
+// +checklocksignore
+func (m *RWMutex) DowngradeLock() {
+	m.mu.DowngradeLock()
+}
+
+var genericMarkIndex *locking.MutexClass
+
+func init() {
+	genericMarkIndex = locking.NewMutexClass(reflect.TypeOf(RWMutex{}))
+}

pkg/sync/locking/lockdep.go

@@ -0,0 +1,136 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build lockdep
+// +build lockdep
+
+package locking
+
+import (
+	"fmt"
+	"reflect"
+	"strings"
+
+	"gvisor.dev/gvisor/pkg/goid"
+	"gvisor.dev/gvisor/pkg/log"
+)
+
+var classMap classAtomicPtrMap
+
+// NewMutexClass allocates a new mutex class.
+func NewMutexClass(t reflect.Type) *MutexClass {
+	c := &MutexClass{}
+	classMap.Store(c, &t)
+	return c
+}
+
+// MutexClass describes dependencies of a specific class.
+type MutexClass struct {
+	// ancestors are locks that are locked before the current class.
+	ancestors ancestorsAtomicPtrMap
+	// subclasses is the list of sub-classes that are used to handle nested locks.
+	subclasses subclassAtomicPtrMap
+}
+
+type goroutineLocks map[*MutexClass]bool
+
+var routineLocks goroutineLocksAtomicPtrMap
+
+// checkLock checks that class isn't in ancestors of prevClass.
+func checkLock(class *MutexClass, prevClass *MutexClass, chain []*MutexClass) {
+	chain = append(chain, prevClass)
+	if c := prevClass.ancestors.Load(class); c != nil {
+		var b strings.Builder
+		fmt.Fprintf(&b, "WARNING: circular locking detected: %s -> %s:\n%s\n",
+			*classMap.Load(chain[0]), *classMap.Load(class), log.Stacks(false))
+
+		fmt.Fprintf(&b, "known lock chain: ")
+		c := class
+		for i := len(chain) - 1; i >= 0; i-- {
+			fmt.Fprintf(&b, "%s -> ", *classMap.Load(c))
+			c = chain[i]
+		}
+		fmt.Fprintf(&b, "%s\n", *classMap.Load(chain[0]))
+		c = class
+		for i := len(chain) - 1; i >= 0; i-- {
+			fmt.Fprintf(&b, "\n====== %s -> %s =====\n%s",
+				*classMap.Load(c), *classMap.Load(chain[i]), *chain[i].ancestors.Load(c))
+			c = chain[i]
+		}
+		panic(b.String())
+	}
+	prevClass.ancestors.Range(func(parentClass *MutexClass, stacks *string) bool {
+		// The recursion is fine here. If it fails, you need to reduce
+		// a number of nested locks.
+		checkLock(class, parentClass, chain)
+		return true
+	})
+}
+
+// AddGLock records a lock to the current goroutine and updates dependences.
+func AddGLock(class *MutexClass, subclass uint32) {
+	gid := goid.Get()
+
+	if subclass != 0 {
+		var c *MutexClass
+		if c = class.subclasses.Load(subclass); c == nil {
+			t := classMap.Load(class)
+			c = NewMutexClass(*t)
+			class.subclasses.Store(subclass, c)
+		}
+		class = c
+	}
+	currentLocks := routineLocks.Load(gid)
+	if currentLocks == nil {
+		locks := goroutineLocks(make(map[*MutexClass]bool))
+		locks[class] = true
+		routineLocks.Store(gid, &locks)
+		return
+	}
+
+	// Check dependencies and add locked mutexes to the ancestors list.
+	for prevClass, _ := range *currentLocks {
+		if prevClass == class {
+			panic(fmt.Sprintf("nested locking: %s:\n%s", *classMap.Load(class), log.Stacks(false)))
+		}
+		checkLock(class, prevClass, nil)
+
+		if c := class.ancestors.Load(prevClass); c == nil {
+			stacks := string(log.Stacks(false))
+			class.ancestors.Store(prevClass, &stacks)
+		}
+	}
+	(*currentLocks)[class] = true
+
+}
+
+// DelGLock deletes a lock from the current goroutine.
+func DelGLock(class *MutexClass, subclass uint32) {
+	if subclass != 0 {
+		class = class.subclasses.Load(subclass)
+	}
+	gid := goid.Get()
+	currentLocks := routineLocks.Load(gid)
+	if currentLocks == nil {
+		panic("the current goroutine doesn't have locks")
+	}
+	if _, ok := (*currentLocks)[class]; !ok {
+		panic("unlock of an unknow lock")
+	}
+
+	delete(*currentLocks, class)
+	if len(*currentLocks) == 0 {
+		routineLocks.Store(gid, nil)
+	}
+}

pkg/sync/locking/lockdep_nolockdep_test.go

@@ -0,0 +1,38 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build !lockdep
+// +build !lockdep
+
+package locking_test
+
+import (
+	"testing"
+)
+
+// Ensure that lockdep is not enabled unless "lockdep" tag is set.
+func TestDummy(t *testing.T) {
+	m := testMutex{}
+	m2 := test2RWMutex{}
+	m.Lock()
+	m2.Lock()
+	t.Logf("m->m2")
+	m2.Unlock()
+	m.Unlock()
+	m2.Lock()
+	m.Lock()
+	t.Logf("m2->m")
+	m.Unlock()
+	m2.Unlock()
+}

pkg/sync/locking/lockdep_norace.go

@@ -0,0 +1,40 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build !lockdep
+// +build !lockdep
+
+package locking
+
+import (
+	"reflect"
+)
+
+type goroutineLocks map[*MutexClass]bool
+
+// MutexClass is a stub class without the lockdep tag.
+type MutexClass struct{}
+
+// NewMutexClass is no-op without the lockdep tag.
+func NewMutexClass(t reflect.Type) *MutexClass {
+	return nil
+}
+
+// AddGLock is no-op without the lockdep tag.
+//go:inline
+func AddGLock(class *MutexClass, subclass uint32) {}
+
+// DelGLock is no-op without the lockdep tag.
+//go:inline
+func DelGLock(class *MutexClass, subclass uint32) {}

pkg/sync/locking/lockdep_test.go

@@ -0,0 +1,106 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build lockdep
+// +build lockdep
+
+package locking_test
+
+import (
+	"testing"
+)
+
+func TestReverse(t *testing.T) {
+	m := testMutex{}
+	m2 := test2RWMutex{}
+	m.Lock()
+	m2.Lock()
+	m2.Unlock()
+	m.Unlock()
+
+	defer func() {
+		if r := recover(); r != nil {
+			t.Logf("Got expected panic: %s", r)
+		}
+	}()
+
+	m2.Lock()
+	m.Lock()
+	m.Unlock()
+	m2.Unlock()
+	t.Error("The reverse lock order hasn't been detected")
+}
+
+func TestIndirect(t *testing.T) {
+	m1 := testMutex{}
+	m2 := test2RWMutex{}
+	m3 := test3Mutex{}
+
+	m1.Lock()
+	m2.Lock()
+	m2.Unlock()
+	m1.Unlock()
+	m2.Lock()
+	m3.Lock()
+	m3.Unlock()
+	m2.Unlock()
+	defer func() {
+		if r := recover(); r != nil {
+			t.Logf("Got expected panic: %s", r)
+		}
+	}()
+
+	m3.Lock()
+	m1.Lock()
+	m1.Unlock()
+	m3.Unlock()
+	t.Error("The reverse lock order hasn't been detected")
+}
+
+func TestSame(t *testing.T) {
+	defer func() {
+		if r := recover(); r != nil {
+			t.Logf("Got expected panic: %s", r)
+		}
+	}()
+
+	m := testMutex{}
+	m.Lock()
+	m.Lock()
+	m.Unlock()
+	m.Unlock()
+	t.Error("The same lock has been locked twice, and was not detected.")
+}
+
+func TestReverseNested(t *testing.T) {
+	m1 := testMutex{}
+	m2 := testMutex{}
+	m1.Lock()
+	m2.NestedLock()
+	m1.Unlock()
+	m2.NestedUnlock()
+
+	defer func() {
+		if r := recover(); r != nil {
+			t.Logf("Got expected panic: %s", r)
+		}
+	}()
+
+	m2.NestedLock()
+	m1.Lock()
+	m1.NestedUnlock()
+	m2.Unlock()
+
+	t.Error("The reverse lock order hasn't been detected")
+}

pkg/sync/locking/locking.bzl

@@ -0,0 +1,27 @@
+"""Mutex-s rules."""
+
+load("//tools/go_generics:defs.bzl", "go_template_instance")
+
+def declare_mutex(package, name, out, prefix):
+    go_template_instance(
+        name = name,
+        out = out,
+        package = package,
+        prefix = prefix,
+        substrs = {
+            "genericMark": "prefix",
+        },
+        template = "//pkg/sync/locking:generic_mutex",
+    )
+
+def declare_rwmutex(package, name, out, prefix):
+    go_template_instance(
+        name = name,
+        out = out,
+        package = package,
+        prefix = prefix,
+        substrs = {
+            "genericMark": "prefix",
+        },
+        template = "//pkg/sync/locking:generic_rwmutex",
+    )

pkg/sync/locking/locking.go

@@ -0,0 +1,28 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package locking implements lock primitives with the correctness validator.
+//
+// All mutexes are divided on classes and the validator check following conditions:
+// * Mutexes of the same class are not taken more than once except cases when
+//   that is expected.
+// * Mutexes are never locked in a reverse order. Lock dependencies are tracked
+//   on the class level.
+//
+// The validator is implemented in a very straightforward way. For each mutex
+// class, we maintain the ancestors list of all classes that have ever been
+// taken before the target one. For each goroutine, we have the list of
+// currently locked mutexes. And finally, all lock methods check that
+// ancestors of currently locked mutexes don't contain the target one.
+package locking

runsc/BUILD

@@ -36,6 +36,7 @@ go_binary(
         "main.go",
         "version.go",
     ],
+    gotags = ["lockdep"],
     static = True,
     visibility = [
         "//visibility:public",

tools/bazel.mk

@@ -202,6 +202,9 @@ clean: ## Cleans the bazel cache.
 	@$(call clean)
 .PHONY: clean
 
+runsc-race:
+	@$(call build,--@io_bazel_rules_go//go/config:race runsc:runsc-race)
+
 testlogs: ## Returns the most recent set of test logs.
 	@if test -f .build_events.json; then \
 	  cat .build_events.json | jq -r \

tools/checklocks/analysis.go

@@ -474,11 +474,12 @@ func (pc *passContext) checkFunctionCall(call callCommon, fn *types.Func, lff *l
 
 	// Check if it's a method dispatch for something in the sync package.
 	// See: https://godoc.org/golang.org/x/tools/go/ssa#Function
-	if fn.Pkg() != nil && fn.Pkg().Name() == "sync" && len(args) > 0 {
+
+	if (lockerRE.MatchString(fn.FullName()) || mutexRE.MatchString(fn.FullName())) && len(args) > 0 {
 		rv := makeResolvedValue(args[0], nil)
 		isExclusive := false
 		switch fn.Name() {
-		case "Lock":
+		case "Lock", "NestedLock":
 			isExclusive = true
 			fallthrough
 		case "RLock":
@@ -488,7 +489,7 @@ func (pc *passContext) checkFunctionCall(call callCommon, fn *types.Func, lff *l
 					pc.maybeFail(call.Pos(), "%s already locked (locks: %s)", s, ls.String())
 				}
 			}
-		case "Unlock":
+		case "Unlock", "NestedUnlock":
 			isExclusive = true
 			fallthrough
 		case "RUnlock":

tools/checklocks/facts.go

@@ -469,9 +469,9 @@ func (pc *passContext) findField(structType *types.Struct, fieldName string) (fl
 }
 
 var (
-	mutexRE   = regexp.MustCompile("((.*/)|^)sync.(CrossGoroutineMutex|Mutex)")
-	rwMutexRE = regexp.MustCompile("((.*/)|^)sync.(CrossGoroutineRWMutex|RWMutex)")
-	lockerRE  = regexp.MustCompile("((.*/)|^)sync.Locker")
+	mutexRE   = regexp.MustCompile(".*Mutex")
+	rwMutexRE = regexp.MustCompile(".*RWMutex")
+	lockerRE  = regexp.MustCompile(".*sync.Locker")
 )
 
 // validateMutex validates the mutex type.
@@ -482,6 +482,9 @@ func (pc *passContext) validateMutex(pos token.Pos, obj types.Object, exclusive
 	// Check that it is indeed a mutex.
 	s := obj.Type().String()
 	switch {
+	case rwMutexRE.MatchString(s):
+		// Safe for all cases.
+		return true
 	case mutexRE.MatchString(s), lockerRE.MatchString(s):
 		// Safe for exclusive cases.
 		if !exclusive {
@@ -489,9 +492,6 @@ func (pc *passContext) validateMutex(pos token.Pos, obj types.Object, exclusive
 			return false
 		}
 		return true
-	case rwMutexRE.MatchString(s):
-		// Safe for all cases.
-		return true
 	default:
 		// Not a mutex at all?
 		pc.maybeFail(pos, "field %s is not a Mutex or an RWMutex", obj.Name())

tools/go_generics/defs.bzl

@@ -93,6 +93,7 @@ def _go_template_instance_impl(ctx):
     args += [("-t=%s=%s" % (p[0], p[1])) for p in ctx.attr.types.items()]
     args += [("-c=%s=%s" % (p[0], p[1])) for p in ctx.attr.consts.items()]
     args += [("-import=%s=%s" % (p[0], p[1])) for p in ctx.attr.imports.items()]
+    args += [("-s=%s=%s" % (p[0], p[1])) for p in ctx.attr.substrs.items()]
 
     if ctx.attr.anon:
         args.append("-anon")
@@ -119,6 +120,7 @@ go_template_instance = rule(
         "types": attr.string_dict(doc = "the map from generic type names to concrete ones"),
         "consts": attr.string_dict(doc = "the map from constant names to their values"),
         "imports": attr.string_dict(doc = "the map from imports used in types/consts to their import paths"),
+        "substrs": attr.string_dict(doc = "the map from sub-strings to their replacements"),
         "anon": attr.bool(doc = "whether anoymous fields should be processed", mandatory = False, default = False),
         "package": attr.string(doc = "the package for the generated source file", mandatory = False),
         "out": attr.output(doc = "output file", mandatory = True),

tools/go_generics/main.go

@@ -117,6 +117,7 @@ var (
 	types       = make(mapValue)
 	consts      = make(mapValue)
 	imports     = make(mapValue)
+	substr      = make(mapValue)
 )
 
 // mapValue implements flag.Value. We use a mapValue flag instead of a regular
@@ -165,6 +166,7 @@ func main() {
 	flag.Var(types, "t", "rename type A to B when `A=B` is passed in. Multiple such mappings are allowed.")
 	flag.Var(consts, "c", "reassign constant A to value B when `A=B` is passed in. Multiple such mappings are allowed.")
 	flag.Var(imports, "import", "specifies the import libraries to use when types are not local. `name=path` specifies that 'name', used in types as name.type, refers to the package living in 'path'.")
+	flag.Var(substr, "s", "replace sub-string A with B when `A=B` is passed in. Multiple such mappings are allowed.")
 	flag.Parse()
 
 	if *input == "" || *output == "" {
@@ -279,7 +281,12 @@ func main() {
 		os.Exit(1)
 	}
 
-	if err := ioutil.WriteFile(*output, buf.Bytes(), 0644); err != nil {
+	byteBuf := buf.Bytes()
+	for old, new := range substr {
+		byteBuf = bytes.ReplaceAll(byteBuf, []byte(old), []byte(new))
+	}
+
+	if err := ioutil.WriteFile(*output, byteBuf, 0644); err != nil {
 		fmt.Fprintf(os.Stderr, "%v\n", err)
 		os.Exit(1)
 	}