Add support for writing .at3/.wav files with ATRAC3 (#30)

* Add support for writing .at3/.wav files with ATRAC3

* Check file size limit

* AT3/WAV export: Add big endian support, force packed structs

5 files changed, 213 insertions, 2 deletions

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index c0c22fe..6f63ef0 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -70,6 +70,7 @@ set(SOURCE_ATRACDENC_IMPL
     atrac/atrac1_bitalloc.cpp
     oma.cpp
     rm.cpp
+    at3.cpp
     atrac3denc.cpp
     atrac/atrac3.cpp
     atrac/atrac3_bitstream.cpp
diff --git a/src/at3.cpp b/src/at3.cpp
new file mode 100644
index 0000000..f4afc4f
--- /dev/null
+++ b/src/at3.cpp
@@ -0,0 +1,164 @@
+/*
+ * This file is part of AtracDEnc.
+ *
+ * AtracDEnc is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * AtracDEnc is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with AtracDEnc; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+#include "at3.h"
+
+#include "lib/endian_tools.h"
+#include <cstring>
+#include <iostream>
+#include <cmath>
+#include <assert.h>
+
+/*
+ * ATRAC3-in-WAV file format.
+ *
+ * Documented for example here:
+ *   - ffmpeg: libavcodec/atrac3.c (atrac3_decode_init() talks about "extradata")
+ *   - libnetmd: libnetmd/secure.c (netmd_write_wav_header() has "ATRAC extensions")
+ */
+
+namespace {
+
+// Based on http://soundfile.sapp.org/doc/WaveFormat/ + ffmpeg/libnetmd docs
+#ifdef _MSC_VER
+#pragma pack(push, 1)
+struct
+#else
+struct __attribute__((packed))
+#endif
+At3WaveHeader {
+    // "RIFF" "WAVE" header
+    char riff_chunk_id[4];
+    uint32_t chunk_size;
+    char riff_format[4];
+
+    // "fmt " subchunk
+    char subchunk1_id[4];
+    uint32_t subchunk1_size;
+
+    // WAVEFORMAT
+    uint16_t audio_format;
+    uint16_t num_channels;
+    uint32_t sample_rate;
+    uint32_t byte_rate;
+    uint16_t block_align;
+    uint16_t bits_per_sample;
+
+    // WAVEFORMATEX cbSize
+    uint16_t extradata_size; // 14
+
+    // atrac3 extradata
+    uint16_t unknown0; // always 1
+    uint32_t bytes_per_frame; // samples per channel (ffmpeg) or bytes per frame (libnetmd)
+    uint16_t coding_mode; // 1 = joint stereo, 0 = stereo
+    uint16_t coding_mode2; // same as <coding_mode>
+    uint16_t unknown1; // always 1
+    uint16_t unknown2; // always 0
+
+    // "data" subchunk
+    char subchunk2_id[4];
+    uint32_t subchunk2_size;
+};
+#ifdef _MSC_VER
+#pragma pack(pop)
+#endif
+
+class TAt3 : public ICompressedOutput {
+public:
+    TAt3(const std::string &filename, uint8_t numChannels,
+        uint32_t numFrames, uint32_t frameSize, bool jointStereo)
+        : fp(fopen(filename.c_str(), "wb"))
+    {
+        if (!fp) {
+            throw std::runtime_error("Cannot open file to write");
+        }
+
+        struct At3WaveHeader header;
+        memset(&header, 0, sizeof(header));
+
+        uint64_t file_size = sizeof(struct At3WaveHeader) + uint64_t(numFrames) * uint64_t(frameSize);
+
+        if (file_size >= UINT32_MAX) {
+            throw std::runtime_error("File size is too big for this file format");
+        }
+
+        memcpy(header.riff_chunk_id, "RIFF", 4);
+        header.chunk_size = swapbyte32_on_be(file_size);
+        memcpy(header.riff_format, "WAVE", 4);
+
+        memcpy(header.subchunk1_id, "fmt ", 4);
+        header.subchunk1_size = swapbyte32_on_be(offsetof(struct At3WaveHeader, subchunk2_id) -
+                                                 offsetof(struct At3WaveHeader, audio_format));
+
+        // libnetmd: #define NETMD_RIFF_FORMAT_TAG_ATRAC3 0x270
+        // mmreg.h (mingw-w64): WAVE_FORMAT_SONY_SCX 0x270
+        // riff.c (ffmpeg): AV_CODEC_ID_ATRAC3 0x0270
+        header.audio_format = swapbyte16_on_be(0x270);
+        header.num_channels = swapbyte16_on_be(numChannels);
+        header.sample_rate = swapbyte32_on_be(44100);
+        header.byte_rate = swapbyte32_on_be(frameSize * header.sample_rate / 1024);
+        header.block_align = swapbyte16_on_be(frameSize);
+        header.bits_per_sample = swapbyte16_on_be(0);
+        header.extradata_size = swapbyte16_on_be(offsetof(struct At3WaveHeader, subchunk2_id) -
+                                                 offsetof(struct At3WaveHeader, unknown0));
+
+        header.unknown0 = swapbyte16_on_be(1);
+        header.bytes_per_frame = swapbyte32_on_be(0x0010); // XXX
+        header.coding_mode = swapbyte16_on_be(jointStereo ? 0x0001 : 0x0000);
+        header.coding_mode2 = header.coding_mode; // already byte-swapped (if needed)
+        header.unknown1 = swapbyte16_on_be(1);
+        header.unknown2 = swapbyte16_on_be(0);
+
+        memcpy(header.subchunk2_id, "data", 4);
+        header.subchunk2_size = swapbyte32_on_be(numFrames * frameSize); // TODO
+
+        if (fwrite(&header, 1, sizeof(header), fp) != sizeof(header)) {
+            throw std::runtime_error("Cannot write WAV header to file");
+        }
+    }
+
+    virtual ~TAt3() override {
+        fclose(fp);
+    }
+
+    virtual void WriteFrame(std::vector<char> data) override {
+        if (fwrite(data.data(), 1, data.size(), fp) != data.size()) {
+            throw std::runtime_error("Cannot write AT3 data to file");
+        }
+    }
+
+    std::string GetName() const override {
+        return {};
+    }
+
+    uint8_t GetChannelNum() const override {
+        return 2;
+    }
+
+private:
+    FILE *fp;
+};
+
+} //namespace
+
+TCompressedOutputPtr
+CreateAt3Output(const std::string& filename, uint8_t numChannel,
+        uint32_t numFrames, uint32_t framesize, bool jointStereo)
+{
+    return std::unique_ptr<TAt3>(new TAt3(filename, numChannel, numFrames, framesize, jointStereo));
+}
diff --git a/src/at3.h b/src/at3.h
new file mode 100644
index 0000000..85f7a9e
--- /dev/null
+++ b/src/at3.h
@@ -0,0 +1,25 @@
+/*
+ * This file is part of AtracDEnc.
+ *
+ * AtracDEnc is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * AtracDEnc is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with AtracDEnc; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+#pragma once
+
+#include "compressed_io.h"
+
+TCompressedOutputPtr
+CreateAt3Output(const std::string& filename, uint8_t numChannel,
+        uint32_t numFrames, uint32_t framesize, bool jointStereo);
diff --git a/src/lib/endian_tools.h b/src/lib/endian_tools.h
index 519cca1..191b424 100644
--- a/src/lib/endian_tools.h
+++ b/src/lib/endian_tools.h
@@ -37,6 +37,22 @@ static inline uint16_t swapbyte16_on_le(uint16_t in) {
 #endif
 }
 
+static inline uint32_t swapbyte32_on_be(uint32_t in) {
+#ifdef BIGENDIAN_ORDER
+    return ((in & 0xff) << 24) | ((in & 0xff00) << 8) | ((in & 0xff0000) >> 8) | ((in & 0xff000000) >> 24);
+#else
+    return in;
+#endif
+}
+
+static inline uint16_t swapbyte16_on_be(uint16_t in) {
+#ifdef BIGENDIAN_ORDER
+    return ((in & 0xff) << 8) | ((in & 0xff00) >> 8);
+#else
+    return in;
+#endif
+}
+
 #ifdef __cplusplus
 
 static inline int16_t conv_ntoh(int16_t in) {
@@ -57,4 +73,4 @@ static inline uint32_t conv_ntoh(uint32_t in) {
 
 #endif
 
-#endif
\ No newline at end of file
+#endif
diff --git a/src/main.cpp b/src/main.cpp
index de087d2..708c24a 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -26,6 +26,7 @@
 #include "wav.h"
 #include "aea.h"
 #include "rm.h"
+#include "at3.h"
 #include "config.h"
 #include "atrac1denc.h"
 #include "atrac3denc.h"
@@ -217,7 +218,11 @@ static void PrepareAtrac3Encoder(const string& inFile,
 
     TCompressedOutputPtr omaIO;
 
-    if (ext == "rm") {
+    if (ext == "wav" || ext == "at3") {
+        omaIO = CreateAt3Output(outFile, numChannels, numFrames,
+                encoderSettings.ConteinerParams->FrameSz,
+                encoderSettings.ConteinerParams->Js);
+    } else if (ext == "rm") {
         omaIO = CreateRmOutput(outFile, "test", numChannels,
             numFrames, encoderSettings.ConteinerParams->FrameSz,
             encoderSettings.ConteinerParams->Js);