Safer way of constructing LPWSTR

Question

I'm using the WinAPI GetLogicalDriveStrings() function that requires a LPWSTR and am wondering if there is a safer way to do this to ensure there is no memory leak.

Currently I construct an initial pointer to a buffer buf using:

auto buf = GetLogicalDriveStrings(0, nullptr);

Then I create the LPWSTR to be used in lieu of the null pointer in my actual call using:

auto driveStrings = static_cast<LPWSTR>(malloc((buf + 1) * sizeof(WCHAR)));

Next I create a pointer to driveStrings to free it later. After checking if driveStrings is a null pointer or if the buffer (buf) is NULL (in case memory couldn't be allocated), I call GetLogicalDriveStrings() using driveStrings.

After I get the result I manually free() the LPWSTR using the pointer I made after allocating it.

How can I use a smart pointer for LPWSTR instead so I don't have to use malloc() and free(), but so that it will still work with the GetLogicalDriveStrings() function?

Mininum working example:

    auto buf = GetLogicalDriveStrings(0, nullptr);

    auto driveStrings = static_cast<LPWSTR>(malloc((buf + 1) * sizeof(WCHAR)));
    auto pDriveStrings = driveStrings;

    if (driveStrings == nullptr || buf == NULL)
    {
        std::stringstream msg;
        msg << "Can't allocate memory for drive list: ";
        msg << GetLastError();
        throw std::runtime_error(msg.str());
    }

    // get drive strings
    if (GetLogicalDriveStrings(buf, driveStrings) == NULL)
    {
        std::stringstream msg;
        msg << "GetLogicalDriveStrings error: ";
        msg << GetLastError();
        throw std::runtime_error(msg.str());
    }

    // iterate over results
    while (*driveStrings)
    {
        // GetDriveType() requires a LPCWSTR
        if (GetDriveType(driveStrings) == DRIVE_FIXED || GetDriveType(driveStrings) == DRIVE_REMOVABLE)
        {
            std::wcout << driveStrings << std::endl;
        }
        driveStrings += lstrlen(driveStrings) + 1;
    }

    free(pDriveStrings);

If I use a std::wstring, I can't figure out how to iterate over each string in the driveStrings buffer. If I use a std::vector<WCHAR>, I can't figure out how to cast each element to an LPCWSTR for GetDriveType().

This works fine as is, but is there a better/safer way to do this? I'm open to any and all improvements.

Answer 1

How can I use a smart pointer for LPWSTR instead so I don't have to use malloc() and free(), but so that it will still work with the GetLogicalDriveStrings() function?

You may use std::unique_ptr for that. It can be used to allocate an array of characters like this:

std::unique_ptr<wchar_t[]> buffer( new wchar_t[ size ] );

An example showing how to use it with GetLogicalDriveStrings() follows. The example also shows how to call GetLastError() correctly. It must be called immediately after the function that sets the last error value. Any other system call in between (possibly hidden in C or C++ standard code) may invalidate the last error value. For easier use I have wrapped it into ThrowLastError() function but the rules still apply.

#include <Windows.h>
#include <iostream>
#include <string>
#include <set>
#include <memory>

void ThrowLastError( const char* msg ) {
    DWORD err = ::GetLastError();
    throw std::system_error( static_cast<int>( err ), std::system_category(), msg );
}

std::set< std::wstring > GetLogicalDriveSet() {
    // Call GetLogicalDriveStrings() to get required buffer size.
    DWORD bufSize = ::GetLogicalDriveStrings( 0, nullptr );
    if( bufSize == 0 )
        ThrowLastError( "Could not get logical drives" );

    // Allocate an array of wchar_t and manage it using unique_ptr.
    // Make sure to allocate space for last '\0'.
    std::unique_ptr<wchar_t[]> buffer( new wchar_t[ bufSize + 1 ] ); 

    // Call GetLogicalDriveStrings() 2nd time to actually receive the strings.
    DWORD len = ::GetLogicalDriveStrings( bufSize, buffer.get() );
    if( len == 0 )
        ThrowLastError( "Could not get logical drives" );

    // In a rare case the number of drives may have changed after 
    // the first call to GetLogicalDriveStrings().
    if( len > bufSize )
        throw std::runtime_error( "Could not get logical drives - buffer size mismatch" );

    std::set< std::wstring > result;

    // Split the string returned by GetLogicalDriveStrings() at '\0'.
    auto p = buffer.get();
    while( *p ) {
        std::wstring path( p );
        result.insert( path );
        p += path.size() + 1;
    }

    return result;
}

int main(int argc, char* argv[]) {
    std::set< std::wstring > drives;
    try {
        drives = GetLogicalDriveSet();
    }
    catch( std::exception& e ) {
        std::cout << "Error: " << e.what() << std::endl;
        return 1;
    }

    std::cout << "Fixed and removable drives:\n";
    for( const auto& drv : drives ) {
        DWORD driveType = ::GetDriveType( drv.c_str() );
        if( driveType == DRIVE_FIXED || driveType == DRIVE_REMOVABLE ){
            std::wcout << drv << std::endl;
        }
    }
    return 0;
}

Personally I would go with GetLogicalDrives() though which avoids the hassles of buffer management completely. In addition the error handling is simplified because you have to call this function only once. For completeness I'm providing an example how to use GetLogicalDrives() below.

#include <Windows.h>
#include <iostream>
#include <string>
#include <set>

void ThrowLastError( const char* msg ) {
    DWORD err = ::GetLastError();
    throw std::system_error( static_cast<int>( err ), std::system_category(), msg );
}

std::set< std::wstring > GetLogicalDriveSet() {
    std::set< std::wstring > result;

    DWORD mask = GetLogicalDrives();
    if( mask == 0 )
        ThrowLastError( "Could not get logical drives" );       

    for( wchar_t drive = 'A'; drive <= 'Z'; ++drive ) {
        if( mask & 1 ) {
            // Build a complete root path like "C:\\" that can be used
            // with GetDriveType().
            wchar_t path[]{ drive, ':', '\\', 0 };
            result.insert( path );
        }
        // Shift all bits to the right so next "mask & 1" will test for
        // next drive letter.
        mask >>= 1;
    }

    return result;
}

int main(int argc, char* argv[]){
    std::set< std::wstring > drives;
    try {
        drives = GetLogicalDriveSet();
    }
    catch( std::exception& e ){
        std::cout << "Error: " << e.what() << std::endl;
        return 1;
    }

    std::cout << "Fixed and removable drives:\n";
    for( const auto& drv : drives ) {
        DWORD driveType = ::GetDriveType( drv.c_str() );
        if( driveType == DRIVE_FIXED || driveType == DRIVE_REMOVABLE ){
            std::wcout << drv << std::endl;
        }
    }

    return 0;
}

Answer 2

I think I'd do something like this:

std::wstring s(buf+1, '\0');

auto len = GetLogicalDriveStrings(buf, &s[0]);
s.resize(len);

This creates a wstring containing NULs, then GetLogicalDriveStrings overwrites the content with what it produces. Finally, we resize the string down to the number of characters that GetLogicalDriveStrings actually wrote.

From there, we have a perfectly normal string that will free its memory when it goes out of scope, just like any other string.