Merge branch 'master' into ip-forwarding

- Merges aleksej-paschenko's with HEAD
- Adds vfs2 support for ip_forward

3 files changed, 145 insertions, 156 deletions

diff --git a/tools/go_stateify/BUILD b/tools/go_stateify/BUILD
index bb53f8ae9..913558b4e 100644
--- a/tools/go_stateify/BUILD
+++ b/tools/go_stateify/BUILD
@@ -1,9 +1,16 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_binary")
+load("//tools:defs.bzl", "bzl_library", "go_binary")
 
 package(licenses = ["notice"])
 
 go_binary(
     name = "stateify",
     srcs = ["main.go"],
-    visibility = ["//visibility:public"],
+    visibility = ["//:sandbox"],
+    deps = ["//tools/tags"],
+)
+
+bzl_library(
+    name = "defs_bzl",
+    srcs = ["defs.bzl"],
+    visibility = ["//visibility:private"],
 )
diff --git a/tools/go_stateify/defs.bzl b/tools/go_stateify/defs.bzl
index 3ce36c1c8..6a5e666f0 100644
--- a/tools/go_stateify/defs.bzl
+++ b/tools/go_stateify/defs.bzl
@@ -1,41 +1,4 @@
-"""Stateify is a tool for generating state wrappers for Go types.
-
-The recommended way is to use the go_library rule defined below with mostly
-identical configuration as the native go_library rule.
-
-load("//tools/go_stateify:defs.bzl", "go_library")
-
-go_library(
-    name = "foo",
-    srcs = ["foo.go"],
-)
-
-Under the hood, the go_stateify rule is used to generate a file that will
-appear in a Go target; the output file should appear explicitly in a srcs list.
-For example (the above is still the preferred way):
-
-load("//tools/go_stateify:defs.bzl", "go_stateify")
-
-go_stateify(
-    name = "foo_state",
-    srcs = ["foo.go"],
-    out = "foo_state.go",
-    package = "foo",
-)
-
-go_library(
-    name = "foo",
-    srcs = [
-        "foo.go",
-        "foo_state.go",
-    ],
-    deps = [
-        "//pkg/state",
-    ],
-)
-"""
-
-load("@io_bazel_rules_go//go:def.bzl", _go_library = "go_library")
+"""Stateify is a tool for generating state wrappers for Go types."""
 
 def _go_stateify_impl(ctx):
     """Implementation for the stateify tool."""
@@ -43,12 +6,12 @@ def _go_stateify_impl(ctx):
 
     # Run the stateify command.
     args = ["-output=%s" % output.path]
-    args += ["-pkg=%s" % ctx.attr.package]
+    args.append("-fullpkg=%s" % ctx.attr.package)
     if ctx.attr._statepkg:
-        args += ["-statepkg=%s" % ctx.attr._statepkg]
+        args.append("-statepkg=%s" % ctx.attr._statepkg)
     if ctx.attr.imports:
-        args += ["-imports=%s" % ",".join(ctx.attr.imports)]
-    args += ["--"]
+        args.append("-imports=%s" % ",".join(ctx.attr.imports))
+    args.append("--")
     for src in ctx.attr.srcs:
         args += [f.path for f in src.files.to_list()]
     ctx.actions.run(
@@ -80,14 +43,11 @@ for statified types.
             mandatory = False,
         ),
         "package": attr.string(
-            doc = "The package name for the input sources.",
+            doc = "The fully qualified package name for the input sources.",
             mandatory = True,
         ),
         "out": attr.output(
-            doc = """
-The name of the generated file output. This must not conflict with any other
-files and must be added to the srcs of the relevant go_library.
-""",
+            doc = "Name of the generator output file.",
             mandatory = True,
         ),
         "_tool": attr.label(
@@ -98,39 +58,3 @@ files and must be added to the srcs of the relevant go_library.
         "_statepkg": attr.string(default = "gvisor.dev/gvisor/pkg/state"),
     },
 )
-
-def go_library(name, srcs, deps = [], imports = [], **kwargs):
-    """Standard go_library wrapped which generates state source files.
-
-    Args:
-      name: the name of the go_library rule.
-      srcs: sources of the go_library. Each will be processed for stateify
-            annotations.
-      deps: dependencies for the go_library.
-      imports: an optional list of extra non-aliased, Go-style absolute import
-               paths required for stateified types.
-      **kwargs: passed to go_library.
-    """
-    if "encode_unsafe.go" not in srcs and (name + "_state_autogen.go") not in srcs:
-        # Only do stateification for non-state packages without manual autogen.
-        go_stateify(
-            name = name + "_state_autogen",
-            srcs = [src for src in srcs if src.endswith(".go")],
-            imports = imports,
-            package = name,
-            out = name + "_state_autogen.go",
-        )
-        all_srcs = srcs + [name + "_state_autogen.go"]
-        if "//pkg/state" not in deps:
-            all_deps = deps + ["//pkg/state"]
-        else:
-            all_deps = deps
-    else:
-        all_deps = deps
-        all_srcs = srcs
-    _go_library(
-        name = name,
-        srcs = all_srcs,
-        deps = all_deps,
-        **kwargs
-    )
diff --git a/tools/go_stateify/main.go b/tools/go_stateify/main.go
index db7a7107b..4f6ed208a 100644
--- a/tools/go_stateify/main.go
+++ b/tools/go_stateify/main.go
@@ -23,13 +23,16 @@ import (
 	"go/parser"
 	"go/token"
 	"os"
+	"path/filepath"
 	"reflect"
 	"strings"
 	"sync"
+
+	"gvisor.dev/gvisor/tools/tags"
 )
 
 var (
-	pkg      = flag.String("pkg", "", "output package")
+	fullPkg  = flag.String("fullpkg", "", "fully qualified output package")
 	imports  = flag.String("imports", "", "extra imports for the output file")
 	output   = flag.String("output", "", "output file")
 	statePkg = flag.String("statepkg", "", "state import package; defaults to empty")
@@ -100,7 +103,7 @@ type scanFunctions struct {
 // skipped if nil.
 //
 // Fields tagged nosave are skipped.
-func scanFields(ss *ast.StructType, fn scanFunctions) {
+func scanFields(ss *ast.StructType, prefix string, fn scanFunctions) {
 	if ss.Fields.List == nil {
 		// No fields.
 		return
@@ -124,7 +127,16 @@ func scanFields(ss *ast.StructType, fn scanFunctions) {
 			continue
 		}
 
-		switch tag := extractStateTag(field.Tag); tag {
+		// Is this a anonymous struct? If yes, then continue the
+		// recursion with the given prefix. We don't pay attention to
+		// any tags on the top-level struct field.
+		tag := extractStateTag(field.Tag)
+		if anon, ok := field.Type.(*ast.StructType); ok && tag == "" {
+			scanFields(anon, name+".", fn)
+			continue
+		}
+
+		switch tag {
 		case "zerovalue":
 			if fn.zerovalue != nil {
 				fn.zerovalue(name)
@@ -168,7 +180,7 @@ func main() {
 		flag.Usage()
 		os.Exit(1)
 	}
-	if *pkg == "" {
+	if *fullPkg == "" {
 		fmt.Fprintf(os.Stderr, "Error: package required.")
 		os.Exit(1)
 	}
@@ -198,33 +210,25 @@ func main() {
 	// initCalls is dumped at the end.
 	var initCalls []string
 
-	// Declare our emission closures.
+	// Common closures.
 	emitRegister := func(name string) {
-		initCalls = append(initCalls, fmt.Sprintf("%sRegister(\"%s.%s\", (*%s)(nil), state.Fns{Save: (*%s).save, Load: (*%s).load})", statePrefix, *pkg, name, name, name, name))
+		initCalls = append(initCalls, fmt.Sprintf("%sRegister((*%s)(nil))", statePrefix, name))
 	}
 	emitZeroCheck := func(name string) {
-		fmt.Fprintf(outputFile, "	if !%sIsZeroValue(x.%s) { m.Failf(\"%s is %%v, expected zero\", x.%s) }\n", statePrefix, name, name, name)
-	}
-	emitLoadValue := func(name, typName string) {
-		fmt.Fprintf(outputFile, "	m.LoadValue(\"%s\", new(%s), func(y interface{}) { x.load%s(y.(%s)) })\n", name, typName, camelCased(name), typName)
-	}
-	emitLoad := func(name string) {
-		fmt.Fprintf(outputFile, "	m.Load(\"%s\", &x.%s)\n", name, name)
-	}
-	emitLoadWait := func(name string) {
-		fmt.Fprintf(outputFile, "	m.LoadWait(\"%s\", &x.%s)\n", name, name)
-	}
-	emitSaveValue := func(name, typName string) {
-		fmt.Fprintf(outputFile, "	var %s %s = x.save%s()\n", name, typName, camelCased(name))
-		fmt.Fprintf(outputFile, "	m.SaveValue(\"%s\", %s)\n", name, name)
+		fmt.Fprintf(outputFile, "	if !%sIsZeroValue(&x.%s) { %sFailf(\"%s is %%#v, expected zero\", &x.%s) }\n", statePrefix, name, statePrefix, name, name)
 	}
-	emitSave := func(name string) {
-		fmt.Fprintf(outputFile, "	m.Save(\"%s\", &x.%s)\n", name, name)
+
+	// Automated warning.
+	fmt.Fprint(outputFile, "// automatically generated by stateify.\n\n")
+
+	// Emit build tags.
+	if t := tags.Aggregate(flag.Args()); len(t) > 0 {
+		fmt.Fprintf(outputFile, "%s\n\n", strings.Join(t.Lines(), "\n"))
 	}
 
 	// Emit the package name.
-	fmt.Fprint(outputFile, "// automatically generated by stateify.\n\n")
-	fmt.Fprintf(outputFile, "package %s\n\n", *pkg)
+	_, pkg := filepath.Split(*fullPkg)
+	fmt.Fprintf(outputFile, "package %s\n\n", pkg)
 
 	// Emit the imports lazily.
 	var once sync.Once
@@ -256,6 +260,7 @@ func main() {
 			fmt.Fprintf(os.Stderr, "Input %q can't be parsed: %v\n", filename, err)
 			os.Exit(1)
 		}
+
 		files = append(files, f)
 	}
 
@@ -317,87 +322,140 @@ func main() {
 				continue
 			}
 
-			// Only generate code for types marked
-			// "// +stateify savable" in one of the proceeding
-			// comment lines.
+			// Only generate code for types marked "// +stateify
+			// savable" in one of the proceeding comment lines. If
+			// the line is marked "// +stateify type" then only
+			// generate type information and register the type.
 			if d.Doc == nil {
 				continue
 			}
-			savable := false
+			var (
+				generateTypeInfo    = false
+				generateSaverLoader = false
+			)
 			for _, l := range d.Doc.List {
 				if l.Text == "// +stateify savable" {
-					savable = true
+					generateTypeInfo = true
+					generateSaverLoader = true
 					break
 				}
+				if l.Text == "// +stateify type" {
+					generateTypeInfo = true
+				}
 			}
-			if !savable {
+			if !generateTypeInfo && !generateSaverLoader {
 				continue
 			}
 
 			for _, gs := range d.Specs {
 				ts := gs.(*ast.TypeSpec)
-				switch ts.Type.(type) {
-				case *ast.InterfaceType, *ast.ChanType, *ast.FuncType, *ast.ParenExpr, *ast.StarExpr:
-					// Don't register.
-					break
+				switch x := ts.Type.(type) {
 				case *ast.StructType:
 					maybeEmitImports()
 
-					ss := ts.Type.(*ast.StructType)
+					// Record the slot for each field.
+					fieldCount := 0
+					fields := make(map[string]int)
+					emitField := func(name string) {
+						fmt.Fprintf(outputFile, "		\"%s\",\n", name)
+						fields[name] = fieldCount
+						fieldCount++
+					}
+					emitFieldValue := func(name string, _ string) {
+						emitField(name)
+					}
+					emitLoadValue := func(name, typName string) {
+						fmt.Fprintf(outputFile, "	m.LoadValue(%d, new(%s), func(y interface{}) { x.load%s(y.(%s)) })\n", fields[name], typName, camelCased(name), typName)
+					}
+					emitLoad := func(name string) {
+						fmt.Fprintf(outputFile, "	m.Load(%d, &x.%s)\n", fields[name], name)
+					}
+					emitLoadWait := func(name string) {
+						fmt.Fprintf(outputFile, "	m.LoadWait(%d, &x.%s)\n", fields[name], name)
+					}
+					emitSaveValue := func(name, typName string) {
+						fmt.Fprintf(outputFile, "	var %s %s = x.save%s()\n", name, typName, camelCased(name))
+						fmt.Fprintf(outputFile, "	m.SaveValue(%d, %s)\n", fields[name], name)
+					}
+					emitSave := func(name string) {
+						fmt.Fprintf(outputFile, "	m.Save(%d, &x.%s)\n", fields[name], name)
+					}
+
+					// Generate the type name method.
+					fmt.Fprintf(outputFile, "func (x *%s) StateTypeName() string {\n", ts.Name.Name)
+					fmt.Fprintf(outputFile, "	return \"%s.%s\"\n", *fullPkg, ts.Name.Name)
+					fmt.Fprintf(outputFile, "}\n\n")
+
+					// Generate the fields method.
+					fmt.Fprintf(outputFile, "func (x *%s) StateFields() []string {\n", ts.Name.Name)
+					fmt.Fprintf(outputFile, "	return []string{\n")
+					scanFields(x, "", scanFunctions{
+						normal: emitField,
+						wait:   emitField,
+						value:  emitFieldValue,
+					})
+					fmt.Fprintf(outputFile, "	}\n")
+					fmt.Fprintf(outputFile, "}\n\n")
 
-					// Define beforeSave if a definition was not found. This
-					// prevents the code from compiling if a custom beforeSave
-					// was defined in a file not provided to this binary and
-					// prevents inherited methods from being called multiple times
-					// by overriding them.
-					if _, ok := simpleMethods[method{ts.Name.Name, "beforeSave"}]; !ok {
-						fmt.Fprintf(outputFile, "func (x *%s) beforeSave() {}\n", ts.Name.Name)
+					// Define beforeSave if a definition was not found. This prevents
+					// the code from compiling if a custom beforeSave was defined in a
+					// file not provided to this binary and prevents inherited methods
+					// from being called multiple times by overriding them.
+					if _, ok := simpleMethods[method{ts.Name.Name, "beforeSave"}]; !ok && generateSaverLoader {
+						fmt.Fprintf(outputFile, "func (x *%s) beforeSave() {}\n\n", ts.Name.Name)
 					}
 
 					// Generate the save method.
-					fmt.Fprintf(outputFile, "func (x *%s) save(m %sMap) {\n", ts.Name.Name, statePrefix)
-					fmt.Fprintf(outputFile, "	x.beforeSave()\n")
-					scanFields(ss, scanFunctions{zerovalue: emitZeroCheck})
-					scanFields(ss, scanFunctions{value: emitSaveValue})
-					scanFields(ss, scanFunctions{normal: emitSave, wait: emitSave})
-					fmt.Fprintf(outputFile, "}\n\n")
+					//
+					// N.B. For historical reasons, we perform the value saves first,
+					// and perform the value loads last. There should be no dependency
+					// on this specific behavior, but the ability to specify slots
+					// allows a manual implementation to be order-dependent.
+					if generateSaverLoader {
+						fmt.Fprintf(outputFile, "func (x *%s) StateSave(m %sSink) {\n", ts.Name.Name, statePrefix)
+						fmt.Fprintf(outputFile, "	x.beforeSave()\n")
+						scanFields(x, "", scanFunctions{zerovalue: emitZeroCheck})
+						scanFields(x, "", scanFunctions{value: emitSaveValue})
+						scanFields(x, "", scanFunctions{normal: emitSave, wait: emitSave})
+						fmt.Fprintf(outputFile, "}\n\n")
+					}
 
-					// Define afterLoad if a definition was not found. We do this
-					// for the same reason that we do it for beforeSave.
+					// Define afterLoad if a definition was not found. We do this for
+					// the same reason that we do it for beforeSave.
 					_, hasAfterLoad := simpleMethods[method{ts.Name.Name, "afterLoad"}]
-					if !hasAfterLoad {
-						fmt.Fprintf(outputFile, "func (x *%s) afterLoad() {}\n", ts.Name.Name)
+					if !hasAfterLoad && generateSaverLoader {
+						fmt.Fprintf(outputFile, "func (x *%s) afterLoad() {}\n\n", ts.Name.Name)
 					}
 
 					// Generate the load method.
 					//
-					// Note that the manual loads always follow the
-					// automated loads.
-					fmt.Fprintf(outputFile, "func (x *%s) load(m %sMap) {\n", ts.Name.Name, statePrefix)
-					scanFields(ss, scanFunctions{normal: emitLoad, wait: emitLoadWait})
-					scanFields(ss, scanFunctions{value: emitLoadValue})
-					if hasAfterLoad {
-						// The call to afterLoad is made conditionally, because when
-						// AfterLoad is called, the object encodes a dependency on
-						// referred objects (i.e. fields). This means that afterLoad
-						// will not be called until the other afterLoads are called.
-						fmt.Fprintf(outputFile, "	m.AfterLoad(x.afterLoad)\n")
+					// N.B. See the comment above for the save method.
+					if generateSaverLoader {
+						fmt.Fprintf(outputFile, "func (x *%s) StateLoad(m %sSource) {\n", ts.Name.Name, statePrefix)
+						scanFields(x, "", scanFunctions{normal: emitLoad, wait: emitLoadWait})
+						scanFields(x, "", scanFunctions{value: emitLoadValue})
+						if hasAfterLoad {
+							// The call to afterLoad is made conditionally, because when
+							// AfterLoad is called, the object encodes a dependency on
+							// referred objects (i.e. fields). This means that afterLoad
+							// will not be called until the other afterLoads are called.
+							fmt.Fprintf(outputFile, "	m.AfterLoad(x.afterLoad)\n")
+						}
+						fmt.Fprintf(outputFile, "}\n\n")
 					}
-					fmt.Fprintf(outputFile, "}\n\n")
 
 					// Add to our registration.
 					emitRegister(ts.Name.Name)
+
 				case *ast.Ident, *ast.SelectorExpr, *ast.ArrayType:
 					maybeEmitImports()
 
-					_, val := resolveTypeName(ts.Name.Name, ts.Type)
-
-					// Dispatch directly.
-					fmt.Fprintf(outputFile, "func (x *%s) save(m %sMap) {\n", ts.Name.Name, statePrefix)
-					fmt.Fprintf(outputFile, "	m.SaveValue(\"\", (%s)(*x))\n", val)
+					// Generate the info methods.
+					fmt.Fprintf(outputFile, "func (x *%s) StateTypeName() string {\n", ts.Name.Name)
+					fmt.Fprintf(outputFile, "	return \"%s.%s\"\n", *fullPkg, ts.Name.Name)
 					fmt.Fprintf(outputFile, "}\n\n")
-					fmt.Fprintf(outputFile, "func (x *%s) load(m %sMap) {\n", ts.Name.Name, statePrefix)
-					fmt.Fprintf(outputFile, "	m.LoadValue(\"\", new(%s), func(y interface{}) { *x = (%s)(y.(%s)) })\n", val, ts.Name.Name, val)
+					fmt.Fprintf(outputFile, "func (x *%s) StateFields() []string {\n", ts.Name.Name)
+					fmt.Fprintf(outputFile, "	return nil\n")
 					fmt.Fprintf(outputFile, "}\n\n")
 
 					// See above.