4 files changed, 180 insertions, 16 deletions
diff --git a/gr-blocks/include/gnuradio/blocks/wavfile_sink.h b/gr-blocks/include/gnuradio/blocks/wavfile_sink.h
index d996339300..144f446cce 100644
--- a/gr-blocks/include/gnuradio/blocks/wavfile_sink.h
+++ b/gr-blocks/include/gnuradio/blocks/wavfile_sink.h
@@ -40,7 +40,8 @@ public:
     static sptr make(const char* filename,
                      int n_channels,
                      unsigned int sample_rate,
-                     int bits_per_sample = 16);
+                     int bits_per_sample = 16,
+                     bool append = false);
 
     /*!
      * \brief Opens a new file and writes a WAV header. Thread-safe.
@@ -67,6 +68,13 @@ public:
      * is kept.
      */
     virtual void set_bits_per_sample(int bits_per_sample) = 0;
+
+    /*!
+     * \brief Enable appending to an existing file instead of
+     * creating it. This will not affect the WAV file currently
+     * opened (see set_sample_rate()).
+     */
+    virtual void set_append(bool append) = 0;
 };
 
 } /* namespace blocks */
diff --git a/gr-blocks/lib/wavfile_sink_impl.cc b/gr-blocks/lib/wavfile_sink_impl.cc
index 431919e7ed..09ff801822 100644
--- a/gr-blocks/lib/wavfile_sink_impl.cc
+++ b/gr-blocks/lib/wavfile_sink_impl.cc
@@ -46,20 +46,23 @@ namespace blocks {
 wavfile_sink::sptr wavfile_sink::make(const char* filename,
                                       int n_channels,
                                       unsigned int sample_rate,
-                                      int bits_per_sample)
+                                      int bits_per_sample,
+                                      bool append)
 {
     return gnuradio::get_initial_sptr(new wavfile_sink_impl(
-        filename, n_channels, sample_rate, bits_per_sample));
+        filename, n_channels, sample_rate, bits_per_sample, append));
 }
 
 wavfile_sink_impl::wavfile_sink_impl(const char* filename,
                                      int n_channels,
                                      unsigned int sample_rate,
-                                     int bits_per_sample)
+                                     int bits_per_sample,
+                                     bool append)
     : sync_block("wavfile_sink",
                  io_signature::make(1, n_channels, sizeof(float)),
                  io_signature::make(0, 0, 0)),
       d_h{}, // Init with zeros
+      d_append(append),
       d_fp(nullptr),
       d_new_fp(nullptr),
       d_updated(false)
@@ -81,12 +84,23 @@ bool wavfile_sink_impl::open(const char* filename)
     gr::thread::scoped_lock guard(d_mutex);
 
     // we use the open system call to get access to the O_LARGEFILE flag.
-    int flags = OUR_O_LARGEFILE | OUR_O_BINARY | O_CREAT | O_WRONLY | O_TRUNC;
-    int fd;
+    int flags = OUR_O_LARGEFILE | OUR_O_BINARY;
+
+    if (!d_append) {
+        // We are generating a new file.
+        flags |= O_CREAT | O_WRONLY | O_TRUNC;
+    } else {
+        flags |= O_RDWR;
+    }
 
+    int fd;
     if ((fd = ::open(filename, flags, 0664)) < 0) {
-        GR_LOG_ERROR(d_logger,
-                     boost::format("::open: %s: %s") % filename % strerror(errno));
+        if (errno == ENOENT) {
+            throw std::runtime_error("WAV append mode requires target file to exist");
+        } else {
+            GR_LOG_ERROR(d_logger,
+                         boost::format("::open: %s: %s") % filename % strerror(errno));
+        }
         return false;
     }
 
@@ -95,7 +109,7 @@ bool wavfile_sink_impl::open(const char* filename)
         d_new_fp = nullptr;
     }
 
-    if (!(d_new_fp = fdopen(fd, "wb"))) {
+    if (!(d_new_fp = fdopen(fd, d_append ? "r+b" : "wb"))) {
         GR_LOG_ERROR(d_logger,
                      boost::format("fdopen: %s: %s") % filename % strerror(errno));
 
@@ -103,12 +117,20 @@ bool wavfile_sink_impl::open(const char* filename)
         return false;
     }
 
-    d_h.first_sample_pos = 44;
-    if (!wavheader_write(
-            d_new_fp, d_h.sample_rate, d_h.nchans, d_bytes_per_sample_new)) {
-        GR_LOG_ERROR(d_logger, boost::format("could not save WAV header"));
-        fclose(d_new_fp);
-        return false;
+    if (d_append) {
+        // We are appending to an existing file, be extra careful here.
+        if (!check_append_compat_file(d_new_fp)) {
+            fclose(d_new_fp);
+            return false;
+        }
+    } else {
+        d_h.first_sample_pos = 44;
+        if (!wavheader_write(
+                d_new_fp, d_h.sample_rate, d_h.nchans, d_bytes_per_sample_new)) {
+            GR_LOG_ERROR(d_logger, boost::format("could not save WAV header"));
+            fclose(d_new_fp);
+            return false;
+        }
     }
 
     d_updated = true;
@@ -116,6 +138,47 @@ bool wavfile_sink_impl::open(const char* filename)
     return true;
 }
 
+bool wavfile_sink_impl::check_append_compat_file(FILE* fp)
+{
+
+    if (d_bytes_per_sample_new != d_h.bytes_per_sample) {
+        GR_LOG_ERROR(d_logger,
+                     "bytes_per_sample is not allowed to change in append mode");
+        return false;
+    }
+
+    wav_header_info h_tmp{};
+    std::swap(d_h, h_tmp);
+
+    if (!wavheader_parse(fp, d_h)) {
+        GR_LOG_ERROR(d_logger, "invalid or incompatible WAV file");
+        return false;
+    }
+
+    if (d_h.sample_rate != h_tmp.sample_rate || d_h.nchans != h_tmp.nchans ||
+        d_h.bytes_per_sample != h_tmp.bytes_per_sample) {
+        GR_LOG_ERROR(d_logger,
+                     "existing WAV file is incompatible with configured options");
+        return false;
+    }
+
+    // TODO: use GR_FSEEK, GR_FTELL.
+    if (fseek(d_new_fp, 0, SEEK_END) != 0) {
+        return false; // This can only happen if the file disappears under our feet.
+    }
+
+    long file_size = ftell(fp);
+    if (file_size - d_h.first_sample_pos != d_h.data_chunk_size) {
+        // This is complicated to properly implement for too little benefit.
+        GR_LOG_ERROR(d_logger,
+                     "existing WAV file is incompatible (extra chunks at the end)");
+        return false;
+    }
+
+    return true;
+}
+
+
 void wavfile_sink_impl::close()
 {
     gr::thread::scoped_lock guard(d_mutex);
@@ -230,6 +293,12 @@ void wavfile_sink_impl::set_bits_per_sample_unlocked(int bits_per_sample)
     d_bytes_per_sample_new = bits_per_sample / 8;
 }
 
+void wavfile_sink_impl::set_append(bool append)
+{
+    gr::thread::scoped_lock guard(d_mutex);
+    d_append = append;
+}
+
 void wavfile_sink_impl::set_sample_rate(unsigned int sample_rate)
 {
     gr::thread::scoped_lock guard(d_mutex);
diff --git a/gr-blocks/lib/wavfile_sink_impl.h b/gr-blocks/lib/wavfile_sink_impl.h
index c81c9658a7..d6e687dd03 100644
--- a/gr-blocks/lib/wavfile_sink_impl.h
+++ b/gr-blocks/lib/wavfile_sink_impl.h
@@ -22,6 +22,7 @@ class wavfile_sink_impl : public wavfile_sink
 private:
     wav_header_info d_h;
     int d_bytes_per_sample_new;
+    bool d_append;
 
     float d_max_sample_val;
     float d_min_sample_val;
@@ -59,6 +60,13 @@ private:
      */
     void close_wav();
 
+    /*!
+     * \brief Checks if the given WAV file is compatible with the current
+     * configuration in order to open it to append information.
+     * This also finds the value of d_first_sample_pos.
+     */
+    bool check_append_compat_file(FILE* fp);
+
 protected:
     bool stop();
 
@@ -66,7 +74,8 @@ public:
     wavfile_sink_impl(const char* filename,
                       int n_channels,
                       unsigned int sample_rate,
-                      int bits_per_sample);
+                      int bits_per_sample,
+                      bool append);
     ~wavfile_sink_impl();
 
     bool open(const char* filename);
@@ -74,6 +83,7 @@ public:
 
     void set_sample_rate(unsigned int sample_rate);
     void set_bits_per_sample(int bits_per_sample);
+    void set_append(bool append) override;
 
     int bits_per_sample();
     unsigned int sample_rate();
diff --git a/gr-blocks/python/blocks/qa_wavfile.py b/gr-blocks/python/blocks/qa_wavfile.py
index ba2d80dc66..f22934d71b 100644
--- a/gr-blocks/python/blocks/qa_wavfile.py
+++ b/gr-blocks/python/blocks/qa_wavfile.py
@@ -71,5 +71,82 @@ class test_wavefile(gr_unittest.TestCase):
         self.assertEqual(in_data[:g_extra_header_offset] + \
                          in_data[g_extra_header_offset + g_extra_header_len:], out_data)
 
+    def test_003_checkwav_append_copy(self):
+        infile  = g_in_file
+        outfile = "test_out_append.wav"
+
+        # 1. Copy input to output
+        from shutil import copyfile
+        copyfile(infile, outfile)
+
+        # 2. append copy
+        wf_in  = blocks.wavfile_source(infile)
+        wf_out = blocks.wavfile_sink(outfile,
+                                     wf_in.channels(),
+                                     wf_in.sample_rate(),
+                                     wf_in.bits_per_sample(),
+                                     True)
+        self.tb.connect(wf_in, wf_out)
+        self.tb.run()
+        wf_out.close()
+
+        # 3. append halved copy
+        wf_in  = blocks.wavfile_source(infile)
+        halver = blocks.multiply_const_ff(0.5)
+        wf_out = blocks.wavfile_sink(outfile,
+                                     wf_in.channels(),
+                                     wf_in.sample_rate(),
+                                     wf_in.bits_per_sample(),
+                                     True)
+        self.tb.connect(wf_in, halver, wf_out)
+        self.tb.run()
+        wf_out.close()
+
+        # Test file validity and read data.
+        import wave
+        try:
+            # In
+            with wave.open(infile, 'rb') as w_in:
+                in_params = w_in.getparams()
+                data_in = wav_read_frames(w_in)
+            # Out
+            with wave.open(outfile, 'rb') as w_out:
+                out_params = w_out.getparams()
+                data_out = wav_read_frames(w_out)
+        except:
+            raise AssertionError('Invalid WAV file')
+
+        # Params must be equal except in size:
+        expected_params = in_params._replace(nframes=3*in_params.nframes)
+        self.assertEqual(out_params, expected_params)
+
+        # Part 1
+        self.assertEqual(data_in, data_out[:len(data_in)])
+
+        # Part 2
+        self.assertEqual(data_in, data_out[len(data_in):2*len(data_in)])
+
+        # Part 3
+        data_in_halved = [int(round(d/2)) for d in data_in]
+        self.assertEqual(data_in_halved, data_out[2*len(data_in):])
+
+    def test_003_checkwav_append_non_existent_should_error(self):
+        outfile = "no_file.wav"
+
+        with self.assertRaisesRegex(RuntimeError, 'WAV append mode requires target file to exist'):
+            blocks.wavfile_sink(outfile, 1, 44100, 16, True)
+
+
+def wav_read_frames(w):
+    import struct
+    # grouper from itertools recipes.
+    grouper = lambda iterable, n: list(zip(* ([iter(iterable)] * n) ))
+    assert w.getsampwidth() == 2  # Assume 16 bits
+    return [
+        struct.unpack('<h', bytes(frame_g))[0]
+        for frame_g in grouper(w.readframes(w.getnframes()), 2)
+    ]
+
+
 if __name__ == '__main__':
     gr_unittest.run(test_wavefile)