From f674cbd60f4f229cd97d2e2cf7211ef1f7c1488d Mon Sep 17 00:00:00 2001
From: Pavel Emelyanov <xemul@virtuozzo.com>
Date: Thu, 24 Mar 2016 16:01:58 +0300
Subject: [PATCH] unix: Collect packets via cinfo engine

Signed-off-by: Pavel Emelyanov <xemul@virtuozzo.com>
---
 criu/cr-restore.c       |  5 +++--
 criu/include/sk-queue.h |  2 +-
 criu/include/sockets.h  |  1 -
 criu/sk-queue.c         | 50 +++++++++++++++--------------------------
 criu/sk-unix.c          |  5 -----
 5 files changed, 22 insertions(+), 41 deletions(-)

diff --git a/criu/cr-restore.c b/criu/cr-restore.c
index 3a407c7be..9db73ff7a 100644
--- a/criu/cr-restore.c
+++ b/criu/cr-restore.c
@@ -76,6 +76,8 @@
 #include "seccomp.h"
 #include "bitmap.h"
 #include "fault-injection.h"
+#include "sk-queue.h"
+
 #include "parasite-syscall.h"
 
 #include "protobuf.h"
@@ -178,6 +180,7 @@ static struct collect_image_info *cinfos[] = {
 	&ext_file_cinfo,
 	&timerfd_cinfo,
 	&file_locks_cinfo,
+	&sk_queues_cinfo,
 };
 
 struct post_prepare_cb {
@@ -242,8 +245,6 @@ static int root_prepare_shared(void)
 		return -1;
 	if (collect_fifo())
 		return -1;
-	if (collect_unix_sockets())
-		return -1;
 
 	if (tty_verify_active_pairs())
 		return -1;
diff --git a/criu/include/sk-queue.h b/criu/include/sk-queue.h
index 9044de0b0..e0a47af1b 100644
--- a/criu/include/sk-queue.h
+++ b/criu/include/sk-queue.h
@@ -1,7 +1,7 @@
 #ifndef __CR_SK_QUEUE_H__
 #define __CR_SK_QUEUE_H__
 
-extern int read_sk_queues(void);
+extern struct collect_image_info sk_queues_cinfo;
 extern int dump_sk_queue(int sock_fd, int sock_id);
 extern int restore_sk_queue(int fd, unsigned int peer_id);
 
diff --git a/criu/include/sockets.h b/criu/include/sockets.h
index b566fae2b..b3bacb750 100644
--- a/criu/include/sockets.h
+++ b/criu/include/sockets.h
@@ -38,7 +38,6 @@ struct ns_id;
 extern int collect_sockets(struct ns_id *);
 extern int collect_inet_sockets(void);
 extern struct collect_image_info unix_sk_cinfo;
-extern int collect_unix_sockets(void);
 extern int fix_external_unix_sockets(void);
 
 extern struct collect_image_info netlink_sk_cinfo;
diff --git a/criu/sk-queue.c b/criu/sk-queue.c
index 208d70fd9..c278e0a12 100644
--- a/criu/sk-queue.c
+++ b/criu/sk-queue.c
@@ -33,43 +33,29 @@ struct sk_packet {
 
 static LIST_HEAD(packets_list);
 
-int read_sk_queues(void)
+static int collect_one_packet(void *obj, ProtobufCMessage *msg, struct cr_img *img)
 {
-	struct sk_packet *pkt;
-	int ret;
-	struct cr_img *img;
+	struct sk_packet *pkt = obj;
 
-	pr_info("Trying to read socket queues image\n");
+	pkt->entry = pb_msg(msg, SkPacketEntry);
+	pkt->img_off = lseek(img_raw_fd(img), 0, SEEK_CUR);
+	/*
+	 * NOTE: packet must be added to the tail. Otherwise sequence
+	 * will be broken.
+	 */
+	list_add_tail(&pkt->list, &packets_list);
+	lseek(img_raw_fd(img), pkt->entry->length, SEEK_CUR);
 
-	img = open_image(CR_FD_SK_QUEUES, O_RSTR);
-	if (!img)
-		return -1;
-
-	while (1) {
-		ret = -1;
-		pkt = xmalloc(sizeof(*pkt));
-		if (!pkt) {
-			pr_err("Failed to allocate packet header\n");
-			break;
-		}
-		ret = pb_read_one_eof(img, &pkt->entry, PB_SK_QUEUES);
-		if (ret <= 0)
-			break;
-
-		pkt->img_off = lseek(img_raw_fd(img), 0, SEEK_CUR);
-		/*
-		 * NOTE: packet must be added to the tail. Otherwise sequence
-		 * will be broken.
-		 */
-		list_add_tail(&pkt->list, &packets_list);
-		lseek(img_raw_fd(img), pkt->entry->length, SEEK_CUR);
-	}
-	close_image(img);
-	xfree(pkt);
-
-	return ret;
+	return 0;
 }
 
+struct collect_image_info sk_queues_cinfo = {
+	.fd_type = CR_FD_SK_QUEUES,
+	.pb_type = PB_SK_QUEUES,
+	.priv_size = sizeof(struct sk_packet),
+	.collect = collect_one_packet,
+};
+
 int dump_sk_queue(int sock_fd, int sock_id)
 {
 	SkPacketEntry pe = SK_PACKET_ENTRY__INIT;
diff --git a/criu/sk-unix.c b/criu/sk-unix.c
index 24510cab1..1396d76ff 100644
--- a/criu/sk-unix.c
+++ b/criu/sk-unix.c
@@ -1330,11 +1330,6 @@ struct collect_image_info unix_sk_cinfo = {
 	.flags = COLLECT_SHARED,
 };
 
-int collect_unix_sockets(void)
-{
-	return read_sk_queues();
-}
-
 static int resolve_unix_peers(void *unused)
 {
 	struct unix_sk_info *ui, *peer;