Skip to content

Commit

Permalink
crit: add memory page content retrieval features
Browse files Browse the repository at this point in the history
This commit adds the MemoryReader struct with methods to retrieve
the memory pages associated with a process.

Signed-off-by: Kouame Behouba Manasse <behouba@gmail.com>
  • Loading branch information
behouba committed Jun 11, 2023
1 parent 748b90b commit 57984f9
Show file tree
Hide file tree
Showing 2 changed files with 344 additions and 0 deletions.
144 changes: 144 additions & 0 deletions crit/mempages.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,144 @@
package crit

import (
"bytes"
"errors"
"fmt"
"os"
"path/filepath"

"github.com/checkpoint-restore/go-criu/v6/crit/images/pagemap"
)

var sysPageSize = os.Getpagesize()

// MemoryReader is a struct used to retrieve
// the content of memory associated with a specific process ID (pid).
// New instances should be created with NewMemoryReader()
type MemoryReader struct {
checkpointDir string
pid uint32
pagesID uint32
pageSize int
pagemapEntries []*pagemap.PagemapEntry
}

// NewMemoryReader creates a new instance of MemoryReader with all the fields populated
func NewMemoryReader(checkpointDir string, pid uint32, pageSize int) (*MemoryReader, error) {
if pageSize == 0 {
pageSize = sysPageSize
}

// Check if the given page size is a positive power of 2, otherwise return an error
if (pageSize & (pageSize - 1)) != 0 {
return nil, errors.New("page size should be a positive power of 2")
}

pagemapImg, err := getImg(filepath.Join(checkpointDir, fmt.Sprintf("pagemap-%d.img", pid)), &pagemap.PagemapHead{})
if err != nil {
return nil, err
}

pagesID := pagemapImg.Entries[0].Message.(*pagemap.PagemapHead).GetPagesId()

pagemapEntries := make([]*pagemap.PagemapEntry, 0)

for _, entry := range pagemapImg.Entries[1:] {
pagemapEntries = append(pagemapEntries, entry.Message.(*pagemap.PagemapEntry))
}

return &MemoryReader{
checkpointDir: checkpointDir,
pid: pid,
pageSize: pageSize,
pagesID: pagesID,
pagemapEntries: pagemapEntries,
}, nil
}

// GetMemPages retrieves the content of memory pages
// associated with a given process ID (pid).
// It retrieves the memory content within the
// specified range defined by the start and end addresses.
func (mr *MemoryReader) GetMemPages(start, end uint64) (*bytes.Buffer, error) {
size := end - start

startPage := start / uint64(mr.pageSize)
endPage := end / uint64(mr.pageSize)

var buffer bytes.Buffer

for pageNumber := startPage; pageNumber <= endPage; pageNumber++ {
var page []byte = nil

pageMem, err := mr.getPage(pageNumber)
if err != nil {
return nil, err
}

if pageMem != nil {
page = pageMem
} else {
page = bytes.Repeat([]byte("\x00"), int(mr.pageSize))
}

var nSkip, nRead uint64

if pageNumber == startPage {
nSkip = start - pageNumber*uint64(mr.pageSize)
if startPage == endPage {
nRead = size
} else {
nRead = uint64(mr.pageSize) - nSkip
}
} else if pageNumber == endPage {
nSkip = 0
nRead = end - pageNumber*uint64(mr.pageSize)
} else {
nSkip = 0
nRead = uint64(mr.pageSize)
}

if _, err := buffer.Write(page[nSkip : nSkip+nRead]); err != nil {
return nil, err
}

Check warning on line 104 in crit/mempages.go

View check run for this annotation

Codecov / codecov/patch

crit/mempages.go#L103-L104

Added lines #L103 - L104 were not covered by tests
}

return &buffer, nil
}

// getPage retrieves a memory page from the pages.img file.
func (mr *MemoryReader) getPage(pageNo uint64) ([]byte, error) {
var offset uint64 = 0

// Iterate over pagemap entries to find the corresponding page
for _, m := range mr.pagemapEntries {
found := false
for i := 0; i < int(*m.NrPages); i++ {
if m.GetVaddr()+uint64(i)*uint64(mr.pageSize) == pageNo*uint64(mr.pageSize) {
found = true
break
}
offset += uint64(mr.pageSize)
}

if !found {
continue
}
f, err := os.Open(filepath.Join(mr.checkpointDir, fmt.Sprintf("pages-%d.img", mr.pagesID)))
if err != nil {
return nil, err
}

defer f.Close()

buff := make([]byte, mr.pageSize)

if _, err := f.ReadAt(buff, int64(offset)); err != nil {
return nil, err
}

return buff, nil
}
return nil, nil
}
200 changes: 200 additions & 0 deletions crit/mempages_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,200 @@
package crit

import (
"errors"
"os"
"path/filepath"
"strings"
"testing"

"github.com/checkpoint-restore/go-criu/v6/crit/images/pstree"
)

const (
testImgsDir = "test-imgs"
)

func TestNewMemoryReader(t *testing.T) {
pid, err := getTestImgPID(testImgsDir)
if err != nil {
t.Fatal(err)
}

testCases := []struct {
name string
dir string
pid uint32
pageSize int
expectedError error
}{
{
name: "Page size is 0",
dir: testImgsDir,
pid: pid,
pageSize: 0,
expectedError: nil,
},
{
name: "Invalid page size",
dir: testImgsDir,
pid: pid,
pageSize: 4000,
expectedError: errors.New("page size should be a positive power of 2"),
},
{
name: "Invalid test-imgs directory",
dir: "no test directory",
pid: pid,
expectedError: errors.New("no such file or directory"),
},
{
name: "Valid test-imgs directory, pid and page size",
dir: testImgsDir,
pid: pid,
pageSize: 4096,
expectedError: nil,
},
}

for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
mr, err := NewMemoryReader(tc.dir, tc.pid, tc.pageSize)
if err != nil && !strings.Contains(err.Error(), tc.expectedError.Error()) {
t.Errorf("Expected error: %v, got error: %v", tc.expectedError, err)
}

if mr == nil && tc.expectedError == nil {
t.Errorf("MemoryReader creation failed for checkpoint directory: %s and pid: %d", tc.dir, tc.pid)
}
})
}
}

// TestNewMemoryReader test GetMempages method of MemoryReader.
func TestGetMemPages(t *testing.T) {
type testcase struct {
name string
mr *MemoryReader
start uint64
end uint64
expectedError error
}

pid, err := getTestImgPID(testImgsDir)
if err != nil {
t.Fatal(err)
}

// Create a temporary empty memory pages file for testing
tmpFilePath := filepath.Join(os.TempDir(), "pages-0.img")
tmpFile, err := os.Create(tmpFilePath)
if err != nil {
t.Fatal(err)
}

defer func() {
tmpFile.Close()
if err := os.Remove(tmpFile.Name()); err != nil {
t.Fatal(err)
}
}()

mr, err := NewMemoryReader(testImgsDir, pid, sysPageSize)
if err != nil {
t.Fatal(err)
}

testCases := []testcase{
{
name: "Zero memory area size",
mr: &MemoryReader{
checkpointDir: testImgsDir,
pid: pid,
pageSize: sysPageSize,
pagesID: mr.pagesID,
pagemapEntries: mr.pagemapEntries,
},
start: 0,
end: 0,
expectedError: nil,
},
{
name: "Valid pagemap entry 1",
mr: &MemoryReader{
checkpointDir: testImgsDir,
pid: pid,
pageSize: sysPageSize,
pagesID: mr.pagesID,
pagemapEntries: mr.pagemapEntries,
},
start: mr.pagemapEntries[0].GetVaddr(),
end: mr.pagemapEntries[0].GetVaddr() + uint64(uint32(sysPageSize)*mr.pagemapEntries[0].GetNrPages()),
expectedError: nil,
},
{
name: "Valid pagemap entry 2",
mr: &MemoryReader{
checkpointDir: testImgsDir,
pid: pid,
pageSize: sysPageSize,
pagesID: mr.pagesID,
pagemapEntries: mr.pagemapEntries,
},
start: mr.pagemapEntries[1].GetVaddr(),
end: mr.pagemapEntries[1].GetVaddr() + uint64(uint32(sysPageSize)*mr.pagemapEntries[1].GetNrPages()),
expectedError: nil,
},
{
name: "Invalid pages file",
mr: &MemoryReader{
checkpointDir: testImgsDir,
pid: pid,
pageSize: sysPageSize,
pagesID: mr.pagesID + 1,
pagemapEntries: mr.pagemapEntries,
},
start: mr.pagemapEntries[0].GetVaddr(),
end: mr.pagemapEntries[0].GetVaddr() + uint64(uint32(sysPageSize)*mr.pagemapEntries[0].GetNrPages()),
expectedError: errors.New("no such file or directory"),
},
{
name: "Empty pages file",
mr: &MemoryReader{
checkpointDir: os.TempDir(),
pid: pid,
pageSize: sysPageSize,
pagesID: 0,
pagemapEntries: mr.pagemapEntries,
},
start: mr.pagemapEntries[1].GetVaddr(),
end: mr.pagemapEntries[1].GetVaddr() + uint64(uint32(sysPageSize)*mr.pagemapEntries[1].GetNrPages()),
expectedError: errors.New("EOF"),
},
}

for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
buff, err := tc.mr.GetMemPages(tc.start, tc.end)
if err != nil && tc.expectedError != nil {
if !strings.Contains(err.Error(), tc.expectedError.Error()) {
t.Errorf("Expected error: %v, got error: %v", tc.expectedError, err)
}
}

if tc.expectedError == nil && buff == nil {
t.Errorf("Returned memory chunk is expected to be non-empty")
}
})
}
}

// helper function to get the PID from the test-imgs directory
func getTestImgPID(dir string) (uint32, error) {
psTreeImg, err := getImg(filepath.Join(dir, "pstree.img"), &pstree.PstreeEntry{})
if err != nil {
return 0, err
}
process := psTreeImg.Entries[0].Message.(*pstree.PstreeEntry)

return process.GetPgid(), nil
}

0 comments on commit 57984f9

Please sign in to comment.