Multidevice ASIO output plugin for WinAMP

#include "ASIODevice.hpp"

#include "WTL.h"

#include <algorithm>
#include <functional>
#include <limits>


//************************************
// Method:    ASIODriverList
// FullName:  ASIODriverList::ASIODriverList
// Access:    public 
//************************************

ASIODriverList::ASIODriverList()
    :
    SubKeys( _T( "Software\\ASIO" ) )
{
    assert( ASIODriverData::maxDrvNameLen <= RegistryKey::maximumRegistryKeyNameLength );
}


//************************************
// Method:    getDriverData
// FullName:  ASIODriverList::getDriverData
// Access:    public 
//************************************

ASIODriverData ASIODriverList::getDriverData( TCHAR const * const driverName ) const
{
    assert( !empty() );
    const_iterator pDriverKey( find( driverName ) );
    if ( pDriverKey.currentIndex() >= size() )
        pDriverKey.reset();
    return ASIODriverData( *parentKey(), pDriverKey->c_str() );
}


//************************************
// Method:    allChannelsAreOfSameType
// FullName:  ASIODriver::allChannelsAreOfSameType
// Access:    public 
//************************************

bool ASIODriver::allChannelsAreOfSameType() const
{
    ASIOSampleType const firstChannelType( getSampleType() );

    ASIOChannelInfo	oneChannelInfo;

    long maxInputChannels, maxOutputChannels;
    VERIFY( getNumberOfChannels( maxInputChannels, maxOutputChannels ) == ASE_OK );
    for( ASIOBool inChannel = ASIOFalse; inChannel <= ASIOTrue; ++inChannel )
    {
        oneChannelInfo.isInput = inChannel;
        for( long i( 1 ); i < ( inChannel ? maxInputChannels : maxOutputChannels ); )
        {
            oneChannelInfo.channel = i++;
            VERIFY( getChannelInfo( oneChannelInfo ) == ASE_OK );
            if( oneChannelInfo.type != firstChannelType )
                return false;
        }
    }
    return true;
}


//************************************
// Method:    can
// FullName:  ASIODriver::can
// Access:    public 
//************************************

bool ASIODriver::can( long const kAsioCanXEnum ) const
{
    assert( kAsioCanXEnum >= kAsioCanInputMonitor && kAsioCanXEnum <= kAsioCanOutputMeter );
    // This is a logically const operation.
    return const_cast<ASIODriver &>( *this ).futureCall( kAsioCanXEnum, NULL ) == ASE_SUCCESS;
}


//************************************
// Method:    createBuffers
// FullName:  ASIODriver::createBuffers
// Access:    public 
//************************************

ASIOError ASIODriver::createBuffers( ASIOBufferInfo bufferInfos[], long const numChannels, long const bufferSizeInSamples, ASIOCallbacks & callbacks )
{
    assert( state() == Initialized );
    assert( sizeof( ASIODriver::ASIOCallbacks ) == sizeof( ::ASIOCallbacks ) );
    assert( &callbacks.asioMessage          == (void*)&reinterpret_cast<::ASIOCallbacks &>( callbacks ).asioMessage          );
    assert( &callbacks.bufferSwitch         == (void*)&reinterpret_cast<::ASIOCallbacks &>( callbacks ).bufferSwitch         );
    assert( &callbacks.bufferSwitchTimeInfo == (void*)&reinterpret_cast<::ASIOCallbacks &>( callbacks ).bufferSwitchTimeInfo );
    assert( &callbacks.sampleRateDidChange  == (void*)&reinterpret_cast<::ASIOCallbacks &>( callbacks ).sampleRateDidChange  );
    ASIOError const error( driver().createBuffers( bufferInfos, numChannels, bufferSizeInSamples, reinterpret_cast<::ASIOCallbacks *>( &callbacks ) ) );
    if( error == ASE_OK ) state_ = Prepared;
    return error;
}


//************************************
// Method:    disposeBuffers
// FullName:  ASIODriver::disposeBuffers
// Access:    public 
//************************************

ASIOError ASIODriver::disposeBuffers()
{
    // (todo) rethink this logic and 'error checking'.
    assert( state() > Initialized );
    if( state_ > Initialized ) state_ = Initialized;
    return driver().disposeBuffers();
}


//************************************
// Method:    initializeDriverCOMObject
// FullName:  ASIODriver::initializeDriverCOMObject
// Access:    public 
//************************************

ASIODriver::State ASIODriver::initializeDriverCOMObject( HWND const parentWindow )
{
    assert( state() == Loaded );
    ASIOBool const initResult( driver().init( static_cast<void *>( parentWindow ) ) );
    if ( initResult != ASIOFalse )
        state_ = Initialized;
    return state_;
}


//************************************
// Method:    loadDriverCOMObject
// FullName:  ASIODriver::loadDriverCOMObject
// Access:    public 
//************************************

ASIODriver::State ASIODriver::loadDriverCOMObject( ASIODriverData const & driverData )
{
    pDriver_ = driverData.createInstance();
    state_ = static_cast<State>( static_cast<bool>( pDriver_ ) );
    return state_;
}


//************************************
// Method:    fixBufferSize
// FullName:  ASIODriver::fixBufferSize
// Access:    public 
//************************************
// Returns the value closest to the one requested but supported by the driver.
//************************************

long ASIODriver::fixBufferSize( long const bufferSize ) const
{
    long minSize, maxSize, preferredSize, granularity;
    VERIFY( getBufferSize( minSize, maxSize, preferredSize, granularity ) == ASE_OK );

    assert(	(minSize != maxSize) ||
        (minSize == maxSize && maxSize == preferredSize && granularity == 0) );

    if( bufferSize == 0 || granularity == 0 )
        return preferredSize;

    if( bufferSize < minSize )
        return minSize;

    if( bufferSize > maxSize )
        return maxSize;

    if( granularity == -1 )
    {
        long sz( minSize );
        while( ( sz *= 2 ) < maxSize )
            if( sz > bufferSize )
                return sz;
    }
    else
    {
        assert( granularity > 0 );
        long sz( minSize );
        while( ( sz += granularity ) < maxSize )
            if( sz > bufferSize )
                return sz;
    }
    return maxSize;
}


//************************************
// Method:    getDriverName
// FullName:  ASIODriver::getDriverName
// Access:    public 
//************************************

char const * ASIODriver::getDriverName() const
{
    static char driverName[ 48 ];
    driver().getDriverName( driverName );
    return driverName;
}


//************************************
// Method:    getErrorMessage
// FullName:  ASIODriver::getErrorMessage
// Access:    public 
//************************************

char const * ASIODriver::getErrorMessage() const
{
    static char driverErrorMessage[ 128 ];
    driver().getErrorMessage( driverErrorMessage );
    return driverErrorMessage;
}


//************************************
// Method:    getInitializationErrorMessage
// FullName:  ASIODriver::getInitializationErrorMessage
// Access:    public 
//************************************

char const * ASIODriver::getInitializationErrorMessage()
{
    switch ( state() )
    {
        case Unloaded: return "ASIO driver COM object creation failed!";
        case Loaded  : return getErrorMessage();
        default:
            assert( ( state() == Initialized ) && "This call makes sense only on an unused/freshly constructed and 'possibly initialized' object" );
            return NULL;
    }
}


//************************************
// Method:    getSamplePosition
// FullName:  ASIODriver::getSamplePosition
// Access:    public 
//************************************

ASIOError ASIODriver::getSamplePosition( unsigned long long & samplesPlayed, unsigned long long & nsTimeStamp ) const
{
    // Yes, this produces the smallest code (with MSVC).
    ASIOSamples   & samples  ( reinterpret_cast<ASIOSamples   &>( samplesPlayed ) );
    ASIOTimeStamp & timeStamp( reinterpret_cast<ASIOTimeStamp &>( nsTimeStamp   ) );
    ASIOError const error( driver().getSamplePosition( &samples, &timeStamp ) );
    std::swap( samples  .hi, samples  .lo );
    std::swap( timeStamp.hi, timeStamp.lo );
    return error;
}

ASIOError ASIODriver::getSamplePosition( ASIOTime & asioTime ) const
{
    return driver().getSamplePosition( &asioTime.timeInfo.samplePosition, &asioTime.timeInfo.systemTime );
}


//************************************
// Method:    getSampleType
// FullName:  ASIODriver::getSampleType
// Access:    public 
//************************************

ASIOSampleType ASIODriver::getSampleType( ASIOBool const forInputChannel ) const
{
    ASIOChannelInfo	firstChannelInfo;
    firstChannelInfo.channel = 0;
    firstChannelInfo.isInput = forInputChannel;
    ATLVERIFY( getChannelInfo( firstChannelInfo ) == ASE_OK );
    return firstChannelInfo.type;
}


//************************************
// Method:    initialize
// FullName:  ASIODriver::initialize
// Access:    public 
//************************************

ASIODriver::State ASIODriver::initialize( ASIODriverData const & driverData, HWND const parentWindow )
{
    if ( loadDriverCOMObject( driverData ) == Loaded )
        initializeDriverCOMObject( parentWindow );
    return state();
}


//************************************
// Method:    setChannelGain
// FullName:  ASIODriver::setChannelGain
// Access:    public 
//************************************

bool ASIODriver::setChannelGain( long const channel, long const gain, bool const isInput )
{
    ASIOChannelControls parameters;
    parameters.channel = channel;
    parameters.isInput = isInput;
    parameters.gain    = gain;
    //parameters.meter   = 0;
    return futureCall( isInput ? kAsioSetInputGain : kAsioSetOutputGain, &parameters ) == ASE_SUCCESS;
}


//************************************
// Method:    setInputMonitor
// FullName:  ASIODriver::setInputMonitor
// Access:    public 
//************************************

bool ASIODriver::setInputMonitor( long const inChannel, long const outChannel, long const gain /*= zero_dB*/, bool turnOn /*= true*/, long const pan /*= centre */ )
{
    assert( 0x7fffffff == (std::numeric_limits<long>::max)() );
    assert( outChannel % 2 == 0 );
    ASIOInputMonitor parameters =
    {
        inChannel,      // parameters.input
        outChannel,     // parameters.output
        gain,           // parameters.gain
        turnOn,         // parameters.state
        pan             // parameters.pan
    };
    return futureCall( kAsioSetInputMonitor, &parameters ) == ASE_SUCCESS;
}


//************************************
// Method:    start
// FullName:  ASIODriver::start
// Access:    public
//************************************

bool ASIODriver::start()
{
    assert( state() == Prepared );
    ASIOError const error( driver().start() );
    if( error == ASE_OK ) state_ = Running;
    return error == ASE_OK;
}


//************************************
// Method:    stop
// FullName:  ASIODriver::stop
// Access:    public
//************************************

bool ASIODriver::stop()
{
    assert( state() == Running );
    ASIOError const error( driver().stop() );
    if( error == ASE_OK ) state_ = Prepared;
    return error == ASE_OK;
}


//************************************
// Method:    unloadDriverCOMObject
// FullName:  ASIODriver::unloadDriverCOMObject
// Access:    public
//************************************

void ASIODriver::unloadDriverCOMObject()
{
#if 0  // (todo) Further investigate if this is actually "wrong"/a problem.
    assert( ( state() <= Initialized ) && "Forcibly unloading an active driver." );
#endif // (todo)
    pDriver_.release();
    state_ = Unloaded;
}


//************************************
// Method:    ASIODriverData
// FullName:  ASIODriverData::ASIODriverData
// Access:    private
//************************************

ASIODriverData::ASIODriverData( HKEY const hASIORootKey, TCHAR const * const driverKeyName )
{
    readDriverData( *RegistryKey( driverKeyName, hASIORootKey, KEY_READ ) );
}

ASIODriverData::ASIODriverData( ASIODriverList const & driverList, ASIODriverList::size_type const driverIndex )
{
    readDriverData( *RegistryKey( driverList[ driverIndex ]->c_str(), *driverList.parentKey() ) );
}

ASIODriverData::ASIODriverData( ASIODriverList::const_iterator const & driverIterator )
{
    readDriverData( *RegistryKey( driverIterator->c_str(), driverIterator.parentKey() ) );
}


//************************************
// Method:    createInstance
// FullName:  ASIODriverData::createInstance
// Access:    public
//************************************

IASIOPtr ASIODriverData::createInstance() const
{
    return IASIOPtr::createInstance( classID(), classID() );
}


//************************************
// Method:    path
// FullName:  ASIODriverData::path
// Access:    public
//************************************
//  Returns a pointer to a statically allocated buffer
//************************************

WinString ASIODriverData::path() const
{
    static OLECHAR const inprocServer[] = L"InprocServer32";
    size_t  const driverKeyPathLen( clsidStrLen + 1 + sizeof( inprocServer ) / sizeof( OLECHAR ) );
    OLECHAR driverInprocServerKeyName[ driverKeyPathLen ];
    // MSDN claims that the number of bytes not characters is returned but it
    // seems that is not true.
    VERIFY( ::StringFromGUID2( classID(), driverInprocServerKeyName, driverKeyPathLen ) == clsidStrLen + 1 );
    assert( driverInprocServerKeyName[ clsidStrLen ] == L'\0' );
    driverInprocServerKeyName[ clsidStrLen ] = '\\';
    std::memcpy( driverInprocServerKeyName + clsidStrLen + 1, inprocServer, sizeof( inprocServer ) );
    assert( std::wcslen( driverInprocServerKeyName ) == driverKeyPathLen - 1 );

    static TCHAR driverDLLPath[ MAX_PATH ] = { 0 };
    DWORD dllPathLength( sizeof( driverDLLPath ) );

    RegistryKey driverInprocServerKey( driverInprocServerKeyName, *RegistryKey( "CLSID", HKEY_CLASSES_ROOT ) );
    if
    (
        driverInprocServerKey &&
        ( ::RegQueryValueEx( *driverInprocServerKey, NULL, 0, NULL, reinterpret_cast<LPBYTE>( driverDLLPath ), &dllPathLength ) == ERROR_SUCCESS )
    )  
        return WinString( driverDLLPath, dllPathLength );
    else
        return WinString( NULL, 0 );
}


//************************************
// Method:    path2
// FullName:  ASIODriverData::path2
// Access:    public
//************************************

tstring ASIODriverData::path2() const
{
    WinString const driverPath( path() );
    return tstring( driverPath.c_str(), driverPath.size() );
}


//************************************
// Method:    readDriverData
// FullName:  ASIODriverData::readDriverData
// Access:    public
//************************************

void ASIODriverData::readDriverData( HKEY const driverKey )
{
    // (todo) Add better error handling here.
    DWORD classIDBufferSize( sizeof( classIDStr_ ) );
    ATLVERIFY( ::RegQueryValueExW( driverKey, L"CLSID", 0, NULL, reinterpret_cast<LPBYTE>( classIDStr_ ), &classIDBufferSize ) == ERROR_SUCCESS );
    assert( classIDBufferSize == sizeof( classIDStr_ ) );
    ATLVERIFY( ::CLSIDFromString( classIDStr_, &classID_ ) == NOERROR );

    DWORD drvNameBufSize( maxDrvNameLen );
    ATLVERIFY( ::RegQueryValueEx( driverKey, _T( "Description" ), 0, NULL, reinterpret_cast<LPBYTE>( description_ ), &drvNameBufSize ) == ERROR_SUCCESS );
}


//************************************
// Method:    convertASIOSampleTypeToString
// FullName:  convertASIOSampleTypeToString
// Access:    public
//************************************

TCHAR const * convertASIOSampleTypeToString( ASIOSampleType const sampleType )
{
    assert( sampleType >= 0 && sampleType < ASIOSTLastEntry );

    switch( sampleType )
    {
        case ASIOSTInt16MSB:
        case ASIOSTInt32MSB16:
        case ASIOSTInt16LSB:
        case ASIOSTInt32LSB16:
            return _T( "16 bit" );

        case ASIOSTInt32MSB18:
        case ASIOSTInt32LSB18:
            return _T( "18 bit" );

        case ASIOSTInt32LSB20:
        case ASIOSTInt32MSB20:
            return _T( "20 bit" );

        case ASIOSTInt24MSB:
        case ASIOSTInt24LSB:
            return _T( "20 or 24 bit" );

        case ASIOSTInt32MSB24:
        case ASIOSTInt32LSB24:
            return _T( "24 bit" );

        case ASIOSTInt32MSB:
        case ASIOSTInt32LSB:
        case ASIOSTFloat32MSB:
        case ASIOSTFloat32LSB:
            return _T( "32 bit" );

        case ASIOSTFloat64MSB:
        case ASIOSTFloat64LSB:
            return _T( "64 bit" );

        case ASIOSTDSDInt8LSB1:
        case ASIOSTDSDInt8MSB1:
        case ASIOSTDSDInt8NER8:
            return _T( "DSD" );

        default:
            assert( !"Unknown sample type" );
            __assume( false );
    }
}


//************************************
// Method:    Create
// FullName:  ImportsHooker::Create
// Access:    public
//************************************

#pragma comment( lib, "dbghelp.lib" )

ImportsHooker ImportsHooker::Create( HMODULE const moduleHandle )
{
    ULONG ulSize;
    PIMAGE_IMPORT_DESCRIPTOR pImportDesc( static_cast<PIMAGE_IMPORT_DESCRIPTOR>( ::ImageDirectoryEntryToData( moduleHandle, TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &ulSize ) ) );
    assert( pImportDesc && "This module has no import section." );
    return ImportsHooker( pImportDesc, reinterpret_cast<BYTE *>( moduleHandle ) );
}

ImportsHooker ImportsHooker::Create( TCHAR const * const moduleName )
{
    HMODULE const moduleHandle( ::GetModuleHandle( moduleName ) );
    assert( moduleHandle && "Requested module is not loaded in the current process." );
    return Create( moduleHandle );
}


//************************************
// Method:    findImportDescriptorForModule
// FullName:  ImportsHooker::findImportDescriptorForModule
// Access:    public
//************************************

PIMAGE_IMPORT_DESCRIPTOR ImportsHooker::findImportDescriptorForModule( PCSTR const module ) const
{
    PIMAGE_IMPORT_DESCRIPTOR pImportDesc( pImportDesc_ );

    while ( pImportDesc->Name )
    {
        PCSTR const pszModName( reinterpret_cast<PCSTR>( pModuleOffset_ + pImportDesc->Name ) );
        if ( ::lstrcmpiA( pszModName, module ) == 0 )
            return pImportDesc;
        ++pImportDesc;
    }
    return NULL;
}


//************************************
// Method:    switchFunctionThunk
// FullName:  ImportsHooker::switchFunctionThunk
// Access:    public
//************************************

bool ImportsHooker::switchFunctionThunk( PCSTR const originalFunctionImplementingModuleName, LPCSTR const pOriginalFunctionName, void * const pNewFunction ) const
{
    assert( pOriginalFunctionName );

    HMODULE const hOriginalImplementationModule( ::GetModuleHandleA( originalFunctionImplementingModuleName ) );
    assert( hOriginalImplementationModule && "Requested module is not loaded in the current process." );

    PROC const pOriginalFuncion( ::GetProcAddress( hOriginalImplementationModule, pOriginalFunctionName ) );
    assert( pOriginalFuncion && "Requested function could not be found in the specified module." );

    return switchFunctionThunk( originalFunctionImplementingModuleName, pOriginalFuncion, pNewFunction );
}

bool ImportsHooker::switchFunctionThunk( PCSTR const originalFunctionImplementingModuleName, void * const pOriginalFunction, void * const pNewFunction ) const
{
    assert( pOriginalFunction );
    assert( pNewFunction      );

    PIMAGE_IMPORT_DESCRIPTOR const pImportDesc( findImportDescriptorForModule( originalFunctionImplementingModuleName ) );
    assert( pImportDesc && "This module doesn't import any functions from this callee." );
    if ( !pImportDesc )
        return false;

    // Get caller's import address table (IAT) for the callee's functions.
    // Replace current function address with new function address.
    for
    (
        PIMAGE_THUNK_DATA pThunk( reinterpret_cast<PIMAGE_THUNK_DATA>( pModuleOffset_ + pImportDesc->FirstThunk ) );
        pThunk->u1.Function;
        ++pThunk
    )
    {
        // Get the address of the function address.
        void * * const ppCurrentFunction = &reinterpret_cast<void * &>( pThunk->u1.Function ); // MSVC bug, doesn't support constructor-like declaration and definition in this case.
        // Is this the function we're looking for?
        if ( *ppCurrentFunction == pOriginalFunction )
        {
            // The addresses match; change the import section address.
            HANDLE const hProcess( ::GetCurrentProcess() );
            bool success( ::WriteProcessMemory( hProcess, ppCurrentFunction, &pNewFunction, sizeof( pNewFunction ), NULL ) != false );
            if ( !success && ( ::GetLastError() == ERROR_NOACCESS ) )
            {
                DWORD dwOrgProtect;
                VERIFY( ::VirtualProtect( ppCurrentFunction, sizeof( pNewFunction ), PAGE_READWRITE, &dwOrgProtect ) );
                success = ( ::WriteProcessMemory( hProcess, ppCurrentFunction, &pNewFunction, sizeof( pNewFunction ), NULL ) != false );
                VERIFY( ::VirtualProtect( ppCurrentFunction, sizeof( pNewFunction ), dwOrgProtect , &dwOrgProtect ) );
            }
            return success;
        }
    }
    return false;
}