Merge pull request #138 from snprajwal/explore-sk

feat(crit): add `crit x sk` for sockets

crit/cli/cli.go

@@ -220,9 +220,9 @@ var infoCmd = &cobra.Command{
 
 // The `crit x` command
 var xCmd = &cobra.Command{
-	Use:   "x DIR {ps|fds|mems|rss}",
+	Use:   "x DIR {ps|fd|mem|rss|sk}",
 	Short: "Explore the image directory",
-	Long:  "Explore the image directory with one of (ps, fds, mems, rss) options",
+	Long:  "Explore the image directory with one of (ps, fd, mem, rss, sk) options",
 	// Exactly two arguments are required:
 	// * Path of the input directory
 	// * Explore type
@@ -242,14 +242,16 @@ var xCmd = &cobra.Command{
 		switch args[1] {
 		case "ps":
 			xData, err = c.ExplorePs()
-		case "fds":
+		case "fd", "fds":
 			xData, err = c.ExploreFds()
-		case "mems":
+		case "mem", "mems":
 			xData, err = c.ExploreMems()
 		case "rss":
 			xData, err = c.ExploreRss()
+		case "sk":
+			xData, err = c.ExploreSk()
 		default:
-			err = errors.New("error exploring directory: invalid explore type")
+			err = errors.New("invalid explore type (supported: {ps|fd|mem|rss|sk})")
 		}
 		if err != nil {
 			log.Fatal(fmt.Errorf("error exploring directory: %w", err))

crit/crit.go

@@ -25,6 +25,7 @@ type Critter interface {
 	ExploreFds() ([]*Fd, error)
 	ExploreMems() ([]*MemMap, error)
 	ExploreRss() ([]*RssMap, error)
+	ExploreSk() ([]*Sk, error)
 }
 
 // crit implements the Critter interface. It contains:

crit/explore.go

@@ -71,7 +71,7 @@ func (c *crit) ExplorePs() (*PsTree, error) {
 // Fd represents the file descriptors opened in a single process
 type Fd struct {
 	PId   uint32  `json:"pId"`
-	Files []*File `json:"files,omitempty"`
+	Files []*File `json:"files"`
 }
 
 // File represents a single opened file
@@ -145,6 +145,11 @@ func (c *crit) ExploreFds() ([]*Fd, error) {
 			Path: filePath,
 		})
 
+		// Omit if the process has no file descriptors
+		if len(fdEntry.Files) == 0 {
+			continue
+		}
+
 		fds = append(fds, &fdEntry)
 	}
 
@@ -368,3 +373,114 @@ func (c *crit) ExploreRss() ([]*RssMap, error) {
 
 	return rssMaps, nil
 }
+
+// Sk represents the sockets associated with a single process
+type Sk struct {
+	PId     uint32    `json:"pId"`
+	Sockets []*Socket `json:"sockets"`
+}
+
+// Socket represents a single socket
+type Socket struct {
+	Fd       uint32 `json:"fd"`
+	FdType   string `json:"fdType"`
+	Family   string `json:"family,omitempty"`
+	Protocol string `json:"protocol,omitempty"`
+	Type     string `json:"type,omitempty"`
+	State    string `json:"state,omitempty"`
+	SrcAddr  string `json:"srcAddr,omitempty"`
+	SrcPort  uint32 `json:"srcPort,omitempty"`
+	DestAddr string `json:"destAddr,omitempty"`
+	DestPort uint32 `json:"destPort,omitempty"`
+	SendBuf  string `json:"sendBuf,omitempty"`
+	RecvBuf  string `json:"recvBuf,omitempty"`
+}
+
+// ExploreSk searches the process tree for sockets
+// and returns a list of PIDs with the associated sockets
+func (c *crit) ExploreSk() ([]*Sk, error) {
+	psTreeImg, err := getImg(filepath.Join(c.inputDirPath, "pstree.img"), &pstree.PstreeEntry{})
+	if err != nil {
+		return nil, err
+	}
+
+	sks := make([]*Sk, 0)
+	for _, entry := range psTreeImg.Entries {
+		process := entry.Message.(*pstree.PstreeEntry)
+		pID := process.GetPid()
+		// Get file with object IDs
+		idsImg, err := getImg(filepath.Join(c.inputDirPath, fmt.Sprintf("ids-%d.img", pID)), &criu_core.TaskKobjIdsEntry{})
+		if err != nil {
+			return nil, err
+		}
+		filesID := idsImg.Entries[0].Message.(*criu_core.TaskKobjIdsEntry).GetFilesId()
+		// Get open file descriptors
+		fdInfoImg, err := getImg(filepath.Join(c.inputDirPath, fmt.Sprintf("fdinfo-%d.img", filesID)), &fdinfo.FdinfoEntry{})
+		if err != nil {
+			return nil, err
+		}
+		skEntry := Sk{PId: pID}
+		for _, fdInfoEntry := range fdInfoImg.Entries {
+			fdInfo := fdInfoEntry.Message.(*fdinfo.FdinfoEntry)
+			file, err := getFile(c.inputDirPath, fdInfo.GetId())
+			if err != nil {
+				return nil, err
+			}
+			socket := Socket{
+				Fd:     fdInfo.GetFd(),
+				FdType: fdInfo.GetType().String(),
+			}
+			switch fdInfo.GetType() {
+			case fdinfo.FdTypes_INETSK:
+				if isk := file.GetIsk(); isk != nil {
+					socket.State = getSkState(tcpState(isk.GetState()))
+					socket.Family = getAddressFamily(isk.GetFamily())
+					socket.Protocol = getSkProtocol(isk.GetProto())
+					socket.Type = getSkType(isk.GetType())
+					socket.SrcAddr = processIP(isk.GetSrcAddr())
+					socket.SrcPort = isk.GetSrcPort()
+					socket.DestAddr = processIP(isk.GetDstAddr())
+					socket.DestPort = isk.GetDstPort()
+					socket.SendBuf = countBytes(int64(isk.GetOpts().GetSoSndbuf()))
+					socket.RecvBuf = countBytes(int64(isk.GetOpts().GetSoRcvbuf()))
+				}
+			case fdinfo.FdTypes_UNIXSK:
+				if usk := file.GetUsk(); usk != nil {
+					socket.State = getSkState(tcpState(usk.GetState()))
+					socket.Type = getSkType(usk.GetType())
+					socket.SrcAddr = string(usk.GetName())
+					socket.SendBuf = countBytes(int64(usk.GetOpts().GetSoSndbuf()))
+					socket.RecvBuf = countBytes(int64(usk.GetOpts().GetSoRcvbuf()))
+				}
+			case fdinfo.FdTypes_PACKETSK:
+				if psk := file.GetPsk(); psk != nil {
+					socket.Type = getSkType(psk.GetProtocol())
+					socket.Protocol = getSkProtocol(psk.GetProtocol())
+					socket.SendBuf = countBytes(int64(psk.GetOpts().GetSoSndbuf()))
+					socket.RecvBuf = countBytes(int64(psk.GetOpts().GetSoRcvbuf()))
+				}
+			case fdinfo.FdTypes_NETLINKSK:
+				if nlsk := file.GetNlsk(); nlsk != nil {
+					socket.State = getSkState(tcpState(nlsk.GetState()))
+					socket.Protocol = getSkProtocol(nlsk.GetProtocol())
+					socket.Type = getSkType(nlsk.GetProtocol())
+					socket.SendBuf = countBytes(int64(nlsk.GetOpts().GetSoSndbuf()))
+					socket.RecvBuf = countBytes(int64(nlsk.GetOpts().GetSoRcvbuf()))
+				}
+			default:
+				continue
+			}
+
+			skEntry.Sockets = append(skEntry.Sockets, &socket)
+		}
+
+		// Omit if the process has no associated sockets
+		if len(skEntry.Sockets) == 0 {
+			continue
+		}
+
+		sks = append(sks, &skEntry)
+	}
+
+	return sks, nil
+}

crit/utils.go

@@ -4,9 +4,11 @@ import (
 	"encoding/binary"
 	"fmt"
 	"io"
+	"net"
 	"os"
 	"path/filepath"
 	"strconv"
+	"syscall"
 
 	"github.com/checkpoint-restore/go-criu/v6/crit/images/fdinfo"
 	"github.com/checkpoint-restore/go-criu/v6/crit/images/pipe"
@@ -44,13 +46,13 @@ func ReadMagic(f *os.File) (string, error) {
 
 // Helper to convert bytes into a more readable unit
 func countBytes(n int64) string {
-	const UNIT int64 = 1024
-	if n < UNIT {
+	const unit int64 = 1024
+	if n < unit {
 		return fmt.Sprint(n, " B")
 	}
-	div, exp := UNIT, 0
-	for i := n / UNIT; i >= UNIT; i /= UNIT {
-		div *= UNIT
+	div, exp := unit, 0
+	for i := n / unit; i >= unit; i /= unit {
+		div *= unit
 		exp++
 	}
 	return fmt.Sprintf("%.1f %cB", float64(n)/float64(div), "KMGTPE"[exp])
@@ -110,27 +112,37 @@ var (
 	filesImg, regImg, pipeImg, unixSkImg *CriuImage
 )
 
-// Helper to get file path for exploring file descriptors
-func getFilePath(dir string, fID uint32, fType fdinfo.FdTypes) (string, error) {
-	var filePath string
+// Helper to fetch a file if it exists in files.img
+func getFile(dir string, fID uint32) (*fdinfo.FileEntry, error) {
 	var err error
-	// Get open files
 	if filesImg == nil {
 		filesImg, err = getImg(filepath.Join(dir, "files.img"), &fdinfo.FileEntry{})
 		if err != nil {
-			return "", err
+			return nil, err
 		}
 	}
 
-	// Check if file entry is present
 	var file *fdinfo.FileEntry
 	for _, entry := range filesImg.Entries {
 		file = entry.Message.(*fdinfo.FileEntry)
 		if file.GetId() == fID {
-			break
+			return file, nil
 		}
 	}
 
+	return nil, nil
+}
+
+// Helper to get file path for exploring file descriptors
+func getFilePath(dir string, fID uint32, fType fdinfo.FdTypes) (string, error) {
+	var filePath string
+	var err error
+	// Fetch the file, if it exists in file.img
+	file, err := getFile(dir, fID)
+	if err != nil {
+		return "", err
+	}
+
 	switch fType {
 	case fdinfo.FdTypes_REG:
 		filePath, err = getRegFilePath(dir, file, fID)
@@ -152,7 +164,7 @@ func getRegFilePath(dir string, file *fdinfo.FileEntry, fID uint32) (string, err
 		if file.GetReg() != nil {
 			return file.GetReg().GetName(), nil
 		}
-		return "Unknown path", nil
+		return "unknown", nil
 	}
 
 	if regImg == nil {
@@ -168,7 +180,7 @@ func getRegFilePath(dir string, file *fdinfo.FileEntry, fID uint32) (string, err
 		}
 	}
 
-	return "Unknown path", nil
+	return "unknown", nil
 }
 
 // Helper to get file path of pipe files
@@ -246,3 +258,129 @@ func (ps *PsTree) FindPs(pid uint32) *PsTree {
 	}
 	return nil
 }
+
+// Helper to convert slice of uint32 into IP address string
+func processIP(parts []uint32) string {
+	// IPv4
+	if len(parts) == 1 {
+		ip := make(net.IP, net.IPv4len)
+		binary.LittleEndian.PutUint32(ip, parts[0])
+		return ip.String()
+	}
+	// IPv6
+	if len(parts) == 4 {
+		ip := make(net.IP, net.IPv6len)
+		for _, part := range parts {
+			binary.LittleEndian.PutUint32(ip, part)
+		}
+		return ip.String()
+	}
+	// Invalid
+	return ""
+}
+
+// tcpState represents the state of a TCP connection.
+type tcpState uint8
+
+// https://github.com/torvalds/linux/blob/999f6631/include/net/tcp_states.h#L12
+const (
+	tcpEstablished tcpState = iota + 1
+	tcpSynSent
+	tcpSynReceived
+	tcpFinWait1
+	tcpFinWait2
+	tcpTimeWait
+	tcpClose
+	tcpCloseWait
+	tcpLastAck
+	tcpListen
+	tcpClosing
+	tcpNewSynRecv
+)
+
+var states = map[tcpState]string{
+	tcpEstablished: "ESTABLISHED",
+	tcpSynSent:     "SYN_SENT",
+	tcpSynReceived: "SYN_RECEIVED",
+	tcpFinWait1:    "FIN_WAIT_1",
+	tcpFinWait2:    "FIN_WAIT_2",
+	tcpTimeWait:    "TIME_WAIT",
+	tcpClose:       "CLOSE",
+	tcpCloseWait:   "CLOSE_WAIT",
+	tcpLastAck:     "LAST_ACK",
+	tcpListen:      "LISTEN",
+	tcpClosing:     "CLOSING",
+	tcpNewSynRecv:  "NEW_SYN_RECV",
+}
+
+// Helper to identify socket state
+func getSkState(state tcpState) string {
+	if stateName, ok := states[state]; ok {
+		return stateName
+	}
+	return ""
+}
+
+// Helper to identify address family
+func getAddressFamily(family uint32) string {
+	switch family {
+	case syscall.AF_UNIX:
+		return "UNIX"
+	case syscall.AF_NETLINK:
+		return "NETLINK"
+	case syscall.AF_BRIDGE:
+		return "BRIDGE"
+	case syscall.AF_KEY:
+		return "KEY"
+	case syscall.AF_PACKET:
+		return "PACKET"
+	case syscall.AF_INET:
+		return "IPV4"
+	case syscall.AF_INET6:
+		return "IPV6"
+	default:
+		return ""
+	}
+}
+
+// Helper to identify socket type
+func getSkType(skType uint32) string {
+	switch skType {
+	case syscall.SOCK_STREAM:
+		return "STREAM"
+	case syscall.SOCK_DGRAM:
+		return "DGRAM"
+	case syscall.SOCK_SEQPACKET:
+		return "SEQPACKET"
+	case syscall.SOCK_RAW:
+		return "RAW"
+	case syscall.SOCK_RDM:
+		return "RDM"
+	case syscall.SOCK_PACKET:
+		return "PACKET"
+	default:
+		return ""
+	}
+}
+
+// Helper to identify socket protocol
+func getSkProtocol(protocol uint32) string {
+	switch protocol {
+	case syscall.IPPROTO_ICMP:
+		return "ICMP"
+	case syscall.IPPROTO_ICMPV6:
+		return "ICMPV6"
+	case syscall.IPPROTO_IGMP:
+		return "IGMP"
+	case syscall.IPPROTO_RAW:
+		return "RAW"
+	case syscall.IPPROTO_TCP:
+		return "TCP"
+	case syscall.IPPROTO_UDP:
+		return "UDP"
+	case syscall.IPPROTO_UDPLITE:
+		return "UDPLITE"
+	default:
+		return ""
+	}
+}