dll/win32/mswsock/msafd/select.c

   1 /*
   2  * COPYRIGHT:   See COPYING in the top level directory
   3  * PROJECT:     ReactOS Winsock 2 SPI
   4  * FILE:        lib/mswsock/lib/init.c
   5  * PURPOSE:     DLL Initialization
   6  */
   7
   8 /* INCLUDES ******************************************************************/
   9 #include "msafd.h"
  10
  11 #define MSAFD_CHECK_EVENT(e, s) \
  12     (!(s->SharedData.AsyncDisabledEvents & e) && \
  13      (s->SharedData.AsyncEvents & e))
  14
  15 #define HANDLES_IN_SET(s) \
  16     s == NULL ? 0 : (s->fd_count & 0xFFFF)
  17
  18 /* DATA **********************************************************************/
  19
  20 HANDLE SockAsyncSelectHelperHandle;
  21 BOOLEAN SockAsyncSelectCalled;
  22
  23 /* FUNCTIONS *****************************************************************/
  24
  25 BOOLEAN
  26 WSPAPI
  27 SockCheckAndInitAsyncSelectHelper(VOID)
  28 {
  29     UNICODE_STRING AfdHelper;
  30     OBJECT_ATTRIBUTES ObjectAttributes;
  31     IO_STATUS_BLOCK IoStatusBlock;
  32     NTSTATUS Status;
  33     FILE_COMPLETION_INFORMATION CompletionInfo;
  34     OBJECT_HANDLE_ATTRIBUTE_INFORMATION HandleFlags;
  35
  36     /* First, make sure we're not already intialized */
  37     if (SockAsyncSelectHelperHandle) return TRUE;
  38
  39     /* Acquire the global lock */
  40     SockAcquireRwLockExclusive(&SocketGlobalLock);
  41
  42     /* Check again, under the lock */
  43     if (SockAsyncSelectHelperHandle)
  44     {
  45         /* Return without lock */
  46         SockReleaseRwLockExclusive(&SocketGlobalLock);
  47         return TRUE;
  48     }
  49
  50     /* Set up Handle Name and Object */
  51     RtlInitUnicodeString(&AfdHelper, L"\\Device\\Afd\\AsyncSelectHlp" );
  52     InitializeObjectAttributes(&ObjectAttributes,
  53                                &AfdHelper,
  54                                OBJ_INHERIT | OBJ_CASE_INSENSITIVE,
  55                                NULL,
  56                                NULL);
  57
  58     /* Open the Handle to AFD */
  59     Status = NtCreateFile(&SockAsyncSelectHelperHandle,
  60                           GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,
  61                           &ObjectAttributes,
  62                           &IoStatusBlock,
  63                           NULL,
  64                           0,
  65                           FILE_SHARE_READ | FILE_SHARE_WRITE,
  66                           FILE_OPEN_IF,
  67                           0,
  68                           NULL,
  69                           0);
  70     if (!NT_SUCCESS(Status))
  71     {
  72         /* Return without lock */
  73         SockReleaseRwLockExclusive(&SocketGlobalLock);
  74         return TRUE;
  75     }
  76
  77     /* Check if the port exists, and if not, create it */
  78     if (SockAsyncQueuePort) SockCreateAsyncQueuePort();
  79
  80     /*
  81      * Now Set up the Completion Port Information
  82      * This means that whenever a Poll is finished, the routine will be executed
  83      */
  84     CompletionInfo.Port = SockAsyncQueuePort;
  85     CompletionInfo.Key = SockAsyncSelectCompletion;
  86     Status = NtSetInformationFile(SockAsyncSelectHelperHandle,
  87                                   &IoStatusBlock,
  88                                   &CompletionInfo,
  89                                   sizeof(CompletionInfo),
  90                                   FileCompletionInformation);
  91
  92     /* Protect the Handle */
  93     HandleFlags.ProtectFromClose = TRUE;
  94     HandleFlags.Inherit = FALSE;
  95     Status = NtSetInformationObject(SockAsyncQueuePort,
  96                                     ObjectHandleFlagInformation,
  97                                     &HandleFlags,
  98                                     sizeof(HandleFlags));
  99
 100     /*
 101      * Set this variable to true so that Send/Recv/Accept will know whether
 102      * to renable disabled events
 103      */
 104     SockAsyncSelectCalled = TRUE;
 105
 106     /* Release lock and return */
 107     SockReleaseRwLockExclusive(&SocketGlobalLock);
 108     return TRUE;
 109 }
 110
 111 VOID
 112 WSPAPI
 113 SockAsyncSelectCompletion(PVOID Context,
 114                           PIO_STATUS_BLOCK IoStatusBlock)
 115 {
 116     PASYNC_DATA AsyncData = Context;
 117     PSOCKET_INFORMATION Socket = AsyncData->ParentSocket;
 118     ULONG Events;
 119     INT ErrorCode;
 120
 121     /* Acquire the socket lock */
 122     EnterCriticalSection(&Socket->Lock);
 123
 124     /* Check if the socket was closed or the I/O cancelled */
 125     if ((Socket->SharedData.State == SocketClosed) ||
 126         (IoStatusBlock->Status == STATUS_CANCELLED))
 127     {
 128         /* Return */
 129         LeaveCriticalSection(&Socket->Lock);
 130         goto error;
 131     }
 132
 133     /* Check if the Sequence  Number Changed behind our back */
 134     if (AsyncData->SequenceNumber != Socket->SharedData.SequenceNumber)
 135     {
 136         /* Return */
 137         LeaveCriticalSection(&Socket->Lock);
 138         goto error;
 139     }
 140
 141     /* Check we were manually called b/c of a failure */
 142     if (!NT_SUCCESS(IoStatusBlock->Status))
 143     {
 144         /* Get the error and tell WPU about it */
 145         ErrorCode = NtStatusToSocketError(IoStatusBlock->Status);
 146         SockUpcallTable->lpWPUPostMessage(Socket->SharedData.hWnd,
 147                                           Socket->SharedData.wMsg,
 148                                           Socket->Handle,
 149                                           WSAMAKESELECTREPLY(0, ErrorCode));
 150
 151         /* Return */
 152         LeaveCriticalSection(&Socket->Lock);
 153         goto error;
 154     }
 155
 156     /* Select the event bits */
 157     Events = AsyncData->AsyncSelectInfo.Handles[0].Events;
 158
 159     /* Check for receive event */
 160     if (MSAFD_CHECK_EVENT(FD_READ, Socket) && (Events & AFD_EVENT_RECEIVE))
 161     {
 162         /* Make the Notifcation */
 163         SockUpcallTable->lpWPUPostMessage(Socket->SharedData.hWnd,
 164                                           Socket->SharedData.wMsg,
 165                                           Socket->Handle,
 166                                           WSAMAKESELECTREPLY(FD_READ, 0));
 167
 168         /* Disable this event until the next read(); */
 169         Socket->SharedData.AsyncDisabledEvents |= FD_READ;
 170     }
 171
 172     /* Check for oob receive event */
 173     if (MSAFD_CHECK_EVENT(FD_OOB, Socket) && (Events & AFD_EVENT_OOB_RECEIVE))
 174     {
 175         /* Make the Notifcation */
 176         SockUpcallTable->lpWPUPostMessage(Socket->SharedData.hWnd,
 177                                           Socket->SharedData.wMsg,
 178                                           Socket->Handle,
 179                                           WSAMAKESELECTREPLY(FD_OOB, 0));
 180
 181         /* Disable this event until the next read(); */
 182         Socket->SharedData.AsyncDisabledEvents |= FD_OOB;
 183     }
 184
 185     /* Check for write event */
 186     if (MSAFD_CHECK_EVENT(FD_WRITE, Socket) && (Events & AFD_EVENT_SEND))
 187     {
 188         /* Make the Notifcation */
 189         SockUpcallTable->lpWPUPostMessage(Socket->SharedData.hWnd,
 190                                           Socket->SharedData.wMsg,
 191                                           Socket->Handle,
 192                                           WSAMAKESELECTREPLY(FD_WRITE, 0));
 193
 194         /* Disable this event until the next read(); */
 195         Socket->SharedData.AsyncDisabledEvents |= FD_WRITE;
 196     }
 197
 198     /* Check for accept event */
 199     if (MSAFD_CHECK_EVENT(FD_ACCEPT, Socket) && (Events & AFD_EVENT_ACCEPT))
 200     {
 201         /* Make the Notifcation */
 202         SockUpcallTable->lpWPUPostMessage(Socket->SharedData.hWnd,
 203                                           Socket->SharedData.wMsg,
 204                                           Socket->Handle,
 205                                           WSAMAKESELECTREPLY(FD_ACCEPT, 0));
 206
 207         /* Disable this event until the next read(); */
 208         Socket->SharedData.AsyncDisabledEvents |= FD_ACCEPT;
 209     }
 210
 211     /* Check for close events */
 212     if (MSAFD_CHECK_EVENT(FD_CLOSE, Socket) && ((Events & AFD_EVENT_ACCEPT) ||
 213                                                 (Events & AFD_EVENT_ABORT) ||
 214                                                 (Events & AFD_EVENT_CLOSE)))
 215     {
 216         /* Make the Notifcation */
 217         SockUpcallTable->lpWPUPostMessage(Socket->SharedData.hWnd,
 218                                           Socket->SharedData.wMsg,
 219                                           Socket->Handle,
 220                                           WSAMAKESELECTREPLY(FD_CLOSE, 0));
 221
 222         /* Disable this event until the next read(); */
 223         Socket->SharedData.AsyncDisabledEvents |= FD_CLOSE;
 224     }
 225
 226     /* Check for QOS event */
 227     if (MSAFD_CHECK_EVENT(FD_QOS, Socket) && (Events & AFD_EVENT_QOS))
 228     {
 229         /* Make the Notifcation */
 230         SockUpcallTable->lpWPUPostMessage(Socket->SharedData.hWnd,
 231                                           Socket->SharedData.wMsg,
 232                                           Socket->Handle,
 233                                           WSAMAKESELECTREPLY(FD_QOS, 0));
 234
 235         /* Disable this event until the next read(); */
 236         Socket->SharedData.AsyncDisabledEvents |= FD_QOS;
 237     }
 238
 239     /* Check for Group QOS event */
 240     if (MSAFD_CHECK_EVENT(FD_GROUP_QOS, Socket) && (Events & AFD_EVENT_GROUP_QOS))
 241     {
 242         /* Make the Notifcation */
 243         SockUpcallTable->lpWPUPostMessage(Socket->SharedData.hWnd,
 244                                           Socket->SharedData.wMsg,
 245                                           Socket->Handle,
 246                                           WSAMAKESELECTREPLY(FD_GROUP_QOS, 0));
 247
 248         /* Disable this event until the next read(); */
 249         Socket->SharedData.AsyncDisabledEvents |= FD_GROUP_QOS;
 250     }
 251
 252     /* Check for Routing Interface Change event */
 253     if (MSAFD_CHECK_EVENT(FD_ROUTING_INTERFACE_CHANGE, Socket) &&
 254         (Events & AFD_EVENT_ROUTING_INTERFACE_CHANGE))
 255     {
 256         /* Make the Notifcation */
 257         SockUpcallTable->lpWPUPostMessage(Socket->SharedData.hWnd,
 258                                           Socket->SharedData.wMsg,
 259                                           Socket->Handle,
 260                                           WSAMAKESELECTREPLY(FD_ROUTING_INTERFACE_CHANGE, 0));
 261
 262         /* Disable this event until the next read(); */
 263         Socket->SharedData.AsyncDisabledEvents |= FD_ROUTING_INTERFACE_CHANGE;
 264     }
 265
 266     /* Check for Address List Change event */
 267     if (MSAFD_CHECK_EVENT(FD_ADDRESS_LIST_CHANGE, Socket) &&
 268         (Events & AFD_EVENT_ADDRESS_LIST_CHANGE))
 269     {
 270         /* Make the Notifcation */
 271         SockUpcallTable->lpWPUPostMessage(Socket->SharedData.hWnd,
 272                                           Socket->SharedData.wMsg,
 273                                           Socket->Handle,
 274                                           WSAMAKESELECTREPLY(FD_ADDRESS_LIST_CHANGE, 0));
 275
 276         /* Disable this event until the next read(); */
 277         Socket->SharedData.AsyncDisabledEvents |= FD_ADDRESS_LIST_CHANGE;
 278     }
 279
 280     /* Check if there are any events left for us to check */
 281     if (!((Socket->SharedData.AsyncEvents) &
 282           (~Socket->SharedData.AsyncDisabledEvents)))
 283     {
 284         /* Nothing left, release lock and return */
 285         LeaveCriticalSection(&Socket->Lock);
 286         goto error;
 287     }
 288
 289     /* Keep Polling */
 290     SockProcessAsyncSelect(Socket, AsyncData);
 291
 292     /* Leave lock and return */
 293     LeaveCriticalSection(&Socket->Lock);
 294     return;
 295
 296 error:
 297     /* Dereference the socket and free the Async Data */
 298     SockDereferenceSocket(Socket);
 299     RtlFreeHeap(SockPrivateHeap, 0, AsyncData);
 300
 301     /* Dereference this thread and return */
 302     InterlockedDecrement(&SockAsyncThreadReferenceCount);
 303     return;
 304 }
 305
 306 VOID
 307 WSPAPI
 308 SockProcessAsyncSelect(PSOCKET_INFORMATION Socket,
 309                        PASYNC_DATA AsyncData)
 310 {
 311     ULONG lNetworkEvents;
 312     NTSTATUS Status;
 313
 314     /* Set up the Async Data Event Info */
 315     AsyncData->AsyncSelectInfo.Timeout.HighPart = 0x7FFFFFFF;
 316     AsyncData->AsyncSelectInfo.Timeout.LowPart = 0xFFFFFFFF;
 317     AsyncData->AsyncSelectInfo.HandleCount = 1;
 318     AsyncData->AsyncSelectInfo.Exclusive = TRUE;
 319     AsyncData->AsyncSelectInfo.Handles[0].Handle = (SOCKET)Socket->WshContext.Handle;
 320     AsyncData->AsyncSelectInfo.Handles[0].Events = 0;
 321
 322     /* Remove unwanted events */
 323     lNetworkEvents = Socket->SharedData.AsyncEvents &
 324                      (~Socket->SharedData.AsyncDisabledEvents);
 325
 326     /* Set Events to wait for */
 327     if (lNetworkEvents & FD_READ)
 328     {
 329         AsyncData->AsyncSelectInfo.Handles[0].Events |= AFD_EVENT_RECEIVE;
 330     }
 331     if (lNetworkEvents & FD_WRITE)
 332     {
 333         AsyncData->AsyncSelectInfo.Handles[0].Events |= AFD_EVENT_SEND;
 334     }
 335     if (lNetworkEvents & FD_OOB)
 336     {
 337         AsyncData->AsyncSelectInfo.Handles[0].Events |= AFD_EVENT_OOB_RECEIVE;
 338     }
 339     if (lNetworkEvents & FD_ACCEPT)
 340     {
 341         AsyncData->AsyncSelectInfo.Handles[0].Events |= AFD_EVENT_ACCEPT;
 342     }
 343     if (lNetworkEvents & FD_CLOSE)
 344     {
 345         AsyncData->AsyncSelectInfo.Handles[0].Events |= AFD_EVENT_DISCONNECT |
 346                                                         AFD_EVENT_ABORT |
 347                                                         AFD_EVENT_CLOSE;
 348     }
 349     if (lNetworkEvents & FD_QOS)
 350     {
 351         AsyncData->AsyncSelectInfo.Handles[0].Events |= AFD_EVENT_QOS;
 352     }
 353     if (lNetworkEvents & FD_GROUP_QOS)
 354     {
 355         AsyncData->AsyncSelectInfo.Handles[0].Events |= AFD_EVENT_GROUP_QOS;
 356     }
 357     if (lNetworkEvents & FD_ROUTING_INTERFACE_CHANGE)
 358     {
 359         AsyncData->AsyncSelectInfo.Handles[0].Events |= AFD_EVENT_ROUTING_INTERFACE_CHANGE;
 360     }
 361     if (lNetworkEvents & FD_ADDRESS_LIST_CHANGE)
 362     {
 363         AsyncData->AsyncSelectInfo.Handles[0].Events |= AFD_EVENT_ADDRESS_LIST_CHANGE;
 364     }
 365
 366     /* Send IOCTL */
 367     Status = NtDeviceIoControlFile(SockAsyncSelectHelperHandle,
 368                                    NULL,
 369                                    NULL,
 370                                    AsyncData,
 371                                    &AsyncData->IoStatusBlock,
 372                                    IOCTL_AFD_SELECT,
 373                                    &AsyncData->AsyncSelectInfo,
 374                                    sizeof(AsyncData->AsyncSelectInfo),
 375                                    &AsyncData->AsyncSelectInfo,
 376                                    sizeof(AsyncData->AsyncSelectInfo));
 377     /* Check for failure */
 378     if (NT_ERROR(Status))
 379     {
 380         /* I/O Manager Won't call the completion routine; do it manually */
 381         AsyncData->IoStatusBlock.Status = Status;
 382         SockAsyncSelectCompletion(AsyncData, &AsyncData->IoStatusBlock);
 383     }
 384 }
 385
 386 VOID
 387 WSPAPI
 388 SockProcessQueuedAsyncSelect(PVOID Context,
 389                              PIO_STATUS_BLOCK IoStatusBlock)
 390 {
 391     PASYNC_DATA AsyncData = Context;
 392     PSOCKET_INFORMATION Socket = AsyncData->ParentSocket;
 393
 394     /* Lock the socket */
 395     EnterCriticalSection(&Socket->Lock);
 396
 397     /* Make sure it's not closed */
 398     if (Socket->SharedData.State == SocketClosed)
 399     {
 400         /* Return */
 401         LeaveCriticalSection(&Socket->Lock);
 402         goto error;
 403     }
 404
 405     /* Check if the Sequence Number changed by now */
 406     if (AsyncData->SequenceNumber != Socket->SharedData.SequenceNumber)
 407     {
 408         /* Return */
 409         LeaveCriticalSection(&Socket->Lock);
 410         goto error;
 411     }
 412
 413     /* Check if select is needed */
 414     if (!((Socket->SharedData.AsyncEvents &
 415           ~Socket->SharedData.AsyncDisabledEvents)))
 416     {
 417         /* Return */
 418         LeaveCriticalSection(&Socket->Lock);
 419         goto error;
 420     }
 421
 422     /* Do the actual select */
 423     SockProcessAsyncSelect(Socket, AsyncData);
 424
 425     /* Return */
 426     LeaveCriticalSection(&Socket->Lock);
 427     return;
 428
 429 error:
 430     /* Dereference the socket and free the async data */
 431     SockDereferenceSocket(Socket);
 432     RtlFreeHeap(SockPrivateHeap, 0, AsyncData);
 433
 434     /* Dereference this thread */
 435     InterlockedDecrement(&SockAsyncThreadReferenceCount);
 436 }
 437
 438 INT
 439 WSPAPI
 440 SockReenableAsyncSelectEvent(IN PSOCKET_INFORMATION Socket,
 441                              IN ULONG Event)
 442 {
 443     PASYNC_DATA AsyncData;
 444     NTSTATUS Status;
 445
 446     /* Make sure the event is actually disabled */
 447     if (!(Socket->SharedData.AsyncDisabledEvents & Event)) return NO_ERROR;
 448
 449     /* Make sure we're not closed */
 450     if (Socket->SharedData.State == SocketClosed) return NO_ERROR;
 451
 452     /* Re-enable it */
 453     Socket->SharedData.AsyncDisabledEvents &= ~Event;
 454
 455     /* Return if no more events are being polled */
 456     if (!((Socket->SharedData.AsyncEvents &
 457           ~Socket->SharedData.AsyncDisabledEvents)))
 458     {
 459         return NO_ERROR;
 460     }
 461
 462     /* Allocate Async Data */
 463     AsyncData = SockAllocateHeapRoutine(SockPrivateHeap, 0, sizeof(ASYNC_DATA));
 464
 465     /* Increase the sequence number to stop anything else */
 466     Socket->SharedData.SequenceNumber++;
 467
 468     /* Set up the Async Data */
 469     AsyncData->ParentSocket = Socket;
 470     AsyncData->SequenceNumber = Socket->SharedData.SequenceNumber;
 471
 472     /* Begin Async Select by using I/O Completion */
 473     Status = NtSetIoCompletion(SockAsyncQueuePort,
 474                                (PVOID)&SockProcessQueuedAsyncSelect,
 475                                AsyncData,
 476                                0,
 477                                0);
 478     if (!NT_SUCCESS(Status))
 479     {
 480         /* Dereference the socket and fail */
 481         SockDereferenceSocket(Socket);
 482         RtlFreeHeap(SockPrivateHeap, 0, AsyncData);
 483         return NtStatusToSocketError(Status);
 484     }
 485
 486     /* All done */
 487     return NO_ERROR;
 488 }
 489
 490 INT
 491 WSPAPI
 492 SockAsyncSelectHelper(IN PSOCKET_INFORMATION Socket,
 493                       IN HWND hWnd,
 494                       IN UINT wMsg,
 495                       IN LONG lEvent)
 496 {
 497     PASYNC_DATA AsyncData = NULL;
 498     BOOLEAN BlockMode;
 499     NTSTATUS Status;
 500     INT ErrorCode;
 501
 502     /* Allocate the Async Data Structure to pass on to the Thread later */
 503     AsyncData = SockAllocateHeapRoutine(SockPrivateHeap, 0, sizeof(*AsyncData));
 504     if (!AsyncData) return WSAENOBUFS;
 505
 506     /* Acquire socket lock */
 507     EnterCriticalSection(&Socket->Lock);
 508
 509     /* Is there an active WSPEventSelect? */
 510     if (Socket->SharedData.AsyncEvents)
 511     {
 512         /* Call the helper to process it */
 513         ErrorCode = SockEventSelectHelper(Socket, NULL, 0);
 514         if (ErrorCode != NO_ERROR) goto error;
 515     }
 516
 517     /* Set Socket to Non-Blocking */
 518     BlockMode = TRUE;
 519     ErrorCode = SockSetInformation(Socket,
 520                                    AFD_INFO_BLOCKING_MODE,
 521                                    &BlockMode,
 522                                    NULL,
 523                                    NULL);
 524     if (ErrorCode != NO_ERROR) goto error;
 525
 526     /* AFD was notified, set it locally as well */
 527     Socket->SharedData.NonBlocking = TRUE;
 528
 529     /* Store Socket Data */
 530     Socket->SharedData.hWnd = hWnd;
 531     Socket->SharedData.wMsg = wMsg;
 532     Socket->SharedData.AsyncEvents = lEvent;
 533     Socket->SharedData.AsyncDisabledEvents = 0;
 534
 535     /* Check if the socket is not connected and not a datagram socket */
 536     if ((!SockIsSocketConnected(Socket)) && !MSAFD_IS_DGRAM_SOCK(Socket))
 537     {
 538         /* Disable FD_WRITE for now, so we don't get it before FD_CONNECT */
 539         Socket->SharedData.AsyncDisabledEvents |= FD_WRITE;
 540     }
 541
 542     /* Increase the sequence number */
 543     Socket->SharedData.SequenceNumber++;
 544
 545     /* Return if there are no more Events */
 546     if (!(Socket->SharedData.AsyncEvents &
 547         (~Socket->SharedData.AsyncDisabledEvents)))
 548     {
 549         /* Release the lock, dereference the async thread and the socket */
 550         LeaveCriticalSection(&Socket->Lock);
 551         InterlockedDecrement(&SockAsyncThreadReferenceCount);
 552         SockDereferenceSocket(Socket);
 553
 554         /* Free the Async Data */
 555         RtlFreeHeap(SockPrivateHeap, 0, AsyncData);
 556         return NO_ERROR;
 557     }
 558
 559     /* Set up the Async Data */
 560     AsyncData->ParentSocket = Socket;
 561     AsyncData->SequenceNumber = Socket->SharedData.SequenceNumber;
 562
 563     /* Release the lock now */
 564     LeaveCriticalSection(&Socket->Lock);
 565
 566     /* Begin Async Select by using I/O Completion */
 567     Status = NtSetIoCompletion(SockAsyncQueuePort,
 568                                (PVOID)&SockProcessQueuedAsyncSelect,
 569                                AsyncData,
 570                                0,
 571                                0);
 572     if (!NT_SUCCESS(Status))
 573     {
 574         /* Dereference the async thread and fail */
 575         InterlockedDecrement(&SockAsyncThreadReferenceCount);
 576         ErrorCode = NtStatusToSocketError(Status);
 577     }
 578
 579 error:
 580     /* Check for error */
 581     if (ErrorCode != NO_ERROR)
 582     {
 583         /* Free the async data */
 584         RtlFreeHeap(SockPrivateHeap, 0, AsyncData);
 585
 586         /* Fail */
 587         return SOCKET_ERROR;
 588     }
 589
 590     /* Increment the socket reference */
 591     InterlockedIncrement(&Socket->RefCount);
 592
 593     /* Return success */
 594     return NO_ERROR;
 595 }
 596
 597 INT
 598 WSPAPI
 599 WSPAsyncSelect(IN SOCKET Handle,
 600                IN HWND hWnd,
 601                IN UINT wMsg,
 602                IN LONG lEvent,
 603                OUT LPINT lpErrno)
 604 {
 605     PSOCKET_INFORMATION Socket;
 606     INT ErrorCode;
 607     PWINSOCK_TEB_DATA ThreadData;
 608
 609     /* Enter prolog */
 610     ErrorCode = SockEnterApiFast(&ThreadData);
 611     if (ErrorCode != NO_ERROR)
 612     {
 613         /* Fail */
 614         *lpErrno = ErrorCode;
 615         return SOCKET_ERROR;
 616     }
 617
 618     /* Check for valid events */
 619     if (lEvent & ~FD_ALL_EVENTS)
 620     {
 621         /* Fail */
 622         ErrorCode = WSAEINVAL;
 623         goto error;
 624     }
 625
 626     /* Check for valid window handle */
 627     if (!IsWindow(hWnd))
 628     {
 629         /* Fail */
 630         ErrorCode = WSAEINVAL;
 631         goto error;
 632     }
 633
 634     /* Create the Asynch Thread if Needed */
 635     if (!SockCheckAndReferenceAsyncThread())
 636     {
 637         /* Fail */
 638         ErrorCode = WSAENOBUFS;
 639         goto error;
 640     }
 641
 642     /* Open a Handle to AFD's Async Helper */
 643     if (!SockCheckAndInitAsyncSelectHelper())
 644     {
 645         /* Dereference async thread and fail */
 646         InterlockedDecrement(&SockAsyncThreadReferenceCount);
 647         ErrorCode = WSAENOBUFS;
 648         goto error;
 649     }
 650
 651     /* Get the socket structure */
 652     Socket = SockFindAndReferenceSocket(Handle, TRUE);
 653     if (!Socket)
 654     {
 655         /* Fail */
 656         ErrorCode = WSAENOTSOCK;
 657         goto error;
 658     }
 659
 660     /* Call the helper to do the work */
 661     ErrorCode = SockAsyncSelectHelper(Socket,
 662                                       hWnd,
 663                                       wMsg,
 664                                       lEvent);
 665
 666     /* Dereference the socket */
 667     SockDereferenceSocket(Socket);
 668
 669 error:
 670     /* Check for error */
 671     if (ErrorCode != NO_ERROR)
 672     {
 673         /* Fail */
 674         *lpErrno = ErrorCode;
 675         return SOCKET_ERROR;
 676     }
 677
 678     /* Return success */
 679     return NO_ERROR;
 680 }
 681
 682 INT
 683 WSPAPI
 684 WSPSelect(INT nfds,
 685           PFD_SET readfds,
 686           PFD_SET writefds,
 687           PFD_SET exceptfds,
 688           CONST LPTIMEVAL timeout,
 689           LPINT lpErrno)
 690 {
 691     IO_STATUS_BLOCK IoStatusBlock;
 692     PAFD_POLL_INFO PollInfo = NULL;
 693     NTSTATUS Status;
 694     CHAR PollBuffer[sizeof(AFD_POLL_INFO) + 3 * sizeof(AFD_HANDLE)];
 695     PAFD_HANDLE HandleArray;
 696     ULONG HandleCount, OutCount = 0;
 697     ULONG PollBufferSize;
 698     ULONG i;
 699     PWINSOCK_TEB_DATA ThreadData;
 700     LARGE_INTEGER uSec;
 701     ULONG BlockType;
 702     INT ErrorCode;
 703
 704     /* Enter prolog */
 705     ErrorCode = SockEnterApiFast(&ThreadData);
 706     if (ErrorCode != NO_ERROR)
 707     {
 708         /* Fail */
 709         *lpErrno = ErrorCode;
 710         return SOCKET_ERROR;
 711     }
 712
 713     /* How many sockets will we check? */
 714     HandleCount = HANDLES_IN_SET(readfds) +
 715                   HANDLES_IN_SET(writefds) +
 716                   HANDLES_IN_SET(exceptfds);
 717
 718     /* Leave if none are */
 719     if (!HandleCount) return NO_ERROR;
 720
 721     /* How much space will they require? */
 722     PollBufferSize = sizeof(*PollInfo) + (HandleCount * sizeof(AFD_HANDLE));
 723
 724     /* Check if our stack is big enough to hold it */
 725     if (PollBufferSize <= sizeof(PollBuffer))
 726     {
 727         /* Use the stack */
 728         PollInfo = (PVOID)PollBuffer;
 729     }
 730     else
 731     {
 732         /* Allocate from heap instead */
 733         PollInfo = SockAllocateHeapRoutine(SockPrivateHeap, 0, PollBufferSize);
 734         if (!PollInfo)
 735         {
 736             /* Fail */
 737             ErrorCode = WSAENOBUFS;
 738             goto error;
 739         }
 740     }
 741
 742     /* Number of handles for AFD to Check */
 743     PollInfo->HandleCount = HandleCount;
 744     PollInfo->Exclusive = FALSE;
 745     HandleArray = PollInfo->Handles;
 746
 747     /* Select the Read Events */
 748     for (i = 0; readfds && i < (readfds->fd_count & 0xFFFF); i++)
 749     {
 750         /* Fill out handle info */
 751         HandleArray->Handle = (SOCKET)readfds->fd_array[i];
 752         HandleArray->Events = AFD_EVENT_RECEIVE |
 753                               AFD_EVENT_DISCONNECT |
 754                               AFD_EVENT_ABORT;
 755
 756         /* Move to the next one */
 757         HandleArray++;
 758     }
 759     for (i = 0; writefds && i < (writefds->fd_count & 0xFFFF); i++)
 760     {
 761         /* Fill out handle info */
 762         HandleArray->Handle = (SOCKET)writefds->fd_array[i];
 763         HandleArray->Events = AFD_EVENT_SEND;
 764
 765         /* Move to the next one */
 766         HandleArray++;
 767     }
 768     for (i = 0; exceptfds && i < (exceptfds->fd_count & 0xFFFF); i++)
 769     {
 770         /* Fill out handle info */
 771         HandleArray->Handle = (SOCKET)exceptfds->fd_array[i];
 772         HandleArray->Events = AFD_EVENT_OOB_RECEIVE | AFD_EVENT_CONNECT_FAIL;
 773
 774         /* Move to the next one */
 775         HandleArray++;
 776     }
 777
 778     /* Check if a timeout was given */
 779     if (timeout)
 780     {
 781         /* Inifinte Timeout */
 782         PollInfo->Timeout.u.LowPart = -1;
 783         PollInfo->Timeout.u.HighPart = 0x7FFFFFFF;
 784     }
 785     else
 786     {
 787         /* Calculate microseconds */
 788         uSec = RtlEnlargedIntegerMultiply(timeout->tv_usec, -10);
 789
 790         /* Calculate seconds */
 791         PollInfo->Timeout = RtlEnlargedIntegerMultiply(timeout->tv_sec,
 792                                                        -1 * 1000 * 1000 * 10);
 793
 794         /* Add microseconds */
 795         PollInfo->Timeout.QuadPart += uSec.QuadPart;
 796     }
 797
 798     /* Send IOCTL */
 799     Status = NtDeviceIoControlFile((HANDLE)PollInfo->Handles[0].Handle,
 800                                    ThreadData->EventHandle,
 801                                    NULL,
 802                                    NULL,
 803                                    &IoStatusBlock,
 804                                    IOCTL_AFD_SELECT,
 805                                    PollInfo,
 806                                    PollBufferSize,
 807                                    PollInfo,
 808                                    PollBufferSize);
 809
 810     /* Check if we have to wait */
 811     if (Status == STATUS_PENDING)
 812     {
 813         /* Check if we'll call the blocking hook */
 814         if (!PollInfo->Timeout.QuadPart) BlockType = NO_BLOCKING_HOOK;
 815
 816         /* Wait for completion */
 817         SockWaitForSingleObject(ThreadData->EventHandle,
 818                                 (SOCKET)PollInfo->Handles[0].Handle,
 819                                 BlockType,
 820                                 NO_TIMEOUT);
 821
 822         /* Get new status */
 823         Status = IoStatusBlock.Status;
 824     }
 825
 826     /* Check for failure */
 827     if (!NT_SUCCESS(Status))
 828     {
 829         /* Fail */
 830         ErrorCode = NtStatusToSocketError(Status);
 831         goto error;
 832     }
 833
 834     /* Clear the Structures */
 835     if(readfds) FD_ZERO(readfds);
 836     if(writefds) FD_ZERO(writefds);
 837     if(exceptfds) FD_ZERO(exceptfds);
 838
 839     /* Get the handle info again */
 840     HandleCount = PollInfo->HandleCount;
 841     HandleArray = PollInfo->Handles;
 842
 843     /* Loop the Handles that got an event */
 844     for (i = 0; i < HandleCount; i++)
 845     {
 846         /* Check for a match */
 847         if (HandleArray->Events & AFD_EVENT_RECEIVE)
 848         {
 849             /* Check if it's not already set */
 850             if (!FD_ISSET((SOCKET)HandleArray->Handle, readfds))
 851             {
 852                 /* Increase Handles with an Event */
 853                 OutCount++;
 854
 855                 /* Set this handle */
 856                 FD_SET((SOCKET)HandleArray->Handle, readfds);
 857             }
 858         }
 859         /* Check for a match */
 860         if (HandleArray->Events & AFD_EVENT_SEND)
 861         {
 862             /* Check if it's not already set */
 863             if (!FD_ISSET((SOCKET)HandleArray->Handle, writefds))
 864             {
 865                 /* Increase Handles with an Event */
 866                 OutCount++;
 867
 868                 /* Set this handle */
 869                 FD_SET((SOCKET)HandleArray->Handle, writefds);
 870             }
 871         }
 872         /* Check for a match */
 873         if (HandleArray->Events & AFD_EVENT_OOB_RECEIVE)
 874         {
 875             /* Check if it's not already set */
 876             if (!FD_ISSET((SOCKET)HandleArray->Handle, exceptfds))
 877             {
 878                 /* Increase Handles with an Event */
 879                 OutCount++;
 880
 881                 /* Set this handle */
 882                 FD_SET((SOCKET)HandleArray->Handle, exceptfds);
 883             }
 884         }
 885         /* Check for a match */
 886         if (HandleArray->Events & AFD_EVENT_ACCEPT)
 887         {
 888             /* Check if it's not already set */
 889             if (!FD_ISSET((SOCKET)HandleArray->Handle, readfds))
 890             {
 891                 /* Increase Handles with an Event */
 892                 OutCount++;
 893
 894                 /* Set this handle */
 895                 FD_SET((SOCKET)HandleArray->Handle, readfds);
 896             }
 897         }
 898         /* Check for a match */
 899         if (HandleArray->Events & AFD_EVENT_CONNECT)
 900         {
 901             /* Check if it's not already set */
 902             if (!FD_ISSET((SOCKET)HandleArray->Handle, writefds))
 903             {
 904                 /* Increase Handles with an Event */
 905                 OutCount++;
 906
 907                 /* Set this handle */
 908                 FD_SET((SOCKET)HandleArray->Handle, writefds);
 909             }
 910         }
 911         /* Check for a match */
 912         if (HandleArray->Events & AFD_EVENT_CONNECT_FAIL)
 913         {
 914             /* Check if it's not already set */
 915             if (!FD_ISSET((SOCKET)HandleArray->Handle, exceptfds))
 916             {
 917                 /* Increase Handles with an Event */
 918                 OutCount++;
 919
 920                 /* Set this handle */
 921                 FD_SET((SOCKET)HandleArray->Handle, exceptfds);
 922             }
 923         }
 924         /* Check for a match */
 925         if (HandleArray->Events & AFD_EVENT_DISCONNECT)
 926         {
 927             /* Check if it's not already set */
 928             if (!FD_ISSET((SOCKET)HandleArray->Handle, readfds))
 929             {
 930                 /* Increase Handles with an Event */
 931                 OutCount++;
 932
 933                 /* Set this handle */
 934                 FD_SET((SOCKET)HandleArray->Handle, readfds);
 935             }
 936         }
 937         /* Check for a match */
 938         if (HandleArray->Events & AFD_EVENT_ABORT)
 939         {
 940             /* Check if it's not already set */
 941             if (!FD_ISSET((SOCKET)HandleArray->Handle, readfds))
 942             {
 943                 /* Increase Handles with an Event */
 944                 OutCount++;
 945
 946                 /* Set this handle */
 947                 FD_SET((SOCKET)HandleArray->Handle, readfds);
 948             }
 949         }
 950         /* Check for a match */
 951         if (HandleArray->Events & AFD_EVENT_CLOSE)
 952         {
 953             /* Check if it's not already set */
 954             if (!FD_ISSET((SOCKET)HandleArray->Handle, readfds))
 955             {
 956                 /* Increase Handles with an Event */
 957                 OutCount++;
 958
 959                 /* Set this handle */
 960                 FD_SET((SOCKET)HandleArray->Handle, readfds);
 961             }
 962         }
 963
 964         /* Move to next entry */
 965         HandleArray++;
 966     }
 967
 968 error:
 969
 970     /* Check if we should free the buffer */
 971     if (PollInfo && (PollInfo != (PVOID)PollBuffer))
 972     {
 973         /* Free it from the heap */
 974         RtlFreeHeap(SockPrivateHeap, 0, PollInfo);
 975     }
 976
 977     /* Check for error */
 978     if (ErrorCode != NO_ERROR)
 979     {
 980         /* Return error */
 981         *lpErrno = ErrorCode;
 982         return SOCKET_ERROR;
 983     }
 984
 985     /* Return the number of handles */
 986     return OutCount;
 987 }
 988