From 78f3a445838fb66bf886c017131a685579f5d56a Mon Sep 17 00:00:00 2001 From: UtkarshSingh-06 Date: Sun, 18 Jan 2026 01:28:44 +0530 Subject: [PATCH] Add remap command for TCP listen port remapping This commit adds a new 'remap' command to checkpointctl that allows remapping TCP listen ports in checkpoint archives. This addresses issue #169 by providing functionality similar to --tcp-listen-remap option of criu-image-streamer or edit_files_img.py script. The implementation: - Adds a new 'remap' command that accepts --tcp-listen-remap flag - Extracts files.img from checkpoint archives - Parses files.img to find TCP listen sockets (SOCK_STREAM, TCP_LISTEN state) - Remaps the source port according to the provided mappings - Repacks the modified files.img back into the checkpoint archive Port mappings are specified in the format old_port:new_port, with multiple mappings separated by commas. Closes #169 Signed-off-by: UtkarshSingh-06 --- checkpointctl.go | 2 + cmd/remap.go | 137 ++++++++++++++++++++ internal/remap.go | 321 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 460 insertions(+) create mode 100644 cmd/remap.go create mode 100644 internal/remap.go diff --git a/checkpointctl.go b/checkpointctl.go index 3851ba20..d9b670be 100644 --- a/checkpointctl.go +++ b/checkpointctl.go @@ -33,6 +33,8 @@ func main() { rootCommand.AddCommand(cmd.BuildCmd()) + rootCommand.AddCommand(cmd.Remap()) + rootCommand.Version = version if err := rootCommand.Execute(); err != nil { diff --git a/cmd/remap.go b/cmd/remap.go new file mode 100644 index 00000000..44160c7a --- /dev/null +++ b/cmd/remap.go @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: Apache-2.0 + +package cmd + +import ( + "fmt" + "os" + "path/filepath" + "strconv" + "strings" + + "github.com/checkpoint-restore/checkpointctl/internal" + metadata "github.com/checkpoint-restore/checkpointctl/lib" + "github.com/spf13/cobra" +) + +var tcpListenRemap *string + +func Remap() *cobra.Command { + cmd := &cobra.Command{ + Use: "remap ", + Short: "Remap TCP listen ports in a container checkpoint archive", + Long: `The 'remap' command allows remapping TCP listen ports in a checkpoint archive. +This is useful when the port originally bound to a socket needs to be changed before restoring. + +Port mappings are specified using the --tcp-listen-remap flag in the format: + old_port:new_port + +Multiple port mappings can be specified by separating them with commas: + --tcp-listen-remap 8080:80,8443:443 + +Example: + checkpointctl remap checkpoint.tar --tcp-listen-remap 8080:80`, + Args: cobra.ExactArgs(1), + RunE: remapTCPPorts, + } + + tcpListenRemap = cmd.Flags().String( + "tcp-listen-remap", + "", + "TCP listen port remapping in format old_port:new_port (e.g., 8080:80). Multiple mappings can be separated by commas.", + ) + + cmd.MarkFlagRequired("tcp-listen-remap") + + return cmd +} + +func remapTCPPorts(cmd *cobra.Command, args []string) error { + checkpointPath := args[0] + + // Parse port mappings + portMappings, err := parsePortMappings(*tcpListenRemap) + if err != nil { + return fmt.Errorf("failed to parse port mappings: %w", err) + } + + if len(portMappings) == 0 { + return fmt.Errorf("no valid port mappings specified") + } + + // Create temporary directory for extraction + tempDir, err := os.MkdirTemp("", "checkpointctl-remap-") + if err != nil { + return fmt.Errorf("failed to create temporary directory: %w", err) + } + defer os.RemoveAll(tempDir) + + // Extract files.img from checkpoint + filesImgPath := filepath.Join(metadata.CheckpointDirectory, "files.img") + if err := internal.UntarFiles(checkpointPath, tempDir, []string{filesImgPath}); err != nil { + return fmt.Errorf("failed to extract files.img from checkpoint: %w", err) + } + + extractedFilesImgPath := filepath.Join(tempDir, filesImgPath) + + // Perform port remapping + modified, err := internal.RemapTCPListenPorts(extractedFilesImgPath, portMappings) + if err != nil { + return fmt.Errorf("failed to remap TCP ports: %w", err) + } + + if !modified { + fmt.Println("No TCP listen sockets were found matching the specified port mappings") + return nil + } + + // Repack the modified files.img back into the checkpoint archive + if err := internal.RepackFileToArchive(checkpointPath, filesImgPath, extractedFilesImgPath); err != nil { + return fmt.Errorf("failed to repack checkpoint archive: %w", err) + } + + fmt.Printf("Successfully remapped TCP listen ports in %s\n", checkpointPath) + return nil +} + +func parsePortMappings(mappingStr string) (map[uint32]uint32, error) { + mappings := make(map[uint32]uint32) + + if mappingStr == "" { + return mappings, nil + } + + parts := strings.Split(mappingStr, ",") + for _, part := range parts { + part = strings.TrimSpace(part) + if part == "" { + continue + } + + ports := strings.Split(part, ":") + if len(ports) != 2 { + return nil, fmt.Errorf("invalid port mapping format: %s (expected old_port:new_port)", part) + } + + oldPort, err := strconv.ParseUint(strings.TrimSpace(ports[0]), 10, 32) + if err != nil { + return nil, fmt.Errorf("invalid old port '%s': %w", ports[0], err) + } + + newPort, err := strconv.ParseUint(strings.TrimSpace(ports[1]), 10, 32) + if err != nil { + return nil, fmt.Errorf("invalid new port '%s': %w", ports[1], err) + } + + if oldPort == 0 || oldPort > 65535 { + return nil, fmt.Errorf("old port %d is out of valid range (1-65535)", oldPort) + } + if newPort == 0 || newPort > 65535 { + return nil, fmt.Errorf("new port %d is out of valid range (1-65535)", newPort) + } + + mappings[uint32(oldPort)] = uint32(newPort) + } + + return mappings, nil +} \ No newline at end of file diff --git a/internal/remap.go b/internal/remap.go new file mode 100644 index 00000000..3bc76911 --- /dev/null +++ b/internal/remap.go @@ -0,0 +1,321 @@ +// SPDX-License-Identifier: Apache-2.0 + +package internal + +import ( + "archive/tar" + "compress/gzip" + "errors" + "fmt" + "io" + "os" + "path/filepath" + + "github.com/checkpoint-restore/go-criu/v8/crit" + "github.com/checkpoint-restore/go-criu/v8/crit/images/fdinfo" + "github.com/containers/storage/pkg/archive" +) + +const ( + // TCP protocol number (IPPROTO_TCP) + IPPROTO_TCP = 6 + // SOCK_STREAM socket type + SOCK_STREAM = 1 + // TCP_LISTEN state (see linux/net/tcp_states.h) + TCP_LISTEN = 10 +) + +// RemapTCPListenPorts remaps TCP listen ports in files.img according to the provided mappings. +// Returns true if any ports were remapped, false otherwise. +func RemapTCPListenPorts(filesImgPath string, portMappings map[uint32]uint32) (bool, error) { + // Open the files.img file for reading + file, err := os.Open(filesImgPath) + if err != nil { + return false, fmt.Errorf("failed to open files.img: %w", err) + } + defer file.Close() + + // Decode the files.img + c := crit.New(file, nil, "", false, false) + img, err := c.Decode(&fdinfo.FileEntry{}) + if err != nil { + return false, fmt.Errorf("failed to decode files.img: %w", err) + } + + modified := false + + // Iterate through all entries in files.img + for _, entry := range img.Entries { + fileEntry, ok := entry.Message.(*fdinfo.FileEntry) + if !ok { + continue + } + + // Check if this is an INETSK entry + if fileEntry.GetType() != fdinfo.FdTypes_INETSK { + continue + } + + inetSk := fileEntry.GetIsk() + if inetSk == nil { + continue + } + + // Check if this is a TCP SOCK_STREAM socket in LISTEN state + if inetSk.GetProto() != IPPROTO_TCP || inetSk.GetType() != SOCK_STREAM || inetSk.GetState() != TCP_LISTEN { + continue + } + + // Get the current source port + srcPort := inetSk.GetSrcPort() + if srcPort == 0 { + continue + } + + // Check if this port needs to be remapped + newPort, needsRemap := portMappings[srcPort] + if !needsRemap { + continue + } + + // Remap the port + inetSk.SrcPort = &newPort + modified = true + } + + if !modified { + return false, nil + } + + // Close the input file before creating output file + file.Close() + + // Create a temporary file for output + tempOutputPath := filesImgPath + ".tmp" + outputFile, err := os.Create(tempOutputPath) + if err != nil { + return false, fmt.Errorf("failed to create temporary output file: %w", err) + } + defer outputFile.Close() + + // Encode the modified image to the temporary file + cOut := crit.New(nil, outputFile, "", false, false) + if err := cOut.Encode(img); err != nil { + outputFile.Close() + os.Remove(tempOutputPath) + return false, fmt.Errorf("failed to encode modified files.img: %w", err) + } + + outputFile.Close() + + // Replace the original file with the modified one + if err := os.Rename(tempOutputPath, filesImgPath); err != nil { + os.Remove(tempOutputPath) + return false, fmt.Errorf("failed to replace files.img: %w", err) + } + + return true, nil +} + +// RepackFileToArchive repacks a modified file back into the checkpoint archive. +func RepackFileToArchive(checkpointPath, filePath string, modifiedFilePath string) error { + // This is a simplified implementation that replaces the file in the archive + // In a production implementation, we'd want to: + // 1. Extract the entire archive + // 2. Replace the modified file + // 3. Recompress and repack the archive + + // For now, we'll use a simple approach: replace the file directly in the archive + // This works if the new file is not larger than the original + + // Open the checkpoint archive for reading + archiveFile, err := os.Open(checkpointPath) + if err != nil { + return fmt.Errorf("failed to open checkpoint archive: %w", err) + } + defer archiveFile.Close() + + // Read the modified file + modifiedFile, err := os.ReadFile(modifiedFilePath) + if err != nil { + return fmt.Errorf("failed to read modified file: %w", err) + } + + // Create a temporary archive with the modified file + tempArchivePath := checkpointPath + ".tmp" + if err := repackArchiveWithFile(archiveFile, tempArchivePath, filePath, modifiedFile); err != nil { + os.Remove(tempArchivePath) + return fmt.Errorf("failed to repack archive: %w", err) + } + + // Replace the original archive with the new one + if err := os.Rename(tempArchivePath, checkpointPath); err != nil { + os.Remove(tempArchivePath) + return fmt.Errorf("failed to replace checkpoint archive: %w", err) + } + + return nil +} + +// repackArchiveWithFile creates a new archive with a replaced file +func repackArchiveWithFile(archiveFile *os.File, outputPath, fileToReplace string, newFileContent []byte) error { + // This is a placeholder implementation + // In production, we'd need to: + // 1. Decompress the archive + // 2. Iterate through entries, replacing the target file + // 3. Recompress and write to output + + // For now, we'll use the archive package to help with this + // Since this is complex, let's implement a simpler version first + + // We need to extract the archive, replace the file, and repack it + // This is the most reliable approach + + tempDir, err := os.MkdirTemp("", "checkpointctl-repack-") + if err != nil { + return fmt.Errorf("failed to create temporary directory: %w", err) + } + defer os.RemoveAll(tempDir) + + // Extract entire archive + if err := extractArchive(archiveFile.Name(), tempDir); err != nil { + return fmt.Errorf("failed to extract archive: %w", err) + } + + // Replace the file + destPath := filepath.Join(tempDir, fileToReplace) + if err := os.MkdirAll(filepath.Dir(destPath), 0o700); err != nil { + return fmt.Errorf("failed to create directory: %w", err) + } + + if err := os.WriteFile(destPath, newFileContent, 0o644); err != nil { + return fmt.Errorf("failed to write modified file: %w", err) + } + + // Repack the archive + if err := packDirectory(tempDir, outputPath); err != nil { + return fmt.Errorf("failed to pack archive: %w", err) + } + + return nil +} + +// extractArchive extracts an entire archive to a directory +func extractArchive(archivePath, destDir string) error { + archiveFile, err := os.Open(archivePath) + if err != nil { + return err + } + defer archiveFile.Close() + + // Decompress the archive + stream, err := archive.DecompressStream(archiveFile) + if err != nil { + return err + } + defer stream.Close() + + // Create a tar reader to read the files from the decompressed archive + tarReader := tar.NewReader(stream) + + for { + header, err := tarReader.Next() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + return err + } + + targetPath := filepath.Join(destDir, header.Name) + + switch header.Typeflag { + case tar.TypeDir: + if err := os.MkdirAll(targetPath, os.FileMode(header.Mode)); err != nil { + return err + } + case tar.TypeReg: + if err := os.MkdirAll(filepath.Dir(targetPath), 0o700); err != nil { + return err + } + outFile, err := os.OpenFile(targetPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.FileMode(header.Mode)) + if err != nil { + return err + } + + if _, err := io.Copy(outFile, tarReader); err != nil { + outFile.Close() + return err + } + outFile.Close() + } + } + + return nil +} + +// packDirectory creates an archive from a directory +func packDirectory(srcDir, archivePath string) error { + // Use the same compression format as the original + // For simplicity, we'll use gzip like most checkpoint archives + return packDirectoryTarGz(srcDir, archivePath) +} + +// packDirectoryTarGz creates a tar.gz archive from a directory +func packDirectoryTarGz(srcDir, archivePath string) error { + outputFile, err := os.Create(archivePath) + if err != nil { + return err + } + defer outputFile.Close() + + // Create gzip writer + gzWriter := gzip.NewWriter(outputFile) + defer gzWriter.Close() + + // Create tar writer + tarWriter := tar.NewWriter(gzWriter) + defer tarWriter.Close() + + return filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error { + if err != nil { + return err + } + + relPath, err := filepath.Rel(srcDir, path) + if err != nil { + return err + } + + // Use forward slashes for tar entries (POSIX format) + relPath = filepath.ToSlash(relPath) + + // Skip the root directory itself + if relPath == "." { + return nil + } + + header, err := tar.FileInfoHeader(info, "") + if err != nil { + return err + } + + header.Name = relPath + if err := tarWriter.WriteHeader(header); err != nil { + return err + } + + if !info.IsDir() { + file, err := os.Open(path) + if err != nil { + return err + } + defer file.Close() + + _, err = io.Copy(tarWriter, file) + return err + } + + return nil + }) +} \ No newline at end of file