crit: add memory page content retrieval features

This commit add MemoryReader struct with methods to retrieve the memory pages
associated with a process.

Signed-off-by: Kouame Behouba Manasse <[email protected]>

crit/mempages.go

@@ -0,0 +1,134 @@
+package crit
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/checkpoint-restore/go-criu/v6/crit/images/pagemap"
+)
+
+const pageSize = 4096
+
+// 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            int
+	pagesID        uint32
+	pagemapEntries []*pagemap.PagemapEntry
+}
+
+// NewMemoryReader creates a new MemoryReader instance with all fields populated.
+func NewMemoryReader(checkpointDir string, pid int) (*MemoryReader, error) {
+	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,
+		pagesID:        pagesID,
+		pagemapEntries: pagemapEntries,
+	}, nil
+}
+
+// GetMemPages retrieves content of memory associated with a pid.
+func (mr *MemoryReader) GetMemPages(start, end uint64) (*bytes.Buffer, error) {
+	// Calculate and check the size of the memory area
+	size := end - start
+	if size == 0 {
+		return nil, fmt.Errorf("memory size must be greater than 0")
+	}
+
+	startPage := start / pageSize
+	endPage := end / 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(pageSize))
+		}
+
+		var nSkip, nRead uint64
+
+		if pageNumber == startPage {
+			nSkip = start - pageNumber*pageSize
+			if startPage == endPage {
+				nRead = size
+			} else {
+				nRead = pageSize - nSkip
+			}
+		} else if pageNumber == endPage {
+			nSkip = 0
+			nRead = end - pageNumber*pageSize
+		} else {
+			nSkip = 0
+			nRead = pageSize
+		}
+
+		if _, err := buffer.Write(page[nSkip : nSkip+nRead]); err != nil {
+			return nil, err
+		}
+	}
+
+	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)*pageSize == pageNo*pageSize {
+				found = true
+				break
+			}
+			offset += 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, pageSize)
+
+		if _, err := f.ReadAt(buff, int64(offset)); err != nil {
+			return nil, err
+		}
+
+		return buff, nil
+	}
+	return nil, nil
+}

crit/mempages_test.go

@@ -0,0 +1,123 @@
+package crit
+
+import (
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"regexp"
+	"strconv"
+	"strings"
+	"testing"
+)
+
+const (
+	testImgsDir = "test-imgs"
+)
+
+// TestNewMemoryReader tests the NewMemoryReader function.
+func TestNewMemoryReader(t *testing.T) {
+	pid, err := getTestImgsPID(testImgsDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	testCases := []struct {
+		name          string
+		dir           string
+		pid           int
+		expectedError error
+	}{
+		{
+			name:          "Valid test-imgs directory and pid",
+			dir:           testImgsDir,
+			pid:           pid,
+			expectedError: nil,
+		},
+		{
+			name:          "Invalid test-imgs directory",
+			dir:           "no test directory",
+			pid:           pid,
+			expectedError: errors.New("no such file or directory"),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			mr, err := NewMemoryReader(tc.dir, tc.pid)
+			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)
+			}
+		})
+	}
+}
+
+func TestGetMemPages(t *testing.T) {
+	pid, err := getTestImgsPID(testImgsDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	type testcase struct {
+		name          string
+		start         uint64
+		end           uint64
+		expectedError error
+	}
+
+	testCases := []testcase{
+		{
+			name:          "Zero memory area size",
+			start:         0,
+			end:           0,
+			expectedError: errors.New("memory size must be greater than 0"),
+		},
+	}
+
+	mr, err := NewMemoryReader(testImgsDir, pid)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for i, pm := range mr.pagemapEntries[:2] {
+		testCases = append(testCases, testcase{
+			name:          fmt.Sprintf("Valid pagemap address %d", i+1),
+			start:         pm.GetVaddr(),
+			end:           pm.GetVaddr() + uint64(uint32(pageSize)*pm.GetNrPages()),
+			expectedError: nil,
+		})
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			buff, err := mr.GetMemPages(tc.start, tc.end)
+			if err != nil && !strings.Contains(err.Error(), tc.expectedError.Error()) {
+				t.Errorf("Expected error: %v, got error: %v", tc.expectedError, err)
+			}
+
+			if tc.expectedError == nil && buff.Len() < 1 {
+				t.Errorf("Returned memory chunk is expected to be non-empty")
+			}
+		})
+	}
+}
+
+func getTestImgsPID(dir string) (int, error) {
+	files, err := ioutil.ReadDir(dir)
+	if err != nil {
+		return 0, fmt.Errorf("Failed to read directory: %w\n", err)
+	}
+
+	re := regexp.MustCompile(`pagemap-(\d+)\.img`)
+
+	for _, file := range files {
+		if re.MatchString(file.Name()) {
+			numberStr := re.FindStringSubmatch(file.Name())[1]
+			return strconv.Atoi(numberStr)
+		}
+	}
+	return 0, fmt.Errorf("no pid found in %s", dir)
+}