c++ - 构建LPWSTR的更安全方法

我正在使用需要GetLogicalDriveStrings()的WinAPI LPWSTR 函数，并且想知道是否存在更安全的方法来确保没有内存泄漏。

目前，我使用以下方法构造指向缓冲区buf的初始指针:

auto buf = GetLogicalDriveStrings(0, nullptr);

然后，我使用以下命令创建LPWSTR来代替我的实际调用中的空指针:

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

接下来，我创建一个指向driveStrings的指针以在以后释放它。在检查driveStrings是否为空指针或缓冲区(buf)是否为NULL(以防无法分配内存)之后，我使用GetLogicalDriveStrings()调用driveStrings。

得到结果后，我使用分配后的指针手动对LPWSTR进行free()编码。

我该如何为LPWSTR使用智能指针，这样我就不必使用malloc()和free()，但是它仍然可以与 GetLogicalDriveStrings() 函数一起使用？

最小工作示例:

    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);

如果我使用std::wstring，则无法弄清楚如何遍历driveStrings缓冲区中的每个字符串。如果我使用std::vector<WCHAR>，则无法弄清楚如何将每个元素都转换为GetDriveType()的LPCWSTR。

这样可以正常工作，但是有没有更好/更安全的方法呢？我愿意接受所有改进。

最佳答案

您可以为此使用 std::unique_ptr 。它可以用来分配这样的字符数组:

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

下面是显示如何与GetLogicalDriveStrings()一起使用的示例。该示例还显示了如何正确调用GetLastError()。必须在设置最后一个错误值的函数之后立即调用它。两者之间的任何其他系统调用(可能隐藏在C或C++标准代码中)都可能使最后的错误值无效。为了更容易使用，我将其包装到ThrowLastError()函数中，但是规则仍然适用。

#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;
}

我个人会使用GetLogicalDrives()，尽管这样可以完全避免缓冲区管理的麻烦。此外，由于只需调用一次此函数，因此简化了错误处理。为了完整起见，我在下面提供了一个如何使用GetLogicalDrives()的示例。

#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;
}