From 8f43eff6da4dc5e7da7da082e08157dd55667b21 Mon Sep 17 00:00:00 2001
From: songweibin <song.weibin@zte.com.cn>
Date: Tue, 27 Feb 2018 17:33:54 +0800
Subject: [PATCH] rbd: import with option --export-format fails to protect
 snapshot

Fixes: http://tracker.ceph.com/issues/23038

Signed-off-by: songweibin <song.weibin@zte.com.cn>
---
 doc/dev/rbd-export.rst         | 14 +++++++++++++-
 src/tools/rbd/Utils.h          |  2 ++
 src/tools/rbd/action/Export.cc | 17 +++++++++++++++--
 src/tools/rbd/action/Import.cc | 23 ++++++++++++++++++++++-
 4 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/doc/dev/rbd-export.rst b/doc/dev/rbd-export.rst
index d009b79953ab4..2edb637f6ef15 100644
--- a/doc/dev/rbd-export.rst
+++ b/doc/dev/rbd-export.rst
@@ -83,9 +83,21 @@ End
 
 
 Diffs records
-~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~
+
 Record the all snapshots and the HEAD in this section. 
 
+Snap Protection status
+----------------------
+
+Record the snapshot's protection status if `--export-format=2`.
+- u8: 'p'
+- le64: length of appending data (8)
+- u8: snap protection status (0 for false, 1 for true)
+
+Others
+------
+
 - le64: number of diffs
 - Diffs ...
 
diff --git a/src/tools/rbd/Utils.h b/src/tools/rbd/Utils.h
index c265249fb3799..28dc377005fea 100644
--- a/src/tools/rbd/Utils.h
+++ b/src/tools/rbd/Utils.h
@@ -46,6 +46,8 @@ static const std::string RBD_DIFF_BANNER_V2 ("rbd diff v2\n");
 #define RBD_DIFF_ZERO		'z'
 #define RBD_DIFF_END		'e'
 
+#define RBD_SNAP_PROTECTION_STATUS     'p'
+
 #define RBD_EXPORT_IMAGE_ORDER		'O'
 #define RBD_EXPORT_IMAGE_FEATURES	'T'
 #define RBD_EXPORT_IMAGE_STRIPE_UNIT	'U'
diff --git a/src/tools/rbd/action/Export.cc b/src/tools/rbd/action/Export.cc
index e8941f8460c34..8c109e17dea80 100644
--- a/src/tools/rbd/action/Export.cc
+++ b/src/tools/rbd/action/Export.cc
@@ -156,12 +156,25 @@ int do_export_diff_fd(librbd::Image& image, const char *fromsnapname,
       encode(tag, bl);
       std::string to(endsnapname);
       if (export_format == 2) {
-	len = to.length() + 4;
-	encode(len, bl);
+        len = to.length() + 4;
+        encode(len, bl);
       }
       encode(to, bl);
     }
 
+    if (endsnapname && export_format == 2) {
+      tag = RBD_SNAP_PROTECTION_STATUS;
+      encode(tag, bl);
+      bool is_protected = false;
+      r = image.snap_is_protected(endsnapname, &is_protected);
+      if (r < 0) {
+        return r;
+      }
+      len = 8;
+      encode(len, bl);
+      encode(is_protected, bl);
+    }
+
     tag = RBD_DIFF_IMAGE_SIZE;
     encode(tag, bl);
     uint64_t endsize = info.size;
diff --git a/src/tools/rbd/action/Import.cc b/src/tools/rbd/action/Import.cc
index ef843f94af808..f727c17457b48 100644
--- a/src/tools/rbd/action/Import.cc
+++ b/src/tools/rbd/action/Import.cc
@@ -184,6 +184,21 @@ static int do_image_snap_to(ImportDiffContext *idiffctx, std::string *tosnap)
   return 0;
 }
 
+static int get_snap_protection_status(ImportDiffContext *idiffctx, bool *is_protected)
+{
+  int r;
+  char buf[sizeof(__u8)];
+  r = safe_read_exact(idiffctx->fd, buf, sizeof(buf));
+  if (r < 0) {
+    return r;
+  }
+
+  *is_protected = (buf[0] != 0);
+  idiffctx->update_progress();
+
+  return 0;
+}
+
 static int do_image_resize(ImportDiffContext *idiffctx)
 {
   int r;
@@ -371,6 +386,7 @@ int do_import_diff_fd(librados::Rados &rados, librbd::Image &image, int fd,
 
   // begin image import
   std::string tosnap;
+  bool is_protected = false;
   ImportDiffContext idiffctx(&image, fd, size, no_progress);
   while (r == 0) {
     __u8 tag;
@@ -385,6 +401,8 @@ int do_import_diff_fd(librados::Rados &rados, librbd::Image &image, int fd,
       r = do_image_snap_from(&idiffctx);
     } else if (tag == RBD_DIFF_TO_SNAP) {
       r = do_image_snap_to(&idiffctx, &tosnap);
+    } else if (tag == RBD_SNAP_PROTECTION_STATUS) {
+      r = get_snap_protection_status(&idiffctx, &is_protected);
     } else if (tag == RBD_DIFF_IMAGE_SIZE) {
       r = do_image_resize(&idiffctx);
     } else if (tag == RBD_DIFF_WRITE || tag == RBD_DIFF_ZERO) {
@@ -399,7 +417,10 @@ int do_import_diff_fd(librados::Rados &rados, librbd::Image &image, int fd,
   int temp_r = idiffctx.throttle.wait_for_ret();
   r = (r < 0) ? r : temp_r; // preserve original error
   if (r == 0 && tosnap.length()) {
-    idiffctx.image->snap_create(tosnap.c_str());
+    r = idiffctx.image->snap_create(tosnap.c_str());
+    if (r == 0 && is_protected) {
+      r = idiffctx.image->snap_protect(tosnap.c_str());
+    } 
   }
 
   idiffctx.finish(r);
-- 
2.39.5