Paul Gofman : winhttp: Don't allow socket transfers after receiving close frame.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: 04fcadacca19c5efbc66e3e44368e576428d8897
URL:    https://source.winehq.org/git/wine.git/?a=commit;h=04fcadacca19c5efbc66e3e44...

Author: Paul Gofman <pgofman(a)codeweavers.com>
Date:   Fri Jan 28 02:06:30 2022 +0300

winhttp: Don't allow socket transfers after receiving close frame.

Signed-off-by: Paul Gofman <pgofman(a)codeweavers.com>
Signed-off-by: Hans Leidekker <hans(a)codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard(a)winehq.org>

---

 dlls/winhttp/request.c            | 14 ++++++++++++--
 dlls/winhttp/tests/notification.c |  3 +++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/dlls/winhttp/request.c b/dlls/winhttp/request.c
index 04df90f74d7..a61355da7b9 100644
--- a/dlls/winhttp/request.c
+++ b/dlls/winhttp/request.c
@@ -3287,6 +3287,16 @@ static BOOL receive_io_complete( struct socket *socket )
     return count >= 0;
 }
 
+static BOOL socket_can_send( struct socket *socket )
+{
+    return socket->state == SOCKET_STATE_OPEN && !socket->close_frame_received;
+}
+
+static BOOL socket_can_receive( struct socket *socket )
+{
+    return socket->state <= SOCKET_STATE_SHUTDOWN && !socket->close_frame_received;
+}
+
 static enum socket_opcode map_buffer_type( WINHTTP_WEB_SOCKET_BUFFER_TYPE type )
 {
     switch (type)
@@ -3364,7 +3374,7 @@ DWORD WINAPI WinHttpWebSocketSend( HINTERNET hsocket, WINHTTP_WEB_SOCKET_BUFFER_
         release_object( &socket->hdr );
         return ERROR_WINHTTP_INCORRECT_HANDLE_TYPE;
     }
-    if (socket->state != SOCKET_STATE_OPEN)
+    if (!socket_can_send( socket ))
     {
         release_object( &socket->hdr );
         return ERROR_INVALID_OPERATION;
@@ -3723,7 +3733,7 @@ DWORD WINAPI WinHttpWebSocketReceive( HINTERNET hsocket, void *buf, DWORD len, D
         release_object( &socket->hdr );
         return ERROR_WINHTTP_INCORRECT_HANDLE_TYPE;
     }
-    if (socket->state > SOCKET_STATE_SHUTDOWN)
+    if (!socket_can_receive( socket ))
     {
         release_object( &socket->hdr );
         return ERROR_INVALID_OPERATION;
diff --git a/dlls/winhttp/tests/notification.c b/dlls/winhttp/tests/notification.c
index 3dc94a909b4..fabc5a03a46 100644
--- a/dlls/winhttp/tests/notification.c
+++ b/dlls/winhttp/tests/notification.c
@@ -1132,6 +1132,9 @@ static void test_websocket(BOOL secure)
     ok( close_status == 1000, "got %u\n", close_status );
     ok( size <= sizeof(buffer), "got %u\n", size );
 
+    err = pWinHttpWebSocketReceive( socket, buffer, sizeof(buffer), &size, &type );
+    ok( err == ERROR_INVALID_OPERATION, "got %u\n", err );
+
     info.buflen = 0xdeadbeef;
     setup_test( &info, winhttp_websocket_close, __LINE__ );
     err = pWinHttpWebSocketClose( socket, 1000, (void *)"success", sizeof("success") );