diff --git a/pkg/safemem/io.go b/pkg/safemem/io.go index 9551ca853..ecbfdf39b 100644 --- a/pkg/safemem/io.go +++ b/pkg/safemem/io.go @@ -46,12 +46,13 @@ type Writer interface { WriteFromBlocks(srcs BlockSeq) (uint64, error) } -// ReadFullToBlocks repeatedly invokes r.ReadToBlocks until dsts.NumBytes() -// bytes have been read or ReadToBlocks returns an error. -func ReadFullToBlocks(r Reader, dsts BlockSeq) (uint64, error) { +// ReadFullToBlocks repeatedly invokes r until dsts.NumBytes() bytes have been +// read or r returns an error. Note that we avoid a Reader interface receiver +// to avoid heap allocation. +func ReadFullToBlocks(r ReaderFunc, dsts BlockSeq) (uint64, error) { var done uint64 for !dsts.IsEmpty() { - n, err := r.ReadToBlocks(dsts) + n, err := r(dsts) done += n if err != nil { return done, err @@ -61,12 +62,13 @@ func ReadFullToBlocks(r Reader, dsts BlockSeq) (uint64, error) { return done, nil } -// WriteFullFromBlocks repeatedly invokes w.WriteFromBlocks until -// srcs.NumBytes() bytes have been written or WriteFromBlocks returns an error. -func WriteFullFromBlocks(w Writer, srcs BlockSeq) (uint64, error) { +// WriteFullFromBlocks repeatedly invokes w until srcs.NumBytes() bytes have +// been written or w returns an error. Note that we avoid a Writer interface +// receiver to avoid heap allocation. +func WriteFullFromBlocks(w WriterFunc, srcs BlockSeq) (uint64, error) { var done uint64 for !srcs.IsEmpty() { - n, err := w.WriteFromBlocks(srcs) + n, err := w(srcs) done += n if err != nil { return done, err @@ -144,18 +146,6 @@ func (r ToIOReader) Read(dst []byte) (int, error) { return int(n), err } -// ToIOWriter implements io.Writer for a (safemem.)Writer. -type ToIOWriter struct { - Writer Writer -} - -// Write implements io.Writer.Write. -func (w ToIOWriter) Write(src []byte) (int, error) { - // io.Writer does not permit partial writes. - n, err := WriteFullFromBlocks(w.Writer, BlockSeqOf(BlockFromSafeSlice(src))) - return int(n), err -} - // FromIOReader implements Reader for an io.Reader by repeatedly invoking // io.Reader.Read until it returns an error or partial read. This is not // thread-safe. diff --git a/pkg/safemem/io_test.go b/pkg/safemem/io_test.go index 629741bee..db4cd6352 100644 --- a/pkg/safemem/io_test.go +++ b/pkg/safemem/io_test.go @@ -100,7 +100,7 @@ func TestSingleByteReader(t *testing.T) { func TestReadFullToBlocks(t *testing.T) { r := FromIOReader{singleByteReader{bytes.NewBufferString("foobar")}} dsts := makeBlocks(make([]byte, 3), make([]byte, 3)) - n, err := ReadFullToBlocks(r, BlockSeqFromSlice(dsts)) + n, err := ReadFullToBlocks(r.ReadToBlocks, BlockSeqFromSlice(dsts)) // ReadFullToBlocks should call into FromIOReader => singleByteReader // repeatedly until dsts is exhausted. if wantN := uint64(6); n != wantN || err != nil { @@ -187,7 +187,7 @@ func TestWriteFullToBlocks(t *testing.T) { srcs := makeBlocks([]byte("foo"), []byte("bar")) var dst bytes.Buffer w := FromIOWriter{singleByteWriter{&dst}} - n, err := WriteFullFromBlocks(w, BlockSeqFromSlice(srcs)) + n, err := WriteFullFromBlocks(w.WriteFromBlocks, BlockSeqFromSlice(srcs)) // WriteFullToBlocks should call into FromIOWriter => singleByteWriter // repeatedly until srcs is exhausted. if wantN := uint64(6); n != wantN || err != nil { diff --git a/pkg/sentry/fsutil/file_range_set.go b/pkg/sentry/fsutil/file_range_set.go index f6a8f1e62..41477b4a0 100644 --- a/pkg/sentry/fsutil/file_range_set.go +++ b/pkg/sentry/fsutil/file_range_set.go @@ -129,7 +129,7 @@ func (s *FileRangeSet) Fill(ctx context.Context, required, optional memmap.Mappa MemCgID: memCgID, } if readAt != nil { - opts.Reader = safemem.ReaderFunc(func(dsts safemem.BlockSeq) (uint64, error) { + opts.ReaderFunc = func(dsts safemem.BlockSeq) (uint64, error) { var done uint64 for !dsts.IsEmpty() { n, err := func() (uint64, error) { @@ -167,7 +167,7 @@ func (s *FileRangeSet) Fill(ctx context.Context, required, optional memmap.Mappa } } return done, nil - }) + } } fr, err := mf.Allocate(gr.Length(), opts) diff --git a/pkg/sentry/kernel/kcov.go b/pkg/sentry/kernel/kcov.go index 27c3515e0..19b81ad7d 100644 --- a/pkg/sentry/kernel/kcov.go +++ b/pkg/sentry/kernel/kcov.go @@ -99,7 +99,7 @@ func (kcov *Kcov) TaskWork(t *Task) { } // Read in the PC count. - if _, err := safemem.ReadFullToBlocks(rw, kcov.countBlock()); err != nil { + if _, err := safemem.ReadFullToBlocks(rw.ReadToBlocks, kcov.countBlock()); err != nil { panic(fmt.Sprintf("Internal error reading count from kcov area: %v", err)) } @@ -111,7 +111,7 @@ func (kcov *Kcov) TaskWork(t *Task) { // output. kcov.count += uint64(n / 8) rw.off = 0 - if _, err := safemem.WriteFullFromBlocks(rw, kcov.countBlock()); err != nil { + if _, err := safemem.WriteFullFromBlocks(rw.WriteFromBlocks, kcov.countBlock()); err != nil { panic(fmt.Sprintf("Internal error writing count to kcov area: %v", err)) } @@ -337,6 +337,6 @@ type kcovIOWriter struct { // Write implements io.Writer.Write. func (w *kcovIOWriter) Write(p []byte) (int, error) { bs := safemem.BlockSeqOf(safemem.BlockFromSafeSlice(p)) - n, err := safemem.WriteFullFromBlocks(w.rw, bs) + n, err := safemem.WriteFullFromBlocks(w.rw.WriteFromBlocks, bs) return int(n), err } diff --git a/pkg/sentry/mm/pma.go b/pkg/sentry/mm/pma.go index 81d06a01e..8a1e02a3c 100644 --- a/pkg/sentry/mm/pma.go +++ b/pkg/sentry/mm/pma.go @@ -371,11 +371,12 @@ func (mm *MemoryManager) getPMAsInternalLocked(ctx context.Context, vseg vmaIter return pstart, pseg.PrevGap(), err } // Copy contents. + reader := safemem.BlockSeqReader{Blocks: mm.internalMappingsLocked(pseg, copyAR)} fr, err := mm.mf.Allocate(uint64(copyAR.Length()), pgalloc.AllocOpts{ - Kind: usage.Anonymous, - Mode: pgalloc.AllocateAndWritePopulate, - MemCgID: memCgID, - Reader: &safemem.BlockSeqReader{mm.internalMappingsLocked(pseg, copyAR)}, + Kind: usage.Anonymous, + Mode: pgalloc.AllocateAndWritePopulate, + MemCgID: memCgID, + ReaderFunc: reader.ReadToBlocks, }) if _, ok := err.(safecopy.BusError); ok { // If we got SIGBUS during the copy, deliver SIGBUS to diff --git a/pkg/sentry/pgalloc/pgalloc.go b/pkg/sentry/pgalloc/pgalloc.go index 186b87e92..ee72c3f9f 100644 --- a/pkg/sentry/pgalloc/pgalloc.go +++ b/pkg/sentry/pgalloc/pgalloc.go @@ -465,12 +465,12 @@ type AllocOpts struct { // that will fill the allocated memory by invoking host system calls should // pass AllocateOnly. Mode AllocationMode - // If Reader is provided, the allocated memory is filled by calling - // ReadToBlocks() repeatedly until either length bytes are read or a non-nil - // error is returned. It returns the allocated memory, truncated down to the - // nearest page. If this is shorter than length bytes due to an error - // returned by ReadToBlocks(), it returns the partially filled fr and error. - Reader safemem.Reader + // If ReaderFunc is provided, the allocated memory is filled by calling it + // repeatedly until either length bytes are read or a non-nil error is + // returned. It returns the allocated memory, truncated down to the nearest + // page. If this is shorter than length bytes due to an error returned by + // ReaderFunc, it returns the partially filled fr and error. + ReaderFunc safemem.ReaderFunc } // Allocate returns a range of initially-zeroed pages of the given length with @@ -514,7 +514,7 @@ func (f *MemoryFile) Allocate(length uint64, opts AllocOpts) (memmap.FileRange, default: panic(fmt.Sprintf("unknown allocation mode: %d", opts.Mode)) } - if opts.Reader != nil { + if opts.ReaderFunc != nil { if dsts.IsEmpty() { dsts, err = f.MapInternal(fr, hostarch.Write) if err != nil { @@ -522,7 +522,7 @@ func (f *MemoryFile) Allocate(length uint64, opts AllocOpts) (memmap.FileRange, return memmap.FileRange{}, err } } - n, err := safemem.ReadFullToBlocks(opts.Reader, dsts) + n, err := safemem.ReadFullToBlocks(opts.ReaderFunc, dsts) un := uint64(hostarch.Addr(n).RoundDown()) if un < length { // Free unused memory and update fr to contain only the memory that is