refactor(crit): pass proto struct as argument

The proto struct required for encoding or decoding the images was previously being inferred from the `images.ProtoHandler()` function. Now, it is inferred using that only in the CLI, and is passed to the actual CRIT functions as a parameter. This eliminates the redundant protobuf bindings being imported when using CRIT as a library. Signed-off-by: Prajwal S N <prajwalnadig21@gmail.com>
checkpoint-restore · Apr 26, 2023 · a2f943a · a2f943a
1 parent 7e3a0d0
commit a2f943a
Show file tree

Hide file tree

Showing 13 changed files with 251 additions and 168 deletions.
diff --git a/crit/Makefile b/crit/Makefile
@@ -1,7 +1,7 @@
 GO ?= go
 CRIU ?= criu
 
-bin/crit: cmd/cli.go
+bin/crit: cmd/main.go
 	$(GO) build ${GOFLAGS} -o $@ $^
 
 ../test/loop/loop:

diff --git a/crit/cmd/cli.go → crit/cli/cli.go b/crit/cmd/cli.go → crit/cli/cli.go
@@ -1,4 +1,4 @@
-package main
+package cli
 
 import (
 	"encoding/json"
@@ -40,11 +40,29 @@ var decodeCmd = &cobra.Command{
 	Long: `Convert the input binary image to JSON and write it to a file.
 If no output file is provided, the JSON is printed to stdout.`,
 	Run: func(cmd *cobra.Command, args []string) {
-		c = crit.NewCli(inputFilePath, outputFilePath,
+		var (
+			inputFile *os.File
+			err       error
+		)
+		if inputFilePath == "" {
+			inputFile = os.Stdin
+		} else {
+			inputFile, err = os.Open(inputFilePath)
+			if err != nil {
+				log.Fatal(fmt.Errorf("error opening input file: %w", err))
+			}
+			defer inputFile.Close()
+		}
+
+		c = crit.New(inputFile, nil,
 			inputDirPath, pretty, noPayload)
-		img, err := c.Decode()
+		entryType, err := GetEntryTypeFromImg(inputFile)
 		if err != nil {
-			log.Fatal(err)
+			log.Fatal(fmt.Errorf("error getting protobuf binding: %w", err))
+		}
+		img, err := c.Decode(entryType)
+		if err != nil {
+			log.Fatal(fmt.Errorf("error decoding image: %w", err))
 		}
 
 		var jsonData []byte
@@ -54,7 +72,7 @@ If no output file is provided, the JSON is printed to stdout.`,
 			jsonData, err = json.Marshal(img)
 		}
 		if err != nil {
-			log.Fatal(errors.New(fmt.Sprint("Error processing data into JSON: ", err)))
+			log.Fatal(fmt.Errorf("error processing data into JSON: %w", err))
 		}
 		// If no output file, print to stdout
 		if outputFilePath == "" {
@@ -64,12 +82,13 @@ If no output file is provided, the JSON is printed to stdout.`,
 		// Write to output file
 		jsonFile, err := os.Create(outputFilePath)
 		if err != nil {
-			log.Fatal(errors.New(fmt.Sprint("Error opening destination file: ", err)))
+			log.Fatal(fmt.Errorf("error opening destination file: %w", err))
 		}
 		defer jsonFile.Close()
+
 		_, err = jsonFile.Write(jsonData)
 		if err != nil {
-			log.Fatal(errors.New(fmt.Sprint("Error writing JSON data: ", err)))
+			log.Fatal(fmt.Errorf("error writing JSON data: %w", err))
 		}
 	},
 }
@@ -80,16 +99,43 @@ var encodeCmd = &cobra.Command{
 	Short: "Convert JSON to binary image file",
 	Long:  "Convert the input JSON to a CRIU image file.",
 	Run: func(cmd *cobra.Command, args []string) {
-		c = crit.NewCli(inputFilePath, outputFilePath,
+		var (
+			inputFile, outputFile *os.File
+			err                   error
+		)
+		if inputFilePath == "" {
+			inputFile = os.Stdin
+		} else {
+			inputFile, err = os.Open(inputFilePath)
+			if err != nil {
+				log.Fatal(fmt.Errorf("error opening input file: %w", err))
+			}
+			defer inputFile.Close()
+		}
+		if outputFilePath == "" {
+			outputFile = os.Stdout
+		} else {
+			outputFile, err = os.Create(outputFilePath)
+			if err != nil {
+				log.Fatal(fmt.Errorf("error opening output file: %w", err))
+			}
+			defer outputFile.Close()
+		}
+
+		c = crit.New(inputFile, outputFile,
 			inputDirPath, pretty, noPayload)
+		entryType, err := GetEntryTypeFromJSON(inputFile)
+		if err != nil {
+			log.Fatal(fmt.Errorf("error getting protobuf binding: %w", err))
+		}
 		// Convert JSON to Go struct
-		img, err := c.Parse()
+		img, err := c.Parse(entryType)
 		if err != nil {
-			log.Fatal(err)
+			log.Fatal(fmt.Errorf("error parsing JSON: %w", err))
 		}
 		// Write Go struct to binary image file
 		if err := c.Encode(img); err != nil {
-			log.Fatal(err)
+			log.Fatal(fmt.Errorf("error writing to file: %w", err))
 		}
 	},
 }
@@ -103,16 +149,34 @@ var showCmd = &cobra.Command{
 	Run: func(cmd *cobra.Command, args []string) {
 		inputFilePath = args[0]
 		pretty = true
-		c = crit.NewCli(inputFilePath, outputFilePath,
+		var (
+			inputFile *os.File
+			err       error
+		)
+		if inputFilePath == "" {
+			inputFile = os.Stdin
+		} else {
+			inputFile, err = os.Open(inputFilePath)
+			if err != nil {
+				log.Fatal(fmt.Errorf("error opening input file: %w", err))
+			}
+			defer inputFile.Close()
+		}
+
+		c = crit.New(inputFile, nil,
 			inputDirPath, pretty, noPayload)
-		img, err := c.Decode()
+		entryType, err := GetEntryTypeFromImg(inputFile)
 		if err != nil {
-			log.Fatal(err)
+			log.Fatal(fmt.Errorf("error getting protobuf binding: %w", err))
+		}
+		img, err := c.Decode(entryType)
+		if err != nil {
+			log.Fatal(fmt.Errorf("error decoding image: %w", err))
 		}
 
 		jsonData, err := json.MarshalIndent(img, "", "    ")
 		if err != nil {
-			log.Fatal(errors.New(fmt.Sprint("Error processing data into JSON: ", err)))
+			log.Fatal(fmt.Errorf("error processing data into JSON: %w", err))
 		}
 		fmt.Println(string(jsonData))
 	},
@@ -125,16 +189,30 @@ var infoCmd = &cobra.Command{
 	Args:  cobra.ExactArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
 		inputFilePath = args[0]
-		c = crit.NewCli(inputFilePath, outputFilePath,
+		var (
+			inputFile *os.File
+			err       error
+		)
+		if inputFilePath == "" {
+			inputFile = os.Stdin
+		} else {
+			inputFile, err = os.Open(inputFilePath)
+			if err != nil {
+				log.Fatal(fmt.Errorf("error opening input file: %w", err))
+			}
+			defer inputFile.Close()
+		}
+
+		c = crit.New(inputFile, nil,
 			inputDirPath, pretty, noPayload)
 		img, err := c.Info()
 		if err != nil {
-			log.Fatal(err)
+			log.Fatal(fmt.Errorf("error decoding image: %w", err))
 		}
 
 		jsonData, err := json.MarshalIndent(img, "", "    ")
 		if err != nil {
-			log.Fatal(errors.New(fmt.Sprint("Error processing data into JSON: ", err)))
+			log.Fatal(fmt.Errorf("error processing data into JSON: %w", err))
 		}
 		fmt.Println(string(jsonData))
 	},
@@ -155,10 +233,10 @@ var xCmd = &cobra.Command{
 		// returned object since we don't really care
 		// about the data itself, as long as we can
 		// marshal it into JSON and display it.
-		var xData interface{}
+		var xData any
 		var err error
 
-		c = crit.NewCli(inputFilePath, outputFilePath,
+		c = crit.New(nil, nil,
 			inputDirPath, pretty, noPayload)
 		// Switch the explore type and call the handler.
 		switch args[1] {
@@ -171,22 +249,22 @@ var xCmd = &cobra.Command{
 		case "rss":
 			xData, err = c.ExploreRss()
 		default:
-			err = errors.New("error exploring directory: Invalid explore type")
+			err = errors.New("error exploring directory: invalid explore type")
 		}
 		if err != nil {
-			log.Fatal(err)
+			log.Fatal(fmt.Errorf("error exploring directory: %w", err))
 		}
 
 		jsonData, err := json.MarshalIndent(xData, "", "    ")
 		if err != nil {
-			log.Fatal(errors.New(fmt.Sprint("Error processing data into JSON: ", err)))
+			log.Fatal(fmt.Errorf("error processing data into JSON: %w", err))
 		}
 		fmt.Println(string(jsonData))
 	},
 }
 
 // Add all commands to the root command and configure flags
-func init() {
+func Init() {
 	// Disable completion generation
 	rootCmd.CompletionOptions.DisableDefaultCmd = true
 	// Decode options
@@ -212,8 +290,8 @@ func init() {
 	rootCmd.AddCommand(xCmd)
 }
 
-func main() {
+func Run() {
 	if err := rootCmd.Execute(); err != nil {
-		log.Fatal(err)
+		log.Fatal(fmt.Errorf("error running CLI: %w", err))
 	}
 }
diff --git a/crit/images/handler.go → crit/cli/handler.go b/crit/images/handler.go → crit/cli/handler.go
@@ -1,8 +1,12 @@
-package images
+package cli
 
 import (
+	"encoding/json"
 	"fmt"
+	"io"
+	"os"
 
+	"github.com/checkpoint-restore/go-criu/v6/crit"
 	"github.com/checkpoint-restore/go-criu/v6/crit/images/apparmor"
 	"github.com/checkpoint-restore/go-criu/v6/crit/images/autofs"
 	binfmt_misc "github.com/checkpoint-restore/go-criu/v6/crit/images/binfmt-misc"
@@ -59,7 +63,41 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
-func ProtoHandler(magic string) (proto.Message, error) {
+func GetEntryTypeFromImg(imgFile *os.File) (proto.Message, error) {
+	magic, err := crit.ReadMagic(imgFile)
+	if err != nil {
+		return nil, err
+	}
+	// Seek to the beginning of the file
+	_, err = imgFile.Seek(0, io.SeekStart)
+	if err != nil {
+		return nil, err
+	}
+
+	return protoHandler(magic)
+}
+
+func GetEntryTypeFromJSON(jsonFile *os.File) (proto.Message, error) {
+	jsonData, err := io.ReadAll(jsonFile)
+	if err != nil {
+		return nil, err
+	}
+	// Seek to the beginning of the file
+	_, err = jsonFile.Seek(0, io.SeekStart)
+	if err != nil {
+		return nil, err
+	}
+
+	var img map[string]any
+	err = json.Unmarshal(jsonData, &img)
+	if err != nil {
+		return nil, err
+	}
+
+	return protoHandler(img["magic"].(string))
+}
+
+func protoHandler(magic string) (proto.Message, error) {
 	switch magic {
 	case "APPARMOR":
 		return &apparmor.ApparmorEntry{}, nil
@@ -189,6 +227,14 @@ func ProtoHandler(magic string) (proto.Message, error) {
 		return &utsns.UtsnsEntry{}, nil
 	case "VMAS":
 		return &vma.VmaEntry{}, nil
+	/* Pagemap and ghost file have custom handlers
+	and cannot use a single proto struct to be
+	encoded or decoded. Hence, for these two
+	image types, nil is returned. */
+	case "PAGEMAP":
+		return nil, nil
+	case "GHOST_FILE":
+		return nil, nil
 	}
-	return nil, fmt.Errorf("no handler found for magic 0x%x", magic)
+	return nil, fmt.Errorf("no protobuf binding found for magic 0x%x", magic)
 }
diff --git a/crit/cmd/main.go b/crit/cmd/main.go
@@ -0,0 +1,8 @@
+package main
+
+import "github.com/checkpoint-restore/go-criu/v6/crit/cli"
+
+func main() {
+	cli.Init()
+	cli.Run()
+}