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

5 files changed, 103 insertions, 86 deletions

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"],