From fb233771b5a24f77f785bb8509f1123f537ea95e Mon Sep 17 00:00:00 2001
From: Birk Tjelmeland <birktjelmeland@yahoo.no>
Date: Sat, 9 Feb 2019 20:36:40 +0100
Subject: [PATCH] Fix small errors and clean up

---
 Cargo.toml             |  3 +++
 src/encoder/mod.rs     | 43 +++++++++++++++++++++---------------------
 tests/encode_images.rs | 16 +++++++---------
 3 files changed, 32 insertions(+), 30 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 90a298e..838e0d6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,3 +19,6 @@ byteorder = "1.2"
 lzw = "0.10"
 num-derive = "0.2"
 num-traits = "0.2"
+
+[dev-dependencies]
+tempfile = "3.0"
diff --git a/src/encoder/mod.rs b/src/encoder/mod.rs
index fbfa9ba..eb28ef8 100644
--- a/src/encoder/mod.rs
+++ b/src/encoder/mod.rs
@@ -1,9 +1,9 @@
-use std::io::{self, Write, Seek};
-use std::collections::{BTreeMap, HashMap};
-use byteorder::{ByteOrder, WriteBytesExt, NativeEndian};
+use std::io::{Write, Seek};
+use std::collections::BTreeMap;
+use byteorder::{NativeEndian};
 
 use crate::decoder::ifd;
-use crate::error::TiffResult;
+use crate::error::{TiffResult, TiffFormatError, TiffError};
 
 mod writer;
 use self::writer::*;
@@ -105,13 +105,14 @@ impl TiffValue for str {
     }
 
     fn write<W: Write>(&self, writer: &mut TiffWriter<W>) -> TiffResult<()> {
-        for x in self.chars() {
-            if x.is_ascii() {
-                writer.write_u8(x as u8)?;
-            }
+        if self.is_ascii() {
+            writer.write_bytes(self.as_bytes())?;
+            writer.write_u8(0)?;
+            Ok(())
+        }
+        else {
+            Err(TiffError::FormatError(TiffFormatError::InvalidTag))
         }
-        writer.write_u8(0)?;
-        Ok(())
     }
 }
 
@@ -143,14 +144,12 @@ impl ColorType for RGBA8 {
 
 pub struct TiffEncoder<W> {
     writer: TiffWriter<W>,
-    offset: usize,
 }
 
 impl<W: Write + Seek> TiffEncoder<W> {
     pub fn new(writer: W) -> TiffResult<TiffEncoder<W>> {
         let mut encoder = TiffEncoder {
             writer: TiffWriter::new(writer),
-            offset: 0,
         };
         
         NativeEndian::write_header(&mut encoder.writer)?;
@@ -200,13 +199,15 @@ impl<'a, W: Write + Seek> DirectoryEncoder<'a, W> {
         self.ifd.insert(tag.to_u16(), (<T>::FIELD_TYPE, value.count(), bytes));
     }
 
-    pub fn modify_tag<T: TiffValue>(&mut self, tag: ifd::Tag, offset: u64, value: T) {
-        // TODO: handle errors
-        let bytes = &mut self.ifd.get_mut(&tag.to_u16()).unwrap().2;
+    pub fn modify_tag<T: TiffValue>(&mut self, tag: ifd::Tag, offset: u64, value: T) -> TiffResult<()> {
+        let bytes = &mut self.ifd.get_mut(&tag.to_u16())
+            .ok_or(TiffError::FormatError(TiffFormatError::RequiredTagNotFound(tag)))?.2;
+
         let mut writer = TiffWriter::new(std::io::Cursor::new(bytes));
 
-        writer.goto_offset(offset).unwrap();
-        value.write(&mut writer).unwrap();
+        writer.goto_offset(offset)?;
+        value.write(&mut writer)?;
+        Ok(())
     }
 
     fn write_directory(&mut self) -> TiffResult<u64> {
@@ -241,7 +242,7 @@ impl<'a, W: Write + Seek> DirectoryEncoder<'a, W> {
 
     pub fn write_data<T: TiffValue>(&mut self, value: T) -> TiffResult<u64> {
         let offset = self.writer.offset();
-        value.write(&mut self.writer);
+        value.write(&mut self.writer)?;
         Ok(offset)
     }
 
@@ -330,14 +331,14 @@ impl<'a, W: Write + Seek, T: ColorType> ImageEncoder<'a, W, T> {
     pub fn write_strip(&mut self, value: &[T::Inner]) -> TiffResult<()> {
         // TODO: Compression
         let offset = self.encoder.write_data(value)?;
-        self.encoder.modify_tag(ifd::Tag::StripOffsets, self.strip_idx as u64 * 4, offset as u32);
-        self.encoder.modify_tag(ifd::Tag::StripByteCounts, self.strip_idx as u64 * 4, value.bytes() as u32);
+        self.encoder.modify_tag(ifd::Tag::StripOffsets, self.strip_idx as u64 * 4, offset as u32)?;
+        self.encoder.modify_tag(ifd::Tag::StripByteCounts, self.strip_idx as u64 * 4, value.bytes() as u32)?;
 
         self.strip_idx += 1;
         Ok(())
     }
 
-    pub fn finish(mut self) -> TiffResult<()> {
+    pub fn finish(self) -> TiffResult<()> {
         self.encoder.finish()
     }
 }
diff --git a/tests/encode_images.rs b/tests/encode_images.rs
index 2dbbf29..c70d909 100644
--- a/tests/encode_images.rs
+++ b/tests/encode_images.rs
@@ -1,11 +1,11 @@
 extern crate tiff;
+extern crate tempfile;
 
 use tiff::ColorType;
-use tiff::decoder::{Decoder, DecodingResult, ifd::Tag, ifd::Value};
+use tiff::decoder::{Decoder, DecodingResult};
 use tiff::encoder::{TiffEncoder, RGB8};
 
-use std::fs::File;
-
+use std::io::{Seek, SeekFrom};
 
 #[test]
 fn encode_decode() {
@@ -18,9 +18,9 @@ fn encode_decode() {
             image_data.push(val);
         }
     }
+    let mut file = tempfile::tempfile().unwrap();
     {
-        let file = std::fs::File::create("test.tiff").unwrap();
-        let mut tiff = TiffEncoder::new(file).unwrap();
+        let mut tiff = TiffEncoder::new(&mut file).unwrap();
 
         let mut image = tiff.new_image::<RGB8>(100, 100).unwrap();
     
@@ -33,8 +33,8 @@ fn encode_decode() {
         image.finish().unwrap();
     }
     {
-        let file = File::open("test.tiff").unwrap();
-        let mut decoder = Decoder::new(file).unwrap();
+        file.seek(SeekFrom::Start(0)).unwrap();
+        let mut decoder = Decoder::new(&mut file).unwrap();
         assert_eq!(decoder.colortype().unwrap(), ColorType::RGB(8));
         assert_eq!(decoder.dimensions().unwrap(), (100, 100));
         if let DecodingResult::U8(img_res) = decoder.read_image().unwrap() {
@@ -44,6 +44,4 @@ fn encode_decode() {
             panic!("Wrong data type");
         }
     }
-    
-    std::fs::remove_file("test.tiff").unwrap();
 }