rainmeter-studio/Plugins/PluginPerfMon/PerfSnap.cpp

//============================================================================
// File: PERFSNAP.CPP
// Author: Matt Pietrek
// From: Microsoft Systems Journal, "Under the Hood", March 1996
//============================================================================

#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif

#include <windows.h>
#include <winperf.h>
#include <stdlib.h>
#include <tchar.h>
#pragma hdrstop
#include "titledb.h"
#include "perfsnap.h"
#include "makeptr.h"

PBYTE CPerfSnapshot::c_pBuffer = NULL;
DWORD CPerfSnapshot::c_cbBufferSize = 0;

CPerfSnapshot::CPerfSnapshot(
    CPerfTitleDatabase * pCounterTitles )
{
    m_pPerfDataHeader = 0;
    m_pCounterTitles = pCounterTitles;
}

CPerfSnapshot::~CPerfSnapshot( void )
{
    DisposeSnapshot();
}

BOOL
CPerfSnapshot::TakeSnapshot( PCTSTR pszSnapshotItems )
{
    DisposeSnapshot();  // Clear out any current snapshot

    // Convert the input string (e.g., "Process") into the form required
    // by the HKEY_PERFORMANCE_DATA key (e.g., "232")

    TCHAR szConvertedItemNames[ 256 ];
    if ( !ConvertSnapshotItemName( pszSnapshotItems, szConvertedItemNames,
                                   sizeof(szConvertedItemNames) ) )
        return FALSE;

    DWORD cbAllocSize = 0;
    LONG retValue;
        
    while ( 1 ) // Loop until we get the data, or we fail unexpectedly
    {
        retValue = RegQueryValueEx( HKEY_PERFORMANCE_DATA,
                                    szConvertedItemNames, 0, 0,
                                    c_pBuffer, &c_cbBufferSize );

        if ( retValue == ERROR_SUCCESS )    // We apparently got the snapshot
        {
			m_pPerfDataHeader = (PPERF_DATA_BLOCK)c_pBuffer;

            // Verify that the signature is a unicode "PERF"
            if ( memcmp( m_pPerfDataHeader->Signature, L"PERF", 8 ) )
                break;

            return TRUE;
        }
        else if ( retValue != ERROR_MORE_DATA ) // Anything other failure
            break;                              // code means something
                                                // bad happened, so bail out.

        // If we get here, our buffer wasn't big enough.  Delete it and
        // try again with a bigger buffer.
        delete [] c_pBuffer;
        
        // The new buffer size will be 4096 bytes bigger than the larger
        // of: 1) The previous allocation size, or 2) The size that the
        // RegQueryValueEx call said was necessary.
        if ( c_cbBufferSize > cbAllocSize )
            cbAllocSize = c_cbBufferSize + 4096;
        else
            cbAllocSize += 4096;

        // Allocate a new, larger buffer in preparation to try again.
        c_pBuffer = new BYTE[ cbAllocSize ];
        if ( !c_pBuffer )
            break;

        c_cbBufferSize = cbAllocSize;
    }

    // If we get here, the RegQueryValueEx failed unexpectedly.
    return FALSE;
}

void
CPerfSnapshot::DisposeSnapshot( void )
{
    m_pPerfDataHeader = 0;
}

void
CPerfSnapshot::CleanUp( void )
{
    delete [] c_pBuffer;
    c_pBuffer = 0;
	c_cbBufferSize = 0;
}

DWORD
CPerfSnapshot::GetNumObjectTypes( void )
{
    return m_pPerfDataHeader ? m_pPerfDataHeader->NumObjectTypes: 0;
}

BOOL
CPerfSnapshot::GetSystemName( PTSTR pszSystemName, DWORD nSize )
{
    if ( !m_pPerfDataHeader )   // If no snapshot data, bail out.
        return FALSE;

    // Make sure the input buffer size is big enough
    if ( nSize < m_pPerfDataHeader->SystemNameLength )
        return FALSE;

    // Make a unicode string point to the system name string
    // that's stored in the PERF_DATA_BLOCK
    LPWSTR lpwszName = MakePtr( LPWSTR, m_pPerfDataHeader,
                                m_pPerfDataHeader->SystemNameOffset );

    #ifdef UNICODE  // Copy the PERF_DATA_BLOCK string to the input buffer
    lstrcpy( pszSystemName, lpwszName );
    #else
    wcstombs( pszSystemName, lpwszName, nSize );
    #endif
        
    return TRUE;
}

PVOID
CPerfSnapshot::GetPostHeaderPointer( void )
{
    // Returns a header to the first byte following the PERF_DATA_BLOCK
    // (including the variable length system name string at the end)
    return m_pPerfDataHeader ?
            MakePtr(PVOID, m_pPerfDataHeader,m_pPerfDataHeader->HeaderLength)
            : 0;
}

BOOL
CPerfSnapshot::ConvertSnapshotItemName( PCTSTR pszIn,
                                        PTSTR pszOut, DWORD nSize )
{
    if ( IsBadStringPtr( pszIn, 0xFFFFFFFF ) )
        return FALSE;


	DWORD objectID = m_pCounterTitles->GetIndexFromTitleString(pszIn);

    if ( objectID )
        pszOut += wsprintf( pszOut, TEXT("%u "), objectID );
    else
        pszOut += wsprintf( pszOut, TEXT("%s "), pszIn );

    return TRUE;
}