netbird/gvisor
0
0
Fork 0
mirror of https://github.com/netbirdio/gvisor.git synced 2026-05-22 17:12:49 -07:00
Code Issues Packages Projects Releases Wiki Activity

Implementing CopyInN

Browse Source 
PiperOrigin-RevId: 554542787
This commit is contained in:
Shambhavi Srivastava
2023-08-07 11:42:20 -07:00
committed by gVisor bot
parent e89e40fded
commit 8ff6816f07
7 changed files with 93 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/marshal/marshal.go
+8
View File
@@ -112,6 +112,14 @@ type Marshallable interface {
// for UnmarshalBytes.
CopyIn(cc CopyContext, addr hostarch.Addr) (int, error)
// CopyInN is like CopyIn, but explicitly requests a partial
// copy-in. Note that this may yield unexpected results for non-packed
// types and the caller may only want to allow this for packed types. See
// comment on UnmarshalBytes.
//
// The limit must be less than or equal to SizeBytes().
CopyInN(cc CopyContext, addr hostarch.Addr, limit int) (int, error)
// CopyOut serializes a Marshallable type to a task's memory. This may only
// be called from a task goroutine. This is more efficient than calling
// MarshalUnsafe on Marshallable.Packed types, as the type being serialized
pkg/marshal/primitive/primitive.go
+5
View File
@@ -105,6 +105,11 @@ func (b *ByteSlice) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, err
return cc.CopyInBytes(addr, *b)
}
// CopyInN implements marshal.Marshallable.CopyInN.
func (b *ByteSlice) CopyInN(cc marshal.CopyContext, addr hostarch.Addr, limit int) (int, error) {
return cc.CopyInBytes(addr, (*b)[:limit])
}
// CopyOut implements marshal.Marshallable.CopyOut.
func (b *ByteSlice) CopyOut(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {
return cc.CopyOutBytes(addr, *b)
tools/go_marshal/gomarshal/generator_interfaces_array_newtype.go
+10 -3
View File
@@ -124,17 +124,24 @@ func (g *interfaceGenerator) emitMarshallableForArrayNewtype(n *ast.Ident, a *as
})
g.emit("}\n\n")
g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n")
g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName())
g.emit("// CopyInN implements marshal.Marshallable.CopyInN.\n")
g.emit("func (%s *%s) CopyInN(cc marshal.CopyContext, addr hostarch.Addr, limit int) (int, error) {\n", g.r, g.typeName())
g.inIndent(func() {
g.emitCastToByteSlice(g.r, "buf", fmt.Sprintf("%s.SizeBytes()", g.r))
g.emit("length, err := cc.CopyInBytes(addr, buf) // escapes: okay.\n")
g.emit("length, err := cc.CopyInBytes(addr, buf[:limit]) // escapes: okay.\n")
g.emitKeepAlive(g.r)
g.emit("return length, err\n")
})
g.emit("}\n\n")
g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n")
g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName())
g.inIndent(func() {
g.emit("return %s.CopyInN(cc, addr, %s.SizeBytes())\n", g.r, g.r)
})
g.emit("}\n\n")
g.emit("// WriteTo implements io.WriterTo.WriteTo.\n")
g.emit("func (%s *%s) WriteTo(writer io.Writer) (int64, error) {\n", g.r, g.typeName())
g.inIndent(func() {
tools/go_marshal/gomarshal/generator_interfaces_dynamic.go
+12 -2
View File
@@ -65,10 +65,11 @@ func (g *interfaceGenerator) emitMarshallableForDynamicType() {
})
g.emit("}\n\n")
g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n")
g.emit("// CopyInN implements marshal.Marshallable.CopyInN.\n")
g.emit("//go:nosplit\n")
g.recordUsedImport("marshal")
g.recordUsedImport("hostarch")
g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName())
g.emit("func (%s *%s) CopyInN(cc marshal.CopyContext, addr hostarch.Addr, limit int) (int, error) {\n", g.r, g.typeName())
g.inIndent(func() {
g.emit("// Type %s doesn't have a packed layout in memory, fall back to UnmarshalBytes.\n", g.typeName())
g.emit("buf := cc.CopyScratchBuffer(%s.SizeBytes()) // escapes: okay.\n", g.r)
@@ -80,6 +81,15 @@ func (g *interfaceGenerator) emitMarshallableForDynamicType() {
})
g.emit("}\n\n")
g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n")
g.recordUsedImport("marshal")
g.recordUsedImport("hostarch")
g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName())
g.inIndent(func() {
g.emit("return %s.CopyInN(cc, addr, %s.SizeBytes())\n", g.r, g.r)
})
g.emit("}\n\n")
g.emit("// WriteTo implements io.WriterTo.WriteTo.\n")
g.recordUsedImport("io")
g.emit("func (%s *%s) WriteTo(writer io.Writer) (int64, error) {\n", g.r, g.typeName())
tools/go_marshal/gomarshal/generator_interfaces_primitive_newtype.go
+10 -3
View File
@@ -184,17 +184,24 @@ func (g *interfaceGenerator) emitMarshallableForPrimitiveNewtype(nt *ast.Ident)
})
g.emit("}\n\n")
g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n")
g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName())
g.emit("// CopyInN implements marshal.Marshallable.CopyInN.\n")
g.emit("func (%s *%s) CopyInN(cc marshal.CopyContext, addr hostarch.Addr, limit int) (int, error) {\n", g.r, g.typeName())
g.inIndent(func() {
g.emitCastToByteSlice(g.r, "buf", fmt.Sprintf("%s.SizeBytes()", g.r))
g.emit("length, err := cc.CopyInBytes(addr, buf) // escapes: okay.\n")
g.emit("length, err := cc.CopyInBytes(addr, buf[:limit]) // escapes: okay.\n")
g.emitKeepAlive(g.r)
g.emit("return length, err\n")
})
g.emit("}\n\n")
g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n")
g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName())
g.inIndent(func() {
g.emit("return %s.CopyInN(cc, addr, %s.SizeBytes())\n", g.r, g.r)
})
g.emit("}\n\n")
g.emit("// WriteTo implements io.WriterTo.WriteTo.\n")
g.emit("func (%s *%s) WriteTo(writer io.Writer) (int64, error) {\n", g.r, g.typeName())
g.inIndent(func() {
tools/go_marshal/gomarshal/generator_interfaces_struct.go
+14 -4
View File
@@ -322,6 +322,7 @@ func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType) {
fallback()
}
})
g.emit("}\n\n")
g.emit("// CopyOutN implements marshal.Marshallable.CopyOutN.\n")
g.recordUsedImport("marshal")
@@ -364,15 +365,15 @@ func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType) {
})
g.emit("}\n\n")
g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n")
g.emit("// CopyInN implements marshal.Marshallable.CopyInN.\n")
g.recordUsedImport("marshal")
g.recordUsedImport("hostarch")
g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName())
g.emit("func (%s *%s) CopyInN(cc marshal.CopyContext, addr hostarch.Addr, limit int) (int, error) {\n", g.r, g.typeName())
g.inIndent(func() {
fallback := func() {
g.emit("// Type %s doesn't have a packed layout in memory, fall back to UnmarshalBytes.\n", g.typeName())
g.emit("buf := cc.CopyScratchBuffer(%s.SizeBytes()) // escapes: okay.\n", g.r)
g.emit("length, err := cc.CopyInBytes(addr, buf) // escapes: okay.\n")
g.emit("length, err := cc.CopyInBytes(addr, buf[:limit]) // escapes: okay.\n")
g.emit("// Unmarshal unconditionally. If we had a short copy-in, this results in a\n")
g.emit("// partially unmarshalled struct.\n")
g.emit("%s.UnmarshalBytes(buf) // escapes: fallback.\n", g.r)
@@ -390,7 +391,7 @@ func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType) {
// Fast deserialization.
g.emitCastToByteSlice(g.r, "buf", fmt.Sprintf("%s.SizeBytes()", g.r))
g.emit("length, err := cc.CopyInBytes(addr, buf) // escapes: okay.\n")
g.emit("length, err := cc.CopyInBytes(addr, buf[:limit]) // escapes: okay.\n")
g.emitKeepAlive(g.r)
g.emit("return length, err\n")
} else {
@@ -399,6 +400,15 @@ func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType) {
})
g.emit("}\n\n")
g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n")
g.recordUsedImport("marshal")
g.recordUsedImport("hostarch")
g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName())
g.inIndent(func() {
g.emit("return %s.CopyInN(cc, addr, %s.SizeBytes())\n", g.r, g.r)
})
g.emit("}\n\n")
g.emit("// WriteTo implements io.WriterTo.WriteTo.\n")
g.recordUsedImport("io")
g.emit("func (%s *%s) WriteTo(writer io.Writer) (int64, error) {\n", g.r, g.typeName())
tools/go_marshal/test/marshal_test.go
+34
View File
@@ -222,6 +222,32 @@ func limitedCopyOut(t *testing.T, src marshal.Marshallable, limit int) {
compareMemory(t, expectedMem, actualMem, n)
}
// copyInN marshals src to task memory, requesting the marshalling to be
// limited to limit bytes.
func copyInN(t *testing.T, src, dst marshal.Marshallable, limit int) {
var cc mockCopyContext
cc.populate(src)
cc.setLimit(limit)
n, err := dst.CopyInN(&cc, hostarch.Addr(0), limit)
if err != nil {
t.Errorf("CopyInN returned unexpected error: %v", err)
}
if n != limit {
t.Errorf("CopyInN copied unexpected number of bytes, expected %d, got %d", limit, n)
}
expectedMem := unsafeMemory(src)
defer runtime.KeepAlive(src)
actualMem := unsafeMemory(dst)
defer runtime.KeepAlive(dst)
t.Logf("Expected: %v + %v\n", expectedMem[:n], expectedMem[n:])
t.Logf("Actual : %v + %v\n", actualMem[:n], actualMem[n:])
compareMemory(t, expectedMem, actualMem, n)
}
// copyOutN marshals src to task memory, requesting the marshalling to be
// limited to limit bytes.
func copyOutN(t *testing.T, src marshal.Marshallable, limit int) {
@@ -296,6 +322,14 @@ func TestLimitedMarshalling(t *testing.T) {
copyOutN(t, expected, expected.SizeBytes()/2)
})
// Explicitly request partial copy-in.
t.Run(fmt.Sprintf("PartialCopyInN_%v", ty), func(t *testing.T) {
expected := reflect.New(ty).Interface().(marshal.Marshallable)
analysis.RandomizeValue(expected)
copyInN(t, expected, expected, expected.SizeBytes()/2)
})
}
}