Fix Go branch for arm64.

This requires several changes:
* Templates must preserve relevant tags.
* Pagetables templates are split into two targets, each preserving tags.
* The binary VDSO is similarly split into two targets, with some juggling.
* The top level tools/go_branch.sh now does a crossbuild of ARM64 as well,
  and checks and merges the results of the two branches together.

Fixes #5178

PiperOrigin-RevId: 351304330

12 files changed, 204 insertions, 122 deletions

diff --git a/.buildkite/pipeline.yaml b/.buildkite/pipeline.yaml
index b9c392ee9..f8bf00b50 100644
--- a/.buildkite/pipeline.yaml
+++ b/.buildkite/pipeline.yaml
@@ -33,7 +33,7 @@ steps:
   - <<: *common
     label: ":golang: Go branch"
     commands:
-      - make go
+      - tools/go_branch.sh
       - git checkout go && git clean -f
       - go build ./...
 
diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml
index e62991691..802fe5ce5 100644
--- a/.github/workflows/go.yml
+++ b/.github/workflows/go.yml
@@ -43,7 +43,7 @@ jobs:
         key: ${{ runner.os }}-bazel-${{ hashFiles('WORKSPACE') }}
         restore-keys: |
           ${{ runner.os }}-bazel-
-    - run: make go
+    - run: tools/go_branch.sh
     - run: |
         git remote add upstream "https://github.com/${{ github.repository }}"
         git push upstream go:go
diff --git a/Makefile b/Makefile
index 549093ffa..92f277278 100644
--- a/Makefile
+++ b/Makefile
@@ -151,11 +151,6 @@ nogo: ## Surfaces all nogo findings.
 	@$(call run,//tools/github $(foreach dir,$(BUILD_ROOTS),-path=$(CURDIR)/$(dir)) -dry-run nogo)
 .PHONY: nogo
 
-go: ## Builds the Go branch.
-	@$(call clean)
-	@$(call build,//:gopath)
-	@tools/go_branch.sh
-
 gazelle: ## Runs gazelle to update WORKSPACE.
 	@$(call run,//:gazelle update-repos -from_file=go.mod -prune)
 .PHONY: gazelle
diff --git a/pkg/sentry/loader/BUILD b/pkg/sentry/loader/BUILD
index 34bdb0b69..ab074b400 100644
--- a/pkg/sentry/loader/BUILD
+++ b/pkg/sentry/loader/BUILD
@@ -1,14 +1,7 @@
-load("//tools:defs.bzl", "go_embed_data", "go_library")
+load("//tools:defs.bzl", "go_library")
 
 package(licenses = ["notice"])
 
-go_embed_data(
-    name = "vdso_bin",
-    src = "//vdso:vdso.so",
-    package = "loader",
-    var = "vdsoBin",
-)
-
 go_library(
     name = "loader",
     srcs = [
@@ -17,7 +10,6 @@ go_library(
         "loader.go",
         "vdso.go",
         "vdso_state.go",
-        ":vdso_bin",
     ],
     visibility = ["//pkg/sentry:internal"],
     deps = [
@@ -33,6 +25,7 @@ go_library(
         "//pkg/sentry/fsbridge",
         "//pkg/sentry/kernel/auth",
         "//pkg/sentry/limits",
+        "//pkg/sentry/loader/vdsodata",
         "//pkg/sentry/memmap",
         "//pkg/sentry/mm",
         "//pkg/sentry/pgalloc",
diff --git a/pkg/sentry/loader/vdso.go b/pkg/sentry/loader/vdso.go
index 241d87835..a32d37d62 100644
--- a/pkg/sentry/loader/vdso.go
+++ b/pkg/sentry/loader/vdso.go
@@ -26,6 +26,7 @@ import (
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/safemem"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
+	"gvisor.dev/gvisor/pkg/sentry/loader/vdsodata"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
 	"gvisor.dev/gvisor/pkg/sentry/mm"
 	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
@@ -177,7 +178,7 @@ type VDSO struct {
 
 // getSymbolValueFromVDSO returns the specific symbol value in vdso.so.
 func getSymbolValueFromVDSO(symbol string) (uint64, error) {
-	f, err := elf.NewFile(bytes.NewReader(vdsoBin))
+	f, err := elf.NewFile(bytes.NewReader(vdsodata.Binary))
 	if err != nil {
 		return 0, err
 	}
@@ -199,19 +200,19 @@ func getSymbolValueFromVDSO(symbol string) (uint64, error) {
 // PrepareVDSO validates the system VDSO and returns a VDSO, containing the
 // param page for updating by the kernel.
 func PrepareVDSO(mfp pgalloc.MemoryFileProvider) (*VDSO, error) {
-	vdsoFile := &byteFullReader{data: vdsoBin}
+	vdsoFile := &byteFullReader{data: vdsodata.Binary}
 
 	// First make sure the VDSO is valid. vdsoFile does not use ctx, so a
 	// nil context can be passed.
-	info, err := validateVDSO(nil, vdsoFile, uint64(len(vdsoBin)))
+	info, err := validateVDSO(nil, vdsoFile, uint64(len(vdsodata.Binary)))
 	if err != nil {
 		return nil, err
 	}
 
 	// Then copy it into a VDSO mapping.
-	size, ok := usermem.Addr(len(vdsoBin)).RoundUp()
+	size, ok := usermem.Addr(len(vdsodata.Binary)).RoundUp()
 	if !ok {
-		return nil, fmt.Errorf("VDSO size overflows? %#x", len(vdsoBin))
+		return nil, fmt.Errorf("VDSO size overflows? %#x", len(vdsodata.Binary))
 	}
 
 	mf := mfp.MemoryFile()
@@ -226,7 +227,7 @@ func PrepareVDSO(mfp pgalloc.MemoryFileProvider) (*VDSO, error) {
 		return nil, fmt.Errorf("unable to map VDSO memory: %v", err)
 	}
 
-	_, err = safemem.CopySeq(ims, safemem.BlockSeqOf(safemem.BlockFromSafeSlice(vdsoBin)))
+	_, err = safemem.CopySeq(ims, safemem.BlockSeqOf(safemem.BlockFromSafeSlice(vdsodata.Binary)))
 	if err != nil {
 		mf.DecRef(vdso)
 		return nil, fmt.Errorf("unable to copy VDSO into memory: %v", err)
diff --git a/pkg/sentry/loader/vdsodata/BUILD b/pkg/sentry/loader/vdsodata/BUILD
new file mode 100644
index 000000000..119199f97
--- /dev/null
+++ b/pkg/sentry/loader/vdsodata/BUILD
@@ -0,0 +1,38 @@
+load("//tools:defs.bzl", "go_add_tags", "go_embed_data", "go_library")
+
+package(licenses = ["notice"])
+
+go_embed_data(
+    name = "vdso_bin",
+    src = "//vdso:vdso.so",
+    package = "vdsodata",
+    var = "Binary",
+)
+
+[
+    # Generate multiple tagged files. Note that the contents of all files
+    # will be the same (i.e. vdso_arm64.go will contain the amd64 vdso), but
+    # the build tags will ensure only one is selected. When we generate the
+    # "Go" branch, we select all archiecture files from the relevant build.
+    # This is a hack around some limitations for "out" being a configurable
+    # attribute and selects for srcs. See also tools/go_branch.sh.
+    go_add_tags(
+        name = "vdso_%s" % arch,
+        src = ":vdso_bin",
+        out = "vdso_%s.go" % arch,
+        go_tags = [arch],
+    )
+    for arch in ("amd64", "arm64")
+]
+
+go_library(
+    name = "vdsodata",
+    srcs = [
+        "vdsodata.go",
+        ":vdso_amd64",
+        ":vdso_arm64",
+    ],
+    marshal = False,
+    stateify = False,
+    visibility = ["//pkg/sentry:internal"],
+)
diff --git a/pkg/sentry/loader/vdsodata/vdsodata.go b/pkg/sentry/loader/vdsodata/vdsodata.go
new file mode 100644
index 000000000..a6dec3b48
--- /dev/null
+++ b/pkg/sentry/loader/vdsodata/vdsodata.go
@@ -0,0 +1,16 @@
+// Copyright 2020 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 vdsodata contains a compiled VDSO object.
+package vdsodata
diff --git a/pkg/sentry/platform/ring0/pagetables/BUILD b/pkg/sentry/platform/ring0/pagetables/BUILD
index 16d5f478b..9e3539e4c 100644
--- a/pkg/sentry/platform/ring0/pagetables/BUILD
+++ b/pkg/sentry/platform/ring0/pagetables/BUILD
@@ -1,74 +1,38 @@
-load("//tools:defs.bzl", "go_library", "go_test", "select_arch")
+load("//tools:defs.bzl", "go_library", "go_test")
 load("//tools/go_generics:defs.bzl", "go_template", "go_template_instance")
 
 package(licenses = ["notice"])
 
-go_template(
-    name = "generic_walker",
-    srcs = select_arch(
-        amd64 = ["walker_amd64.go"],
-        arm64 = ["walker_arm64.go"],
-    ),
-    opt_types = [
-        "Visitor",
-    ],
-    visibility = [":__pkg__"],
-)
-
-go_template_instance(
-    name = "walker_map",
-    out = "walker_map.go",
-    package = "pagetables",
-    prefix = "map",
-    template = ":generic_walker",
-    types = {
-        "Visitor": "mapVisitor",
-    },
-)
-
-go_template_instance(
-    name = "walker_unmap",
-    out = "walker_unmap.go",
-    package = "pagetables",
-    prefix = "unmap",
-    template = ":generic_walker",
-    types = {
-        "Visitor": "unmapVisitor",
-    },
-)
+[
+    # These files are tagged with relevant build architectures. We can always
+    # build all the input files, which will be included only in the relevant
+    # architecture builds.
+    go_template(
+        name = "generic_walker_%s" % arch,
+        srcs = ["walker_%s.go" % arch],
+        opt_types = [
+            "Visitor",
+        ],
+        visibility = [":__pkg__"],
+    )
+    for arch in ("amd64", "arm64")
+]
 
-go_template_instance(
-    name = "walker_lookup",
-    out = "walker_lookup.go",
-    package = "pagetables",
-    prefix = "lookup",
-    template = ":generic_walker",
-    types = {
-        "Visitor": "lookupVisitor",
-    },
-)
-
-go_template_instance(
-    name = "walker_empty",
-    out = "walker_empty.go",
-    package = "pagetables",
-    prefix = "empty",
-    template = ":generic_walker",
-    types = {
-        "Visitor": "emptyVisitor",
-    },
-)
-
-go_template_instance(
-    name = "walker_check",
-    out = "walker_check.go",
-    package = "pagetables",
-    prefix = "check",
-    template = ":generic_walker",
-    types = {
-        "Visitor": "checkVisitor",
-    },
-)
+[
+    # See above.
+    go_template_instance(
+        name = "walker_%s_%s" % (op, arch),
+        out = "walker_%s_%s.go" % (op, arch),
+        package = "pagetables",
+        prefix = op,
+        template = ":generic_walker_%s" % arch,
+        types = {
+            "Visitor": "%sVisitor" % op,
+        },
+    )
+    for op in ("map", "unmap", "lookup", "empty", "check")
+    for arch in ("amd64", "arm64")
+]
 
 go_library(
     name = "pagetables",
@@ -86,10 +50,14 @@ go_library(
         "pcids_x86.go",
         "walker_amd64.go",
         "walker_arm64.go",
-        "walker_empty.go",
-        "walker_lookup.go",
-        "walker_map.go",
-        "walker_unmap.go",
+        ":walker_empty_amd64",
+        ":walker_empty_arm64",
+        ":walker_lookup_amd64",
+        ":walker_lookup_arm64",
+        ":walker_map_amd64",
+        ":walker_map_arm64",
+        ":walker_unmap_amd64",
+        ":walker_unmap_arm64",
     ],
     visibility = [
         "//pkg/sentry/platform/kvm:__subpackages__",
@@ -108,7 +76,8 @@ go_test(
         "pagetables_amd64_test.go",
         "pagetables_arm64_test.go",
         "pagetables_test.go",
-        "walker_check.go",
+        ":walker_check_amd64",
+        ":walker_check_arm64",
     ],
     library = ":pagetables",
     deps = ["//pkg/usermem"],
diff --git a/tools/defs.bzl b/tools/defs.bzl
index 4be767432..d2c697c0d 100644
--- a/tools/defs.bzl
+++ b/tools/defs.bzl
@@ -10,7 +10,7 @@ load("//tools/go_marshal:defs.bzl", "go_marshal", "marshal_deps", "marshal_test_
 load("//tools/nogo:defs.bzl", "nogo_test")
 load("//tools/bazeldefs:defs.bzl", _arch_genrule = "arch_genrule", _build_test = "build_test", _bzl_library = "bzl_library", _coreutil = "coreutil", _default_installer = "default_installer", _default_net_util = "default_net_util", _more_shards = "more_shards", _most_shards = "most_shards", _proto_library = "proto_library", _select_arch = "select_arch", _select_system = "select_system", _short_path = "short_path", _version = "version")
 load("//tools/bazeldefs:cc.bzl", _cc_binary = "cc_binary", _cc_flags_supplier = "cc_flags_supplier", _cc_grpc_library = "cc_grpc_library", _cc_library = "cc_library", _cc_proto_library = "cc_proto_library", _cc_test = "cc_test", _cc_toolchain = "cc_toolchain", _gbenchmark = "gbenchmark", _grpcpp = "grpcpp", _gtest = "gtest", _vdso_linker_option = "vdso_linker_option")
-load("//tools/bazeldefs:go.bzl", _gazelle = "gazelle", _go_binary = "go_binary", _go_embed_data = "go_embed_data", _go_grpc_and_proto_libraries = "go_grpc_and_proto_libraries", _go_library = "go_library", _go_path = "go_path", _go_proto_library = "go_proto_library", _go_test = "go_test", _select_goarch = "select_goarch", _select_goos = "select_goos")
+load("//tools/bazeldefs:go.bzl", _gazelle = "gazelle", _go_binary = "go_binary", _go_embed_data = "go_embed_data", _go_grpc_and_proto_libraries = "go_grpc_and_proto_libraries", _go_library = "go_library", _go_path = "go_path", _go_proto_library = "go_proto_library", _go_rule = "go_rule", _go_test = "go_test", _select_goarch = "select_goarch", _select_goos = "select_goos")
 load("//tools/bazeldefs:pkg.bzl", _pkg_deb = "pkg_deb", _pkg_tar = "pkg_tar")
 load("//tools/bazeldefs:platforms.bzl", _default_platform = "default_platform", _platforms = "platforms")
 load("//tools/bazeldefs:tags.bzl", "go_suffixes")
@@ -43,10 +43,10 @@ vdso_linker_option = _vdso_linker_option
 
 # Go rules.
 gazelle = _gazelle
-go_embed_data = _go_embed_data
 go_path = _go_path
 select_goos = _select_goos
 select_goarch = _select_goarch
+go_embed_data = _go_embed_data
 
 # Packaging rules.
 pkg_deb = _pkg_deb
@@ -56,6 +56,35 @@ pkg_tar = _pkg_tar
 default_platform = _default_platform
 platforms = _platforms
 
+def _go_add_tags(ctx):
+    """ Adds tags to the given source file. """
+    output = ctx.outputs.out
+    runner = ctx.actions.declare_file(ctx.label.name + ".sh")
+    lines = ["#!/bin/bash"]
+    lines += ["echo '// +build %s' >> %s" % (tag, output.path) for tag in ctx.attr.go_tags]
+    lines.append("echo '' >> %s" % output.path)
+    lines += ["cat %s >> %s" % (f.path, output.path) for f in ctx.files.src]
+    lines.append("")
+    ctx.actions.write(runner, "\n".join(lines), is_executable = True)
+    ctx.actions.run(
+        inputs = ctx.files.src,
+        outputs = [output],
+        executable = runner,
+    )
+    return [DefaultInfo(
+        files = depset([output]),
+    )]
+
+go_add_tags = _go_rule(
+    rule,
+    implementation = _go_add_tags,
+    attrs = {
+        "go_tags": attr.string_list(doc = "Go build tags to be added.", mandatory = True),
+        "src": attr.label(doc = "Source file.", allow_single_file = True, mandatory = True),
+        "out": attr.output(doc = "Output file.", mandatory = True),
+    },
+)
+
 def go_binary(name, nogo = True, pure = False, static = False, x_defs = None, **kwargs):
     """Wraps the standard go_binary.
 
diff --git a/tools/go_branch.sh b/tools/go_branch.sh
index 7ef4ddf83..3a6a83f2e 100755
--- a/tools/go_branch.sh
+++ b/tools/go_branch.sh
@@ -16,23 +16,6 @@
 
 set -xeou pipefail
 
-# Discovery the package name from the go.mod file.
-declare module origpwd othersrc
-module=$(cat go.mod | grep -E "^module" | cut -d' ' -f2)
-origpwd=$(pwd)
-othersrc=("go.mod" "go.sum" "AUTHORS" "LICENSE")
-readonly module origpwd othersrc
-
-
-# Check that gopath has been built.
-declare gopath_dir
-gopath_dir="$(pwd)/bazel-bin/gopath/src/${module}"
-readonly gopath_dir
-if ! [[ -d "${gopath_dir}" ]]; then
-  echo "No gopath directory found; build the :gopath target." >&2
-  exit 1
-fi
-
 # Create a temporary working directory, and ensure that this directory and all
 # subdirectories are cleaned up upon exit.
 declare tmp_dir
@@ -44,6 +27,51 @@ finish() {
 }
 trap finish EXIT
 
+# Discover the package name from the go.mod file.
+declare module origpwd othersrc
+module=$(cat go.mod | grep -E "^module" | cut -d' ' -f2)
+origpwd=$(pwd)
+othersrc=("go.mod" "go.sum" "AUTHORS" "LICENSE")
+readonly module origpwd othersrc
+
+# Build an amd64 & arm64 gopath.
+declare -r go_amd64="${tmp_dir}/amd64"
+declare -r go_arm64="${tmp_dir}/arm64"
+rm -rf bazel-bin/gopath
+make build BAZEL_OPTIONS="" TARGETS="//:gopath"
+rsync --recursive --delete --copy-links bazel-bin/gopath/ "${go_amd64}"
+rm -rf bazel-bin/gopath
+make build BAZEL_OPTIONS=--config=cross-aarch64 TARGETS="//:gopath"
+rsync --recursive --delete --copy-links bazel-bin/gopath/ "${go_arm64}"
+
+# Strip irrelevant files, i.e. use only arm64 files from the arm64 build.
+# This is because bazel may generate incorrect files for non-target platforms
+# as a workaround. See pkg/sentry/loader/vdsodata as an example.
+find "${go_amd64}/src/${module}" -name '*_arm64*.go' -exec rm -f {} \;
+find "${go_amd64}/src/${module}" -name '*_arm64*.s' -exec rm -f {} \;
+find "${go_arm64}/src/${module}" -name '*_amd64*.go' -exec rm -f {} \;
+find "${go_arm64}/src/${module}" -name '*_amd64*.s' -exec rm -f {} \;
+
+# See below. The certs.go file is pseudo-random, and therefore will also
+# differ between the branches. Since we merge, it only has to come from one.
+# We arbitrarily keep the one from the amd64 branch, and drop the arm64 one.
+rm -f "${go_arm64}/src/${module}/webhook/pkg/injector/certs.go"
+
+# Check that all files are compatible. This means that if the files exist in
+# both architectures, then they must be identical. The only ones that we expect
+# to exist in a single architecture (due to binary builds) may be different.
+function cross_check() {
+  (cd "${1}" && find "src/${module}" -type f | \
+    xargs -n 1 -I {} sh -c "diff '${1}/{}' '${2}/{}' 2>/dev/null; test \$? -ne 1")
+}
+cross_check "${go_arm64}" "${go_amd64}"
+cross_check "${go_amd64}" "${go_arm64}"
+
+# Merge the two for a complete set of source files.
+declare -r go_merged="${tmp_dir}/merged"
+rsync --recursive --update "${go_amd64}/" "${go_merged}"
+rsync --recursive --update "${go_arm64}/" "${go_merged}"
+
 # Record the current working commit.
 declare head
 head=$(git describe --always)
@@ -89,14 +117,15 @@ git merge --no-commit --strategy ours "${head}" || \
 find . -type f -exec chmod 0644 {} \;
 find . -type d -exec chmod 0755 {} \;
 
-# Sync the entire gopath_dir. Note that we exclude auto-generated source
-# files that will change here. Otherwise, it adds a tremendous amount of noise
-# to commits. If this file disappears in the future, then presumably we will
-# still delete the underlying directory.
+# Sync the entire gopath. Note that we exclude auto-generated source files that
+# will change here. Otherwise, it adds a tremendous amount of noise to commits.
+# If this file disappears in the future, then presumably we will still delete
+# the underlying directory.
+declare -r gopath="${go_merged}/src/${module}"
 rsync --recursive --delete \
   --exclude .git \
   --exclude webhook/pkg/injector/certs.go \
-  -L "${gopath_dir}/" .
+  "${gopath}/" .
 
 # Add additional files.
 for file in "${othersrc[@]}"; do
diff --git a/tools/go_generics/go_merge/BUILD b/tools/go_generics/go_merge/BUILD
index 2fd5a200d..5e0487e93 100644
--- a/tools/go_generics/go_merge/BUILD
+++ b/tools/go_generics/go_merge/BUILD
@@ -6,4 +6,7 @@ go_binary(
     name = "go_merge",
     srcs = ["main.go"],
     visibility = ["//:sandbox"],
+    deps = [
+        "//tools/tags",
+    ],
 )
diff --git a/tools/go_generics/go_merge/main.go b/tools/go_generics/go_merge/main.go
index e0345500f..801f2354f 100644
--- a/tools/go_generics/go_merge/main.go
+++ b/tools/go_generics/go_merge/main.go
@@ -22,10 +22,12 @@ import (
 	"go/format"
 	"go/parser"
 	"go/token"
-	"io/ioutil"
 	"os"
 	"path/filepath"
 	"strconv"
+	"strings"
+
+	"gvisor.dev/gvisor/tools/tags"
 )
 
 var (
@@ -132,10 +134,17 @@ func main() {
 	// Write the output file.
 	var buf bytes.Buffer
 	if err := format.Node(&buf, fset, f); err != nil {
-		fatalf("%v\n", err)
+		fatalf("fomatting: %v\n", err)
 	}
-
-	if err := ioutil.WriteFile(*output, buf.Bytes(), 0644); err != nil {
-		fatalf("%v\n", err)
+	outf, err := os.OpenFile(*output, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
+	if err != nil {
+		fatalf("opening output: %v\n", err)
+	}
+	defer outf.Close()
+	if t := tags.Aggregate(flag.Args()); len(t) > 0 {
+		fmt.Fprintf(outf, "%s\n\n", strings.Join(t.Lines(), "\n"))
+	}
+	if _, err := outf.Write(buf.Bytes()); err != nil {
+		fatalf("write: %v\n", err)
 	}
 }