text: // ------------------------------------------------------------
  // Copyright (c) Microsoft Corporation.  All rights reserved.
  // Licensed under the MIT License (MIT). See License.txt in the repo root for license information.
  // ------------------------------------------------------------
  
  #include "stdafx.h"
  #if NTDDI_VERSION >= NTDDI_WIN8 
  
  using namespace std;
  using namespace Common;
  using namespace HttpCommon;
  using namespace HttpServer;
  
  StringLiteral const TraceType("HttpServerWebSocket");
  
  HttpServerWebSocket::HttpServerWebSocket(
      __in IRequestMessageContext &clientRequest,
      __in KStringViewA const &supportedSubProtocols,
      __in_opt ULONG receiveBufferSize,
      __in_opt ULONG sendBufferSize,
      __in_opt ULONGLONG openHandshakeTimeoutMilliSeconds,
      __in_opt ULONGLONG gracefulCloseTimeoutMilliSeconds,
      __in_opt ULONGLONG pongKeepAlivePeriodMilliSeconds,
      __in_opt ULONGLONG pingIntervalMilliSeconds,
      __in_opt ULONGLONG pongTimeoutMilliSeconds)
      : WebSocketHandler(TraceType, clientRequest.GetUrl())
      , clientRequest_(clientRequest)
      , supportedSubProtocols_(supportedSubProtocols)
      , receiveBufferSize_(receiveBufferSize)
      , sendBufferSize_(sendBufferSize)
      , openHandshakeTimeoutMilliSeconds_(openHandshakeTimeoutMilliSeconds)
      , gracefulCloseTimeoutMilliSeconds_(gracefulCloseTimeoutMilliSeconds)
      , pongKeepAlivePeriodMilliSeconds_(pongKeepAlivePeriodMilliSeconds)
      , pingIntervalMilliSeconds_(pingIntervalMilliSeconds)
      , pongTimeoutMilliSeconds_(pongTimeoutMilliSeconds)
  {}
  
  class HttpServerWebSocket::OpenAsyncOperation : public AsyncOperation
  {
  public:
      OpenAsyncOperation(
          HttpServerWebSocket &owner,
          WebSocketCloseReceivedCallback const &closeReceivedCallback,
          Common::AsyncCallback const &callback,
          Common::AsyncOperationSPtr const &parent)
          : AsyncOperation(callback, parent)
          , owner_(owner)
      {
          owner.closeReceivedCallback_ = closeReceivedCallback;
      }
  
      static ErrorCode End(__in AsyncOperationSPtr const& operation)
      {
          auto thisPtr = AsyncOperation::End<OpenAsyncOperation>(operation);
          return thisPtr->Error;
      }
  
  protected:
      void OnStart(AsyncOperationSPtr const &thisSPtr);
  
  private:
      void OnOpenComplete(
          KAsyncContextBase* const ,
          KAsyncServiceBase& ,
          NTSTATUS openStatus)
      {
          if (!NT_SUCCESS(openStatus))
          {
              WriteWarning(
                  TraceType,
                  "WebSocket open failed with NTSTATUS - {0}, for request url {1}",
                  owner_.clientRequest_.GetUrl(),
                  openStatus);
          }
  
          TryComplete(thisSPtr_, ErrorCode::FromNtStatus(openStatus));
          thisSPtr_.reset();
      }
  
      HttpServerWebSocket &owner_;
      AsyncOperationSPtr thisSPtr_;
  };
  
  void HttpServerWebSocket::OpenAsyncOperation::OnStart(AsyncOperationSPtr const &thisSPtr)
  {
      auto status = KHttpServerWebSocket::Create(owner_.kWebSocket_, owner_.clientRequest_.GetThisAllocator());
      if (!NT_SUCCESS(status))
      {
          WriteWarning(
              TraceType,
              "KHttpServerWebSocket create failed with NTSTATUS - {0}, for request url {1}",
              owner_.clientRequest_.GetUrl(),
              status);
  
          TryComplete(thisSPtr, ErrorCode::FromNtStatus(status));
          return;
      }
  
      status = owner_.kWebSocket_->CreateReceiveFragmentOperation(owner_.receiveFragmentKtlOperation_);
      if (!NT_SUCCESS(status))
      {
          WriteWarning(
              TraceType,
              "create receive fragment operation failed with NTSTATUS - {0}, for request url {1}",
              owner_.clientRequest_.GetUrl(),
              status);
  
          TryComplete(thisSPtr, ErrorCode::FromNtStatus(status));
          return;
      }
  
      status = owner_.kWebSocket_->CreateSendFragmentOperation(owner_.sendFragmentKtlOperation_);
      if (!NT_SUCCESS(status))
      {
          WriteWarning(
              TraceType,
              "Create send fragment operation failed with NTSTATUS - {0}, for request url {1}",
              owner_.clientRequest_.GetUrl(),
              status);
  
          TryComplete(thisSPtr, ErrorCode::FromNtStatus(status));
          return;
      }
  
      auto error = owner_.SetTimingConstants();
      if (!error.IsSuccess())
      {
          TryComplete(thisSPtr, error);
          return;
      }
  
      KStringA::SPtr ktlSubProtocolAnsiString;
      status = KStringA::Create(ktlSubProtocolAnsiString, HttpUtil::GetKtlAllocator(), owner_.supportedSubProtocols_, TRUE);
      if (!NT_SUCCESS(status))
      {
          WriteWarning(
              TraceType,
              "Getting subprotocol string failed with NTSTATUS - {0}, for request url {1}",
              owner_.clientRequest_.GetUrl(),
              status);
  
          TryComplete(thisSPtr, ErrorCode::FromNtStatus(status));
          return;
      }
  
      thisSPtr_ = shared_from_this();
  
      status = owner_.kWebSocket_->StartOpenWebSocket(
          *(static_cast<RequestMessageContext&>(owner_.clientRequest_)).GetInnerRequest(),
          KWebSocket::WebSocketCloseReceivedCallback(&owner_, &WebSocketHandler::CloseReceived),
          nullptr,
          KAsyncServiceBase::OpenCompletionCallback(this, &HttpServerWebSocket::OpenAsyncOperation::OnOpenComplete),
          owner_.receiveBufferSize_,
          owner_.sendBufferSize_,
          ktlSubProtocolAnsiString);
  
      if (!NT_SUCCESS(status))
      {
          WriteWarning(
              TraceType,
              "Starting websocket open failed with NTSTATUS - {0}, for request url {1}",
              owner_.clientRequest_.GetUrl(),
              status);
  
          thisSPtr_.reset();
          TryComplete(thisSPtr, ErrorCode::FromNtStatus(status));
          return;
      }
  }
  
  class HttpServerWebSocket::CloseAsyncOperation : public AsyncOperation
  {
  public:
      CloseAsyncOperation(
          HttpServerWebSocket &owner,
          __in KWebSocket::CloseStatusCode &statusCode,
          __in KSharedBufferStringA::SPtr &statusReason,
          AsyncCallback const &callback,
          AsyncOperationSPtr const &parent)
          : AsyncOperation(callback, parent)
          , owner_(owner)
          , statusCode_(statusCode)
          , statusReason_(statusReason)
      {
      }
  
      static ErrorCode End(AsyncOperationSPtr const& operation)
      {
          auto thisPtr = AsyncOperation::End<CloseAsyncOperation>(operation);
          return thisPtr->Error;
      }
  
  protected:
      void OnStart(AsyncOperationSPtr const &thisSPtr);
  
  private:
      VOID
      OnCloseComplete(
          KAsyncContextBase* const,
          KAsyncServiceBase&,
          NTSTATUS closeStatus)
      {
          if (!NT_SUCCESS(closeStatus))
          {
              WriteWarning(
                  TraceType,
                  "Close complete for request url {0} with status {1}",
                  owner_.clientRequest_.GetUrl(),
                  closeStatus);
          }
  
          TryComplete(thisSPtr_, ErrorCode::FromNtStatus(closeStatus));
          thisSPtr_.reset();
      }
  
      HttpServerWebSocket &owner_;
      AsyncOperationSPtr thisSPtr_;
      KWebSocket::CloseStatusCode statusCode_;
      KSharedBufferStringA::SPtr statusReason_;
  };
  
  void HttpServerWebSocket::CloseAsyncOperation::OnStart(AsyncOperationSPtr const &thisSPtr)
  {
      thisSPtr_ = thisSPtr;
  
      if (owner_.kWebSocket_->GetConnectionStatus() == KWebSocket::ConnectionStatus::CloseInitiated ||
          owner_.kWebSocket_->GetConnectionStatus() == KWebSocket::ConnectionStatus::CloseCompleting ||
          owner_.kWebSocket_->GetConnectionStatus() == KWebSocket::ConnectionStatus::Closed)
      {
          TryComplete(thisSPtr, ErrorCodeValue::ObjectClosed);
          return;
      }
  
      owner_.kWebSocket_->StartCloseWebSocket(
          nullptr,
          KAsyncServiceBase::CloseCompletionCallback(this, &HttpServerWebSocket::CloseAsyncOperation::OnCloseComplete),
          statusCode_,
          statusReason_);
  
  }
  
  AsyncOperationSPtr HttpServerWebSocket::BeginOpen(
      WebSocketCloseReceivedCallback const &closeReceivedCallback,
      AsyncCallback const &callback,
      AsyncOperationSPtr const &parent)
  {
      return AsyncOperation::CreateAndStart<OpenAsyncOperation>(
          *this,
          closeReceivedCallback,
          callback,
          parent);
  }
  
  ErrorCode HttpServerWebSocket::EndOpen(
      __in AsyncOperationSPtr const &operation)
  {
      return OpenAsyncOperation::End(operation);
  }
  
  AsyncOperationSPtr HttpServerWebSocket::BeginClose(
      __in KWebSocket::CloseStatusCode &statusCode,
      __in KSharedBufferStringA::SPtr &statusReason,
      AsyncCallback const &callback,
      AsyncOperationSPtr const &parent)
  {
      return AsyncOperation::CreateAndStart<CloseAsyncOperation>(
          *this,
          statusCode,
          statusReason,
          callback,
          parent);
  }
  
  ErrorCode HttpServerWebSocket::EndClose(
      AsyncOperationSPtr const &operation)
  {
      return CloseAsyncOperation::End(operation);
  }
  
  ErrorCode HttpServerWebSocket::GetRemoteCloseStatus(
      __out KWebSocket::CloseStatusCode &statusCode,
      __out KSharedBufferStringA::SPtr &statusReason)
  {
      if (kWebSocket_->GetConnectionStatus() == KWebSocket::ConnectionStatus::CloseReceived ||
          kWebSocket_->GetConnectionStatus() == KWebSocket::ConnectionStatus::CloseCompleting ||
          kWebSocket_->GetConnectionStatus() == KWebSocket::ConnectionStatus::Closed)
      {
          statusCode = kWebSocket_->GetRemoteWebSocketCloseStatusCode();
          statusReason = kWebSocket_->GetRemoteCloseReason();
  
          return ErrorCodeValue::Success;
      }
  
      return ErrorCodeValue::NotReady;
  }
  
  ErrorCode HttpServerWebSocket::SetTimingConstants()
  {
      auto status = kWebSocket_->SetTimingConstant(
          KWebSocket::TimingConstant::OpenTimeoutMs,
          static_cast<KWebSocket::TimingConstantValue>(openHandshakeTimeoutMilliSeconds_));
  
      if (!NT_SUCCESS(status))
      {
          WriteWarning(
              TraceType,
              "Set OpenTimeout failed with NTSTATUS - {0}, for request url {1}",
              status,
              requestUri_);
  
          return ErrorCode::FromNtStatus(status);
      }
  
      status = kWebSocket_->SetTimingConstant(
          KWebSocket::TimingConstant::CloseTimeoutMs,
          static_cast<KWebSocket::TimingConstantValue>(gracefulCloseTimeoutMilliSeconds_));
  
      if (!NT_SUCCESS(status))
      {
          WriteWarning(
              TraceType,
              "Set CloseTimeout failed with NTSTATUS - {0}, for request url {1}",
              status,
              requestUri_);
  
          return ErrorCode::FromNtStatus(status);
      }
  
      status = kWebSocket_->SetTimingConstant(
          KWebSocket::TimingConstant::PingQuietChannelPeriodMs,
          static_cast<KWebSocket::TimingConstantValue>(pingIntervalMilliSeconds_));
  
      if (!NT_SUCCESS(status))
      {
          WriteWarning(
              TraceType,
              "Set PingQuietChannelPeriodMs failed with NTSTATUS - {0}, for request url {1}",
              status,
              requestUri_);
  
          return ErrorCode::FromNtStatus(status);
      }
  
      status = kWebSocket_->SetTimingConstant(
          KWebSocket::TimingConstant::PongTimeoutMs,
          static_cast<KWebSocket::TimingConstantValue>(pongTimeoutMilliSeconds_));
  
      if (!NT_SUCCESS(status))
      {
          WriteWarning(
              TraceType,
              "Set PongTimeoutMs failed with NTSTATUS - {0}, for request url {1}",
              status,
              requestUri_);
  
          return ErrorCode::FromNtStatus(status);
      }
  
      status = kWebSocket_->SetTimingConstant(
          KWebSocket::TimingConstant::PongKeepalivePeriodMs,
          static_cast<KWebSocket::TimingConstantValue>(pongKeepAlivePeriodMilliSeconds_));
  
      if (!NT_SUCCESS(status))
      {
          WriteWarning(
              TraceType,
              "Set PongKeepalivePeriodMs failed with NTSTATUS - {0}, for request url {1}",
              status,
              requestUri_);
  
          return ErrorCode::FromNtStatus(status);
      }
  
      return ErrorCodeValue::Success;
  }
  
  KWebSocket::ConnectionStatus HttpServerWebSocket::GetConnectionStatus()
  {
      return kWebSocket_->GetConnectionStatus();
  }
  
  #endif
  