Nikolay Sivov : xmllite/writer: Fix Flush() behaviour on partial writes.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: 81d2516449f0c2df39d76b493230853e8afa7a91
URL:    http://source.winehq.org/git/wine.git/?a=commit;h=81d2516449f0c2df39d76b493230853e8afa7a91

Author: Nikolay Sivov <nsivov at codeweavers.com>
Date:   Fri May 16 12:42:52 2014 +0400

xmllite/writer: Fix Flush() behaviour on partial writes.

---

 dlls/xmllite/tests/writer.c |   78 +++++++++++++++++++++++++++++++++++++++++++
 dlls/xmllite/writer.c       |   33 +++++++++---------
 2 files changed, 93 insertions(+), 18 deletions(-)

diff --git a/dlls/xmllite/tests/writer.c b/dlls/xmllite/tests/writer.c
index b6b8fe3..d7cfe44 100644
--- a/dlls/xmllite/tests/writer.c
+++ b/dlls/xmllite/tests/writer.c
@@ -67,6 +67,53 @@ static const IUnknownVtbl testoutputvtbl = {
 
 static IUnknown testoutput = { &testoutputvtbl };
 
+static HRESULT WINAPI teststream_QueryInterface(ISequentialStream *iface, REFIID riid, void **obj)
+{
+    if (IsEqualIID(riid, &IID_IUnknown) || IsEqualIID(riid, &IID_ISequentialStream))
+    {
+        *obj = iface;
+        return S_OK;
+    }
+
+    *obj = NULL;
+    return E_NOINTERFACE;
+}
+
+static ULONG WINAPI teststream_AddRef(ISequentialStream *iface)
+{
+    return 2;
+}
+
+static ULONG WINAPI teststream_Release(ISequentialStream *iface)
+{
+    return 1;
+}
+
+static HRESULT WINAPI teststream_Read(ISequentialStream *iface, void *pv, ULONG cb, ULONG *pread)
+{
+    ok(0, "unexpected call\n");
+    return E_NOTIMPL;
+}
+
+static ULONG g_write_len;
+static HRESULT WINAPI teststream_Write(ISequentialStream *iface, const void *pv, ULONG cb, ULONG *written)
+{
+    g_write_len = cb;
+    *written = cb;
+    return S_OK;
+}
+
+static const ISequentialStreamVtbl teststreamvtbl =
+{
+    teststream_QueryInterface,
+    teststream_AddRef,
+    teststream_Release,
+    teststream_Read,
+    teststream_Write
+};
+
+static ISequentialStream teststream = { &teststreamvtbl };
+
 static void test_writer_create(void)
 {
     HRESULT hr;
@@ -235,6 +282,36 @@ static void test_writestartdocument(void)
     IXmlWriter_Release(writer);
 }
 
+static void test_flush(void)
+{
+    IXmlWriter *writer;
+    HRESULT hr;
+
+    hr = pCreateXmlWriter(&IID_IXmlWriter, (void**)&writer, NULL);
+    ok(hr == S_OK, "Expected S_OK, got %08x\n", hr);
+
+    hr = IXmlWriter_SetOutput(writer, (IUnknown*)&teststream);
+    ok(hr == S_OK, "got 0x%08x\n", hr);
+
+    hr = IXmlWriter_WriteStartDocument(writer, XmlStandalone_Yes);
+    ok(hr == S_OK, "got 0x%08x\n", hr);
+
+    g_write_len = 0;
+    hr = IXmlWriter_Flush(writer);
+    ok(hr == S_OK, "got 0x%08x\n", hr);
+    ok(g_write_len > 0, "got %d\n", g_write_len);
+
+    g_write_len = 1;
+    hr = IXmlWriter_Flush(writer);
+    ok(hr == S_OK, "got 0x%08x\n", hr);
+    ok(g_write_len == 0, "got %d\n", g_write_len);
+
+    /* Release() flushes too */
+    g_write_len = 1;
+    IXmlWriter_Release(writer);
+    ok(g_write_len == 0, "got %d\n", g_write_len);
+}
+
 START_TEST(writer)
 {
     if (!init_pointers())
@@ -243,4 +320,5 @@ START_TEST(writer)
     test_writer_create();
     test_writeroutput();
     test_writestartdocument();
+    test_flush();
 }
diff --git a/dlls/xmllite/writer.c b/dlls/xmllite/writer.c
index 0a1d361..de5f755 100644
--- a/dlls/xmllite/writer.c
+++ b/dlls/xmllite/writer.c
@@ -63,7 +63,6 @@ typedef struct
     IMalloc *imalloc;
     xml_encoding encoding;
     struct output_buffer buffer;
-    ULONG stream_written;
 } xmlwriteroutput;
 
 static const struct IUnknownVtbl xmlwriteroutputvtbl;
@@ -225,11 +224,6 @@ static HRESULT write_output_buffer_quoted(xmlwriteroutput *output, const WCHAR *
     return S_OK;
 }
 
-static inline void reset_output_buffer(xmlwriteroutput *output)
-{
-    output->stream_written = 0;
-}
-
 static void writeroutput_release_stream(xmlwriteroutput *writeroutput)
 {
     if (writeroutput->stream) {
@@ -253,7 +247,7 @@ static inline HRESULT writeroutput_query_for_stream(xmlwriteroutput *writeroutpu
 static HRESULT writeroutput_flush_stream(xmlwriteroutput *output)
 {
     struct output_buffer *buffer;
-    ULONG written = 0, len;
+    ULONG written, offset = 0;
     HRESULT hr;
 
     if (!output || !output->stream)
@@ -261,17 +255,20 @@ static HRESULT writeroutput_flush_stream(xmlwriteroutput *output)
 
     buffer = &output->buffer;
 
-    len = buffer->written - output->stream_written;
-    if (!len)
-        return S_OK;
+    /* It will loop forever until everything is written or an error occured. */
+    do {
+        written = 0;
+        hr = ISequentialStream_Write(output->stream, buffer->data + offset, buffer->written, &written);
+        if (FAILED(hr)) {
+            WARN("write to stream failed (0x%08x)\n", hr);
+            buffer->written = 0;
+            return hr;
+        }
 
-    hr = ISequentialStream_Write(output->stream, buffer->data + output->stream_written, len, &written);
-    if (FAILED(hr)) {
-        WARN("write to stream failed (0x%08x)\n", hr);
-        return hr;
-    }
+        offset += written;
+        buffer->written -= written;
+    } while (buffer->written > 0);
 
-    output->stream_written += written;
     return S_OK;
 }
 
@@ -309,6 +306,8 @@ static ULONG WINAPI xmlwriter_Release(IXmlWriter *iface)
     ref = InterlockedDecrement(&This->ref);
     if (ref == 0) {
         IMalloc *imalloc = This->imalloc;
+
+        IXmlWriter_Flush(iface);
         if (This->output) IUnknown_Release(&This->output->IXmlWriterOutput_iface);
         writer_free(This, This);
         if (imalloc) IMalloc_Release(imalloc);
@@ -327,7 +326,6 @@ static HRESULT WINAPI xmlwriter_SetOutput(IXmlWriter *iface, IUnknown *output)
     TRACE("(%p)->(%p)\n", This, output);
 
     if (This->output) {
-        reset_output_buffer(This->output);
         writeroutput_release_stream(This->output);
         IUnknown_Release(&This->output->IXmlWriterOutput_iface);
         This->output = NULL;
@@ -858,7 +856,6 @@ HRESULT WINAPI CreateXmlWriterOutputWithEncodingName(IUnknown *stream,
     if (imalloc) IMalloc_AddRef(imalloc);
     writeroutput->encoding = parse_encoding_name(encoding ? encoding : utf8W, -1);
     writeroutput->stream = NULL;
-    writeroutput->stream_written = 0;
     hr = init_output_buffer(writeroutput);
     if (FAILED(hr)) {
         IUnknown_Release(&writeroutput->IXmlWriterOutput_iface);