/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim: set ts=8 sts=2 et sw=2 tw=80: */
/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "mozilla/mscom/FastMarshaler.h"

#include "mozilla/mscom/Utils.h"

#include <objbase.h>

namespace mozilla {
namespace mscom {

HRESULT
FastMarshaler::Create(IUnknown* aOuter, IUnknown** aOutMarshalerUnk) {
  MOZ_ASSERT(XRE_IsContentProcess());

  if (!aOuter || !aOutMarshalerUnk) {
    return E_INVALIDARG;
  }

  *aOutMarshalerUnk = nullptr;

  HRESULT hr;
  RefPtr<FastMarshaler> fm(new FastMarshaler(aOuter, &hr));
  if (FAILED(hr)) {
    return hr;
  }

  return fm->InternalQueryInterface(IID_IUnknown, (void**)aOutMarshalerUnk);
}

FastMarshaler::FastMarshaler(IUnknown* aOuter, HRESULT* aResult)
    : mRefCnt(0), mOuter(aOuter), mStdMarshalWeak(nullptr) {
  *aResult =
      ::CoGetStdMarshalEx(aOuter, SMEXF_SERVER, getter_AddRefs(mStdMarshalUnk));
  if (FAILED(*aResult)) {
    return;
  }

  *aResult =
      mStdMarshalUnk->QueryInterface(IID_IMarshal, (void**)&mStdMarshalWeak);
  if (FAILED(*aResult)) {
    return;
  }

  // mStdMarshalWeak is weak
  mStdMarshalWeak->Release();
}

HRESULT
FastMarshaler::InternalQueryInterface(REFIID riid, void** ppv) {
  if (!ppv) {
    return E_INVALIDARG;
  }

  if (riid == IID_IUnknown) {
    RefPtr<IUnknown> punk(static_cast<IUnknown*>(&mInternalUnknown));
    punk.forget(ppv);
    return S_OK;
  }

  if (riid == IID_IMarshal) {
    RefPtr<IMarshal> ptr(this);
    ptr.forget(ppv);
    return S_OK;
  }

  return mStdMarshalUnk->QueryInterface(riid, ppv);
}

ULONG
FastMarshaler::InternalAddRef() { return ++mRefCnt; }

ULONG
FastMarshaler::InternalRelease() {
  ULONG result = --mRefCnt;
  if (!result) {
    delete this;
  }

  return result;
}

DWORD
FastMarshaler::GetMarshalFlags(DWORD aDestContext, DWORD aMshlFlags) {
  // Only worry about local contexts.
  if (aDestContext != MSHCTX_LOCAL) {
    return aMshlFlags;
  }

  if (IsCallerExternalProcess()) {
    return aMshlFlags;
  }

  // The caller is our parent main thread. Disable ping functionality.
  return aMshlFlags | MSHLFLAGS_NOPING;
}

HRESULT
FastMarshaler::GetUnmarshalClass(REFIID riid, void* pv, DWORD dwDestContext,
                                 void* pvDestContext, DWORD mshlflags,
                                 CLSID* pCid) {
  if (!mStdMarshalWeak) {
    return E_POINTER;
  }

  return mStdMarshalWeak->GetUnmarshalClass(
      riid, pv, dwDestContext, pvDestContext,
      GetMarshalFlags(dwDestContext, mshlflags), pCid);
}

HRESULT
FastMarshaler::GetMarshalSizeMax(REFIID riid, void* pv, DWORD dwDestContext,
                                 void* pvDestContext, DWORD mshlflags,
                                 DWORD* pSize) {
  if (!mStdMarshalWeak) {
    return E_POINTER;
  }

  return mStdMarshalWeak->GetMarshalSizeMax(
      riid, pv, dwDestContext, pvDestContext,
      GetMarshalFlags(dwDestContext, mshlflags), pSize);
}

HRESULT
FastMarshaler::MarshalInterface(IStream* pStm, REFIID riid, void* pv,
                                DWORD dwDestContext, void* pvDestContext,
                                DWORD mshlflags) {
  if (!mStdMarshalWeak) {
    return E_POINTER;
  }

  return mStdMarshalWeak->MarshalInterface(
      pStm, riid, pv, dwDestContext, pvDestContext,
      GetMarshalFlags(dwDestContext, mshlflags));
}

HRESULT
FastMarshaler::UnmarshalInterface(IStream* pStm, REFIID riid, void** ppv) {
  if (!mStdMarshalWeak) {
    return E_POINTER;
  }

  return mStdMarshalWeak->UnmarshalInterface(pStm, riid, ppv);
}

HRESULT
FastMarshaler::ReleaseMarshalData(IStream* pStm) {
  if (!mStdMarshalWeak) {
    return E_POINTER;
  }

  return mStdMarshalWeak->ReleaseMarshalData(pStm);
}

HRESULT
FastMarshaler::DisconnectObject(DWORD dwReserved) {
  if (!mStdMarshalWeak) {
    return E_POINTER;
  }

  return mStdMarshalWeak->DisconnectObject(dwReserved);
}

}  // namespace mscom
}  // namespace mozilla