Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
highly optimized version of badblocks for testing USB memory sticks.
#include <iostream>
#include <random>
#include <cassert>
#include <chrono>
#include <sstream>
#include <cstdint>
#include <climits>
#include <algorithm>
#include <functional>
#include <map>
#include <iomanip>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#define SECTOR_SIZE 512
#define FAST_SIZE SECTOR_SIZE*20480
#define RANDOM_START_SEED 137
std::string bin2hex(const std::string &$str)
{
// from https://stackoverflow.com/a/18906469/1067003
std::string ossbuf;
ossbuf.reserve($str.size() * 2);
std::ostringstream out(std::move(ossbuf));
out << std::hex << std::setfill('0');
for (char c : $str)
{
out << std::setw(2) << uint_fast16_t(((unsigned char)c));
}
return out.str();
}
static void write_all(const int handle, const std::string& data)
{
size_t written = 0;
while (written < data.size())
{
const ssize_t written_last = write(handle, &data.at(written),data.size()-written);
if (written_last < 1)
{
throw std::runtime_error("failed to write data to handle after writing "+std::to_string(written)+" bytes! total bytes to write: " + std::to_string(data.size()) + " ftell: " +std::to_string(lseek(handle, 0, SEEK_CUR))+" errno: "+std::to_string(errno)+": "+std::string(strerror(errno)));
}
written += written_last;
}
assert(written==data.size());
}
static void fread_all(const int handle, const size_t bytes, std::string &buf)
{
buf.resize(bytes);
size_t read_total = 0;
size_t remaining = bytes;
while (remaining > 0)
{
const size_t read_last = read(handle, &buf[read_total], remaining);
if (read_last < 1)
{
throw std::runtime_error("failed to read all requested bytes! could only read "+std::to_string(read_total)+" of "+std::to_string(bytes)+" total bytes.");
}
remaining -= read_last;
read_total += read_last;
}
return;
}
static double microtime(void)
{
return (double(std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count()) / double(1000000));
}
std::string number_format(const double $number, const size_t $decimals = 0, const std::string &$dec_point = ".", const std::string &$thousands_sep = ",")
{
// i don't know how to implement this function "in a good way"
std::string ret;
{
std::ostringstream oss;
oss.precision(std::streamsize($decimals));
oss << std::fixed << $number;
ret = oss.str();
}
auto dotpos = ($decimals <= 0 ? std::string::npos : ret.find("."));
assert($decimals <= 0 || dotpos != std::string::npos);
{
uint_fast8_t last = 0;
for (size_t i = ($decimals <= 0 ? ret.size() : (dotpos)); i > 0; --i)
{
if (last == 3)
{
ret.insert(i, $thousands_sep);
dotpos+=$thousands_sep.size();
last = 1;
}
else
{
++last;
}
}
}
if ($decimals > 0)
{
ret.replace(dotpos, 1, $dec_point);
}
return ret;
}
static size_t get_size(const std::string& path)
{
int fp=open(path.c_str(),O_RDONLY);
if (fp==-1)
{
throw std::runtime_error("error: unable to open input file \""+path+"\" for reading!");
}
lseek(fp, 0, SEEK_END);
size_t ret = lseek(fp, 0, SEEK_CUR);// ftell()
close(fp);
if (ret < 2)
{
throw std::runtime_error("unable to get size of file! (ftell() returned bogus value "+std::to_string(ret)+")");
}
return ret;
}
static std::string get_string(const size_t size,const bool reset=false){
// could actually just use size 0 as a magic number for "reset", but...
assert((size==0 && reset==true) || (size>0 && reset==false));
static bool inited=false;
static bool reset_now=false;
static std::condition_variable cv;
static std::mutex mut;
static std::vector<std::string> vec;
if(!inited){
inited=true;
std::thread([](void)->void{
std::independent_bits_engine<std::default_random_engine, CHAR_BIT, uint8_t> rbe;
rbe.seed(RANDOM_START_SEED);
while(true){
{
std::unique_lock<std::mutex> lk(mut);
cv.wait(lk, []{return vec.size() == 0;});
}
if(reset_now){
reset_now=false;
rbe.seed(RANDOM_START_SEED);
}
std::string buf;
buf.resize(FAST_SIZE);
std::generate(buf.begin(), buf.end(), std::ref(rbe));
vec.push_back(std::move(buf));
cv.notify_all();
}
}).detach();
}
{
std::unique_lock<std::mutex> lk(mut);
cv.wait(lk, []{return vec.size() == 1;});
}
if(reset){
assert(size==0);
reset_now=true;
vec.pop_back();
cv.notify_all();
return "";
}
std::string ret=std::move(vec.back());
vec.pop_back();
cv.notify_all();
ret.resize(size);
return ret;
}
int main(int argc, char *argv[])
{
if (argc != 2)
{
std::cerr << "usage: " << argv[0] << " '/dev/usb'" << std::endl;
exit(EXIT_FAILURE);
}
get_string(0,true);// initialize the worker thread.
const int handle=open(argv[1],O_RDWR);
if(handle == -1)
{
std::cerr << "error: unable to open input file \"" << argv[1] << "\"" <<std::endl;
exit(EXIT_FAILURE);
}
size_t error_count = 0;
size_t size = get_size(argv[1]);
const double mb = size / 1024 / 1024;
const double gb = mb / 1024;
std::cout << "apparent size: " << number_format(mb,3) << " megabytes, aka " << number_format(gb,3) << " gigabytes, i hope that is correct. - " << std::endl;
if (true)
{
std::cout << "first running overwrite-everything:" << std::endl;
// already pre-initialized it above, otherwise: get_string(0,true);
const double starttime=microtime();
for (size_t i = 0; i < size;)
{
const double start=microtime();
const std::string buf=get_string(std::min(size_t(FAST_SIZE),size_t(size-i)));
write_all(handle, buf);
i += buf.size();
std::cout << "\r" << number_format(i/1024/1024,0) << "MB total, " << number_format((double(i)/double(size))*double(100),5) << "% - total average: " << number_format((i/(microtime()-starttime))/1024/1024,3) << "MB/s - right now: " << number_format((buf.size()/(microtime()-start))/1024/1024,3 ) << "MB/s" << std::flush;
}
std::cout << std::endl;
}
std::cout << "all data written! ... flushing.." << std::flush;
fsync(handle);
lseek(handle, 0, SEEK_SET);
// echo ". done. attempting dropcaches..";
// if (!is_writable('/proc/sys/vm/drop_caches')) {
// echo " failed because /proc/sys/vm/drop_caches is not writable.\n";
// } elseif (2 !== file_put_contents("/proc/sys/vm/drop_caches", "3\n")) {
// echo "failed because could not write 3.\n";
// } else {
// echo ". success!\n";
// }
if (true)
{
get_string(0,true);
std::cout << "now verifying.." << std::endl;
const double starttime=microtime();
size_t corrupted_bytes_total=0;
std::map<size_t,size_t> corrupted_sectors;
std::string read;
read.reserve(FAST_SIZE);
for (size_t i = 0; i < size;)
{
const double start=microtime();
read.resize(std::min(size_t(FAST_SIZE),size_t(size-i)));
fread_all(handle, read.size(),read);
// if we fread() before getting the comparison string, the string generator thread will have time to generate the next string,
// so make sure to read first then fetch the comparison string later.
const std::string buf=get_string(read.size());
if(buf!=read)
{
//error, and entering slow un-optimized code.
if(buf.size()!=read.size())
{
throw std::runtime_error("WTF buf.size!=read.size !!!!!!!");
}
for(size_t ii=0,max=std::min(buf.size(),read.size()); ii<max; ++ii)
{
if(!(read[ii]==buf[ii]))
{
std::cout << "\rERROR at sector " << size_t(std::floor((i+ii)/512)) << ": byte corrupted: wrote 0x" << bin2hex(buf.substr(ii,1)) << " read 0x" << bin2hex(read.substr(ii,1)) << " " << std::endl;
++error_count;
++corrupted_bytes_total;
corrupted_sectors[size_t(floor((i+ii)/512))]+=1;
}
}
}
i += read.size();
assert(read.size() == buf.size());
std::cout << "\r" << number_format(i/1024/1024,3) << "MB total, " << number_format((double(i)/double(size))*double(100),5) << "% - total average: " << number_format((double(i)/double(microtime()-starttime))/double(1024)/double(1024),3) << "MB/s - right now: " << number_format((double(buf.size())/double(microtime()-start))/double(1024)/double(1024),3 ) << "MB/s" << std::flush;
}
std::cout << std::endl;
std::cout << "verification finished! total corrupted sectors: " << corrupted_sectors.size() << " total corrupted bytes: " << corrupted_bytes_total << std::endl;
std::cout << "list of corrupted sectors: " << std::endl;
for (auto const& sec : corrupted_sectors)
{
std::cout << "sector " << sec.first << ": " << sec.second << "x corrupted byte(s)"<< std::endl;
}
if(true){
if(corrupted_sectors.size()==0){
const std::string cmd="dd conv=notrunc,fsync bs=10M count=1 status=progress if=/dev/zero of="+std::string(argv[1]);
std::cout << "0 corrupted sectors... cleaning up, executing \"" << cmd << "\"" << std::endl;
const int ret=system(cmd.c_str());
if(ret!=0){
std::cerr << "ERROR! dd did not return \"0\" ! it returned \"" << ret << "\"" << std::endl;
}
}
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.