Merge pull request #99 from davidgiven/extensions

Validate image extensions before reading, not after.
davidgiven · Aug 22, 2019 · 4056364 · 4056364
2 parents 28d0ce7 + 60bfe05
commit 4056364
Show file tree

Hide file tree

Showing 6 changed files with 80 additions and 16 deletions.
diff --git a/lib/globals.h b/lib/globals.h
@@ -22,7 +22,7 @@ extern void hexdumpForSrp16(std::ostream& stream, const Bytes& bytes);
 class Error
 {
 public:
-    ~Error()
+    [[ noreturn ]] ~Error()
     {
         std::cerr << "Error: " << _stream.str() << std::endl;
         exit(1);

diff --git a/lib/imagereader/imagereader.cc b/lib/imagereader/imagereader.cc
@@ -7,22 +7,43 @@
 #include "imagereader/imagereader.h"
 #include "fmt/format.h"
 
+std::map<std::string, ImageReader::Constructor> ImageReader::formats =
+{
+	{".adf", ImageReader::createImgImageReader},
+	{".d81", ImageReader::createImgImageReader},
+	{".img", ImageReader::createImgImageReader},
+};
+
 static bool ends_with(const std::string& value, const std::string& ending)
 {
     if (ending.size() > value.size())
         return false;
     return std::equal(ending.rbegin(), ending.rend(), value.rbegin());
 }
 
-std::unique_ptr<ImageReader> ImageReader::create(const ImageSpec& spec)
+ImageReader::Constructor ImageReader::findConstructor(const ImageSpec& spec)
 {
     const auto& filename = spec.filename;
 
-    if (ends_with(filename, ".img") || ends_with(filename, ".adf"))
-        return createImgImageReader(spec);
+	for (const auto& e : formats)
+	{
+		if (ends_with(filename, e.first))
+			return e.second;
+	}
+
+	return NULL;
+}
+
+std::unique_ptr<ImageReader> ImageReader::create(const ImageSpec& spec)
+{
+    verifyImageSpec(spec);
+    return findConstructor(spec)(spec);
+}
 
-    Error() << "unrecognised image filename extension";
-    return std::unique_ptr<ImageReader>();
+void ImageReader::verifyImageSpec(const ImageSpec& spec)
+{
+    if (!findConstructor(spec))
+        Error() << "unrecognised image filename extension";
 }
 
 ImageReader::ImageReader(const ImageSpec& spec):

diff --git a/lib/imagereader/imagereader.h b/lib/imagereader/imagereader.h
@@ -12,10 +12,21 @@ class ImageReader
 
 public:
     static std::unique_ptr<ImageReader> create(const ImageSpec& spec);
+	static void verifyImageSpec(const ImageSpec& spec);
 
 private:
+	typedef 
+		std::function<
+			std::unique_ptr<ImageReader>(const ImageSpec& spec)
+		>
+		Constructor;
+
+	static std::map<std::string, Constructor> formats;
+
     static std::unique_ptr<ImageReader> createImgImageReader(const ImageSpec& spec);
 
+	static Constructor findConstructor(const ImageSpec& spec);
+
 public:
 	virtual SectorSet readImage() = 0;
 

diff --git a/lib/imagewriter/imagewriter.cc b/lib/imagewriter/imagewriter.cc
@@ -7,26 +7,45 @@
 #include "imagewriter/imagewriter.h"
 #include "fmt/format.h"
 
+std::map<std::string, ImageWriter::Constructor> ImageWriter::formats =
+{
+	{".adf", ImageWriter::createImgImageWriter},
+	{".d64", ImageWriter::createD64ImageWriter},
+	{".d81", ImageWriter::createImgImageWriter},
+	{".img", ImageWriter::createImgImageWriter},
+	{".ldbs", ImageWriter::createLDBSImageWriter},
+};
+
 static bool ends_with(const std::string& value, const std::string& ending)
 {
     if (ending.size() > value.size())
         return false;
     return std::equal(ending.rbegin(), ending.rend(), value.rbegin());
 }
 
-std::unique_ptr<ImageWriter> ImageWriter::create(const SectorSet& sectors, const ImageSpec& spec)
+ImageWriter::Constructor ImageWriter::findConstructor(const ImageSpec& spec)
 {
     const auto& filename = spec.filename;
 
-    if (ends_with(filename, ".img") || ends_with(filename, ".adf"))
-        return createImgImageWriter(sectors, spec);
-	else if (ends_with(filename, ".ldbs"))
-		return createLDBSImageWriter(sectors, spec);
-	else if (ends_with(filename, ".d64"))
-		return createD64ImageWriter(sectors, spec);
+	for (const auto& e : formats)
+	{
+		if (ends_with(filename, e.first))
+			return e.second;
+	}
+
+	return NULL;
+}
+
+std::unique_ptr<ImageWriter> ImageWriter::create(const SectorSet& sectors, const ImageSpec& spec)
+{
+	verifyImageSpec(spec);
+	return findConstructor(spec)(sectors, spec);
+}
 
-    Error() << "unrecognised image filename extension";
-    return std::unique_ptr<ImageWriter>();
+void ImageWriter::verifyImageSpec(const ImageSpec& spec)
+{
+	if (!findConstructor(spec))
+		Error() << "unrecognised image filename extension";
 }
 
 ImageWriter::ImageWriter(const SectorSet& sectors, const ImageSpec& spec):

diff --git a/lib/imagewriter/imagewriter.h b/lib/imagewriter/imagewriter.h
@@ -12,15 +12,26 @@ class ImageWriter
 
 public:
     static std::unique_ptr<ImageWriter> create(const SectorSet& sectors, const ImageSpec& spec);
+	static void verifyImageSpec(const ImageSpec& filename);
 
 private:
+	typedef 
+		std::function<
+			std::unique_ptr<ImageWriter>(const SectorSet& sectors, const ImageSpec& spec)
+		>
+		Constructor;
+
+	static std::map<std::string, Constructor> formats;
+
     static std::unique_ptr<ImageWriter> createImgImageWriter(
 		const SectorSet& sectors, const ImageSpec& spec);
     static std::unique_ptr<ImageWriter> createLDBSImageWriter(
 		const SectorSet& sectors, const ImageSpec& spec);
     static std::unique_ptr<ImageWriter> createD64ImageWriter(
 		const SectorSet& sectors, const ImageSpec& spec);
 
+	static Constructor findConstructor(const ImageSpec& spec);
+
 public:
 	virtual void adjustGeometry();
 	void printMap();

diff --git a/lib/reader.cc b/lib/reader.cc
@@ -14,6 +14,7 @@
 #include "bytes.h"
 #include "decoders/rawbits.h"
 #include "track.h"
+#include "imagewriter/imagewriter.h"
 #include "fmt/format.h"
 
 FlagGroup readerFlags { &hardwareFluxSourceFlags, &fluxmapReaderFlags };
@@ -153,7 +154,8 @@ static void replace_sector(std::unique_ptr<Sector>& replacing, Sector& replaceme
 void readDiskCommand(AbstractDecoder& decoder)
 {
 	const ImageSpec outputSpec(output);
-
+	ImageWriter::verifyImageSpec(outputSpec);
+
 	bool failures = false;
 	SectorSet allSectors;
 	auto tracks = readTracks();