Merge pull request trapexit#31 from trapexit/burnin-fix

Fix burnin so it proper restores original data if possible

src/bbf_burnin.cpp

@@ -32,6 +32,9 @@
 #include <errno.h>
 #include <stdint.h>
 
+typedef std::vector<char> CharVec;
+typedef std::vector<CharVec> CharVecVec;
+
 namespace l
 {
   static
@@ -57,28 +60,27 @@ namespace l
   write_read_compare(BlkDev         &blkdev_,
                      const uint64_t  stepping_,
                      const uint64_t  block_,
-                     char           *buf_,
-                     const size_t    buflen_,
+                     CharVec        &tmpbuf_,
                      const int       retries_,
-                     const char     *write_buf_)
+                     const CharVec  &write_buf_)
   {
     int rv;
 
     rv = -1;
     for(uint64_t i = 0; ((i <= retries_) && (rv < 0)); i++)
-      rv = blkdev_.write(block_,stepping_,write_buf_,buflen_);
+      rv = blkdev_.write(block_,stepping_,write_buf_);
 
     if(rv < 0)
       return rv;
 
     rv = -1;
     for(uint64_t i = 0; ((i <= retries_) && (rv < 0)); i++)
-      rv = blkdev_.read(block_,stepping_,buf_,buflen_);
+      rv = blkdev_.read(block_,stepping_,tmpbuf_);
 
     if(rv < 0)
       return rv;
 
-    rv = ::memcmp(write_buf_,buf_,buflen_);
+    rv = ::memcmp(&write_buf_[0],&tmpbuf_[0],tmpbuf_.size());
     if(rv != 0)
       return -EIO;
 
@@ -87,29 +89,29 @@ namespace l
 
   static
   int
-  burn_block(BlkDev         &blkdev_,
-             const uint64_t  stepping_,
-             const uint64_t  block_,
-             char           *buf_,
-             const size_t    buflen_,
-             const uint64_t  retries_,
-             const std::vector<std::vector<char> > &patterns_)
+  burn_block(BlkDev            &blkdev_,
+             const uint64_t     stepping_,
+             const uint64_t     block_,
+             CharVec           &tmpbuf_,
+             CharVec           &savebuf_,
+             const uint64_t     retries_,
+             const CharVecVec  &patterns_)
   {
     int rv;
 
     rv = -1;
     for(uint64_t i = 0; ((i <= retries_) && (rv < 0)); i++)
-      rv = blkdev_.read(block_,stepping_,buf_,buflen_);
+      rv = blkdev_.read(block_,stepping_,savebuf_);
 
     if(rv < 0)
-      ::memset(buf_,0,buflen_);
+      ::memset(&savebuf_[0],0,savebuf_.size());
 
     for(uint64_t i = 0; i < patterns_.size(); i++)
-      rv = l::write_read_compare(blkdev_,stepping_,block_,buf_,buflen_,retries_,&patterns_[i][0]);
+      rv = l::write_read_compare(blkdev_,stepping_,block_,tmpbuf_,retries_,patterns_[i]);
 
     rv = -1;
     for(uint64_t i = 0; ((i <= retries_) && (rv < 0)); i++)
-      rv = blkdev_.write(block_,stepping_,buf_,buflen_);
+      rv = blkdev_.write(block_,stepping_,savebuf_);
 
     return rv;
   }
@@ -120,8 +122,7 @@ namespace l
               const uint64_t         start_block_,
               const uint64_t         end_block_,
               const uint64_t         stepping_,
-              char                  *buf_,
-              const size_t           buflen_,
+              const uint64_t         buflen_,
               std::vector<uint64_t> &badblocks_,
               const uint64_t         max_errors_,
               const int              retries_)
@@ -130,9 +131,14 @@ namespace l
     uint64_t block;
     uint64_t stepping;
     InfoPrinter info;
-    std::vector<std::vector<char> > patterns;
+    CharVec tmpbuf;
+    CharVec savebuf;
+    CharVecVec patterns;
     const double start_time = Time::get_monotonic();
 
+    tmpbuf.resize(buflen_,0x00);
+    savebuf.resize(buflen_,0x00);
+
     patterns.resize(4);
     patterns[0].resize(buflen_,0x00);
     patterns[1].resize(buflen_,0x55);
@@ -156,7 +162,7 @@ namespace l
 
         stepping = l::trim_stepping(blkdev_,block,stepping_);
 
-        rv = l::burn_block(blkdev_,stepping,block,buf_,buflen_,retries_,patterns);
+        rv = l::burn_block(blkdev_,stepping,block,tmpbuf,savebuf,retries_,patterns);
 
         block += stepping;
         if(rv >= 0)
@@ -186,7 +192,6 @@ namespace l
   {
     int       rv;
     int       retries;
-    char     *buf;
     size_t    buflen;
     uint64_t  start_block;
     uint64_t  end_block;
@@ -225,17 +230,14 @@ namespace l
               << end_block
               << std::endl;
 
-    buf = new char[buflen];
     rv = l::burnin_loop(blkdev_,
-                     start_block,
-                     end_block,
-                     stepping,
-                     buf,
-                     buflen,
-                     badblocks_,
-                     opts_.max_errors,
-                     retries);
-    delete[] buf;
+                        start_block,
+                        end_block,
+                        stepping,
+                        buflen,
+                        badblocks_,
+                        opts_.max_errors,
+                        retries);
 
     std::cout << std::endl;
 

src/blkdev.cpp

@@ -256,6 +256,22 @@ BlkDev::read(const uint64_t  lba_,
   return -ENOTSUP;
 }
 
+int64_t
+BlkDev::read(const uint64_t     lba_,
+             const uint64_t     blocks_,
+             std::vector<char> &buf_)
+{
+  switch(_rw_type)
+    {
+    case ATA:
+      return ata_read(lba_,blocks_,&buf_[0],buf_.size());
+    case OS:
+      return os_read(lba_,blocks_,&buf_[0],buf_.size());
+    }
+
+  return -ENOTSUP;
+}
+
 int64_t
 BlkDev::write(const uint64_t  lba_,
               const uint64_t  blocks_,
@@ -273,6 +289,22 @@ BlkDev::write(const uint64_t  lba_,
   return -ENOTSUP;
 }
 
+int64_t
+BlkDev::write(const uint64_t           lba_,
+              const uint64_t           blocks_,
+              const std::vector<char> &buf_)
+{
+  switch(_rw_type)
+    {
+    case ATA:
+      return ata_write(lba_,blocks_,&buf_[0],buf_.size());
+    case OS:
+      return os_write(lba_,blocks_,&buf_[0],buf_.size());
+    }
+
+  return -ENOTSUP;
+}
+
 uint64_t
 BlkDev::block_stepping(void) const
 {

src/blkdev.hpp

@@ -21,10 +21,12 @@
 #include "sg.hpp"
 
 #include <string>
+#include <vector>
 
 #include <stdlib.h>
 #include <stdint.h>
 
+
 class BlkDev
 {
 public:
@@ -80,11 +82,19 @@ class BlkDev
                void           *buf,
                const uint64_t  buflen);
 
+  int64_t read(const uint64_t     lba,
+               const uint64_t     blocks,
+               std::vector<char> &buf);
+
   int64_t write(const uint64_t  lba,
                 const uint64_t  blocks,
                 const void     *buf,
                 const uint64_t  buflen);
 
+  int64_t write(const uint64_t           lba,
+                const uint64_t           blocks,
+                const std::vector<char> &buf);
+
 public:
   int sync(void);
   int write_flagged_uncorrectable(const uint64_t lba_,

src/captcha.cpp

@@ -30,7 +30,8 @@ namespace captcha
 
     if(blkdev.has_identity())
       ss << blkdev.identity().serial_number;
-    else
+
+    if(ss.gcount() == 0)
       ss << blkdev.size_in_bytes();
 
     return ss.str();