建造LPWSTR的更安全方法

Question

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

当前，我使用以下方法构造一个指向缓冲区buf的初始指针：

auto buf = GetLogicalDriveStrings(0, nullptr);

然后，我创建LPWSTR以代替实际调用中的空指针，使用：

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

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

在得到结果后，我使用分配后创建的指针手动free() LPWSTR。

我如何为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。

这样做很好，但是有没有更好/更安全的方法来做到这一点呢？我愿意接受任何改进。

Answer 1

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

为此，您可以使用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;
}

Answer 2

我想我会这样做：

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

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

这将创建一个包含空的wstring，然后GetLogicalDriveStrings用它生成的内容覆盖内容。最后，我们将字符串调整为GetLogicalDriveStrings实际编写的字符数。

在那里，我们有一个完全正常的字符串，当它超出作用域时，它将释放内存，就像任何其他字符串一样。