amstream: Implement IDirectDrawMediaStream::SetDirectDraw().

Signed-off-by: Gijs Vermeulen <gijsvrm@gmail.com>
Signed-off-by: Zebediah Figura <z.figura12@gmail.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>

dlls/amstream/ddrawstream.c

@@ -35,6 +35,7 @@ struct ddraw_stream
     IMemInputPin IMemInputPin_iface;
     IPin IPin_iface;
     LONG ref;
+    LONG sample_refs;
 
     IMultiMediaStream* parent;
     MSPID purpose_id;
@@ -396,11 +397,35 @@ static HRESULT WINAPI ddraw_IDirectDrawMediaStream_GetDirectDraw(IDirectDrawMedi
 }
 
 static HRESULT WINAPI ddraw_IDirectDrawMediaStream_SetDirectDraw(IDirectDrawMediaStream *iface,
-        IDirectDraw *pDirectDraw)
+        IDirectDraw *ddraw)
 {
-    FIXME("(%p)->(%p) stub!\n", iface, pDirectDraw);
+    struct ddraw_stream *stream = impl_from_IDirectDrawMediaStream(iface);
 
-    return E_NOTIMPL;
+    TRACE("stream %p, ddraw %p.\n", stream, ddraw);
+
+    EnterCriticalSection(&stream->cs);
+
+    if (stream->sample_refs)
+    {
+        HRESULT hr = (stream->ddraw == ddraw) ? S_OK : MS_E_SAMPLEALLOC;
+        LeaveCriticalSection(&stream->cs);
+        return hr;
+    }
+
+    if (stream->ddraw)
+        IDirectDraw_Release(stream->ddraw);
+
+    if (ddraw)
+    {
+        IDirectDraw_AddRef(ddraw);
+        stream->ddraw = ddraw;
+    }
+    else
+        stream->ddraw = NULL;
+
+    LeaveCriticalSection(&stream->cs);
+
+    return S_OK;
 }
 
 static HRESULT WINAPI ddraw_IDirectDrawMediaStream_CreateSample(IDirectDrawMediaStream *iface,
@@ -951,6 +976,7 @@ HRESULT ddraw_stream_create(IUnknown *outer, void **out)
     object->IMemInputPin_iface.lpVtbl = &ddraw_meminput_vtbl;
     object->IPin_iface.lpVtbl = &ddraw_sink_vtbl;
     object->ref = 1;
+    object->sample_refs = 0;
 
     InitializeCriticalSection(&object->cs);
 
@@ -1013,6 +1039,10 @@ static ULONG WINAPI ddraw_sample_Release(IDirectDrawStreamSample *iface)
 
     TRACE("(%p)->(): new ref = %u\n", iface, ref);
 
+    EnterCriticalSection(&sample->parent->cs);
+    InterlockedDecrement(&sample->parent->sample_refs);
+    LeaveCriticalSection(&sample->parent->cs);
+
     if (!ref)
     {
         if (sample->surface)
@@ -1131,6 +1161,10 @@ static HRESULT ddrawstreamsample_create(struct ddraw_stream *parent, IDirectDraw
     object->ref = 1;
     object->parent = parent;
 
+    EnterCriticalSection(&parent->cs);
+    InterlockedIncrement(&parent->sample_refs);
+    LeaveCriticalSection(&parent->cs);
+
     if (surface)
     {
         object->surface = surface;

dlls/amstream/tests/amstream.c

@@ -4798,8 +4798,8 @@ static void test_ddrawstream_getsetdirectdraw(void)
 
     /* The current ddraw is released when SetDirectDraw() is called. */
     hr = IDirectDrawMediaStream_SetDirectDraw(ddraw_stream, NULL);
-    todo_wine ok(hr == S_OK, "Got hr %#x.\n", hr);
+    ok(hr == S_OK, "Got hr %#x.\n", hr);
-    todo_wine EXPECT_REF(ddraw, 2);
+    EXPECT_REF(ddraw, 2);
 
     hr = IDirectDrawMediaStream_GetDirectDraw(ddraw_stream, &ddraw3);
     ok(hr == S_OK, "Got hr %#x.\n", hr);
@@ -4807,11 +4807,9 @@ static void test_ddrawstream_getsetdirectdraw(void)
     if (ddraw3) IDirectDraw_Release(ddraw3);
 
     hr = IDirectDrawMediaStream_SetDirectDraw(ddraw_stream, ddraw2);
-    todo_wine ok(hr == S_OK, "Got hr %#x.\n", hr);
+    ok(hr == S_OK, "Got hr %#x.\n", hr);
     EXPECT_REF(ddraw, 3);
 
-    if (hr == S_OK)
-    {
     hr = IDirectDrawMediaStream_GetDirectDraw(ddraw_stream, &ddraw3);
     ok(hr == S_OK, "Got hr %#x.\n", hr);
     ok(ddraw3 == ddraw2, "Expected ddraw %p, got %p.\n", ddraw2, ddraw3);
@@ -4865,7 +4863,6 @@ static void test_ddrawstream_getsetdirectdraw(void)
 
     ref = IDirectDraw_Release(ddraw3);
     ok(!ref, "Got outstanding refcount %d.\n", ref);
-    }
 
     EXPECT_REF(stream, 3);
     IDirectDrawMediaStream_Release(ddraw_stream);