diff --git a/TCDirCore/AliasManager.cpp b/TCDirCore/AliasManager.cpp index 466b132..f991bf8 100644 --- a/TCDirCore/AliasManager.cpp +++ b/TCDirCore/AliasManager.cpp @@ -367,6 +367,8 @@ HRESULT CAliasManager::WriteAliasBlockToFile ( vector rgFileLines; bool fHasBom = false; SAliasBlock targetBlock; + uintmax_t cbOnDisk = 0; + error_code ec; @@ -375,6 +377,14 @@ HRESULT CAliasManager::WriteAliasBlockToFile ( hr = fileMgr.ReadProfileFile (strTargetPath, rgFileLines, fHasBom); CHR (hr); + // + // Defense in depth: a non-empty profile that parsed to nothing means + // the read lost content, and writing that back would destroy the file. + // + + cbOnDisk = filesystem::file_size (strTargetPath, ec); + CBREx (ec || !rgFileLines.empty() || cbOnDisk == 0, HRESULT_FROM_WIN32 (ERROR_READ_FAULT)); + hr = fileMgr.CreateBackup (strTargetPath); CHR (hr); } diff --git a/TCDirCore/ProfileFileAccess.h b/TCDirCore/ProfileFileAccess.h new file mode 100644 index 0000000..9672b9d --- /dev/null +++ b/TCDirCore/ProfileFileAccess.h @@ -0,0 +1,183 @@ +#pragma once + + + + + +// +// Suffix for the scratch file used to make profile writes atomic +// + +static constexpr LPCWSTR k_pszProfileTempSuffix = L".tcdir-tmp"; + + + + + +//////////////////////////////////////////////////////////////////////////////// +// +// IProfileFileAccess +// +// Injectable abstraction over the raw byte-level file I/O used when reading +// and writing PowerShell profile files. Production code uses +// CProfileFileAccessReal; tests substitute CProfileFileAccessMock so that I/O +// failure modes can be exercised without touching the real file system. +// +//////////////////////////////////////////////////////////////////////////////// + +class IProfileFileAccess +{ +public: + virtual ~IProfileFileAccess() = default; + + virtual HRESULT ReadAllBytes (const wstring & strPath, string & strBytes) = 0; + virtual HRESULT WriteAllBytes (const wstring & strPath, const string & strBytes) = 0; +}; + + + + + +//////////////////////////////////////////////////////////////////////////////// +// +// CProfileFileAccessReal +// +// Production implementation — delegates to the C runtime file APIs. +// +//////////////////////////////////////////////////////////////////////////////// + +class CProfileFileAccessReal : public IProfileFileAccess +{ +public: + HRESULT ReadAllBytes (const wstring & strPath, string & strBytes) override; + HRESULT WriteAllBytes (const wstring & strPath, const string & strBytes) override; +}; + + + + + +//////////////////////////////////////////////////////////////////////////////// +// +// CProfileFileAccessReal::ReadAllBytes +// +// Reads a file in full. Succeeds ONLY when every byte was read. +// +// Every seek/tell/read is checked. An unchecked one here is how a profile +// can come back silently empty and then be written straight back over the +// user's real file, destroying it. +// +//////////////////////////////////////////////////////////////////////////////// + +inline HRESULT CProfileFileAccessReal::ReadAllBytes (const wstring & strPath, string & strBytes) +{ + HRESULT hr = S_OK; + FILE * pf = nullptr; + long cbFile = 0; + size_t cbRead = 0; + int iSeek = 0; + int iError = 0; + + + + strBytes.clear(); + + _wfopen_s (&pf, strPath.c_str(), L"rb"); + CBREx (pf != nullptr, HRESULT_FROM_WIN32 (ERROR_FILE_NOT_FOUND)); + + iSeek = fseek (pf, 0, SEEK_END); + CBREx (iSeek == 0, HRESULT_FROM_WIN32 (ERROR_READ_FAULT)); + + cbFile = ftell (pf); + CBREx (cbFile >= 0, HRESULT_FROM_WIN32 (ERROR_READ_FAULT)); + + iSeek = fseek (pf, 0, SEEK_SET); + CBREx (iSeek == 0, HRESULT_FROM_WIN32 (ERROR_READ_FAULT)); + + if (cbFile > 0) + { + strBytes.resize (static_cast(cbFile)); + + cbRead = fread (strBytes.data(), 1, strBytes.size(), pf); + + // + // A short read is a failure, never a silently truncated buffer. + // + + CBREx (cbRead == strBytes.size(), HRESULT_FROM_WIN32 (ERROR_READ_FAULT)); + + iError = ferror (pf); + CBREx (iError == 0, HRESULT_FROM_WIN32 (ERROR_READ_FAULT)); + } + +Error: + if (pf != nullptr) + { + fclose (pf); + } + + if (FAILED (hr)) + { + strBytes.clear(); + } + + return hr; +} + + + + + +//////////////////////////////////////////////////////////////////////////////// +// +// CProfileFileAccessReal::WriteAllBytes +// +// Writes to a scratch file and renames it over the target, so an interrupted +// or failed write cannot leave a partially written profile behind. +// +//////////////////////////////////////////////////////////////////////////////// + +inline HRESULT CProfileFileAccessReal::WriteAllBytes (const wstring & strPath, const string & strBytes) +{ + HRESULT hr = S_OK; + FILE * pf = nullptr; + wstring strTemp; + size_t cbWritten = 0; + int iFlush = 0; + error_code ec; + + + + strTemp = strPath + k_pszProfileTempSuffix; + + _wfopen_s (&pf, strTemp.c_str(), L"wb"); + CBREx (pf != nullptr, HRESULT_FROM_WIN32 (ERROR_ACCESS_DENIED)); + + if (!strBytes.empty()) + { + cbWritten = fwrite (strBytes.data(), 1, strBytes.size(), pf); + CBREx (cbWritten == strBytes.size(), HRESULT_FROM_WIN32 (ERROR_WRITE_FAULT)); + } + + iFlush = fflush (pf); + CBREx (iFlush == 0, HRESULT_FROM_WIN32 (ERROR_WRITE_FAULT)); + + fclose (pf); + pf = nullptr; + + filesystem::rename (strTemp, strPath, ec); + CBREx (!ec, HRESULT_FROM_WIN32 (ERROR_WRITE_FAULT)); + +Error: + if (pf != nullptr) + { + fclose (pf); + } + + if (FAILED (hr)) + { + filesystem::remove (strTemp, ec); + } + + return hr; +} diff --git a/TCDirCore/ProfileFileManager.cpp b/TCDirCore/ProfileFileManager.cpp index c9a2480..42dc39e 100644 --- a/TCDirCore/ProfileFileManager.cpp +++ b/TCDirCore/ProfileFileManager.cpp @@ -20,6 +20,42 @@ static constexpr unsigned char k_rgUtf16BE[] = { 0xFE, 0xFF }; static constexpr LPCWSTR k_pszHeaderMarker = L"# TCDir Aliases"; static constexpr LPCWSTR k_pszFooterMarker = L"# End TCDir Aliases"; +// +// Default file access used when no accessor is injected +// + +static CProfileFileAccessReal s_realFileAccess; + + + + + +//////////////////////////////////////////////////////////////////////////////// +// +// CProfileFileManager::CProfileFileManager +// +//////////////////////////////////////////////////////////////////////////////// + +CProfileFileManager::CProfileFileManager() + : m_fileAccess (s_realFileAccess) +{ +} + + + + + +//////////////////////////////////////////////////////////////////////////////// +// +// CProfileFileManager::CProfileFileManager +// +//////////////////////////////////////////////////////////////////////////////// + +CProfileFileManager::CProfileFileManager (IProfileFileAccess & fileAccess) + : m_fileAccess (fileAccess) +{ +} + @@ -35,9 +71,7 @@ static constexpr LPCWSTR k_pszFooterMarker = L"# End TCDir Aliases"; HRESULT CProfileFileManager::ReadProfileFile (const wstring & strPath, vector & rgLines, bool & fHasBom) { - HRESULT hr = S_OK; - FILE * pf = nullptr; - long cbFile = 0; + HRESULT hr = S_OK; string strRaw; @@ -46,101 +80,121 @@ HRESULT CProfileFileManager::ReadProfileFile (const wstring & strPath, vector 0) +Error: + if (FAILED (hr)) { - strRaw.resize (static_cast(cbFile)); - size_t cbRead = fread (strRaw.data(), 1, strRaw.size(), pf); - strRaw.resize (cbRead); + rgLines.clear(); } - fclose (pf); - pf = nullptr; + return hr; +} - // - // Check BOM - // - if (strRaw.size() >= 2) - { - auto pb = reinterpret_cast(strRaw.data()); - if ((pb[0] == k_rgUtf16LE[0] && pb[1] == k_rgUtf16LE[1]) || - (pb[0] == k_rgUtf16BE[0] && pb[1] == k_rgUtf16BE[1])) - { - // - // UTF-16 — bail with clear error - // - CBRAEx (false, HRESULT_FROM_WIN32 (ERROR_UNSUPPORTED_TYPE)); - } - } - if (strRaw.size() >= 3) +//////////////////////////////////////////////////////////////////////////////// +// +// CProfileFileManager::ParseProfileBytes +// +// Strips a UTF-8 BOM, rejects UTF-16, and decodes the bytes into lines. +// +//////////////////////////////////////////////////////////////////////////////// + +HRESULT CProfileFileManager::ParseProfileBytes (const string & strRaw, vector & rgLines, bool & fHasBom) +{ + HRESULT hr = S_OK; + const unsigned char * pb = nullptr; + string strBody = strRaw; + int cchNeeded = 0; + wstring strWide; + + + + rgLines.clear(); + fHasBom = false; + + pb = reinterpret_cast(strBody.data()); + + if (strBody.size() >= 2) { - auto pb = reinterpret_cast(strRaw.data()); + bool fUtf16 = (pb[0] == k_rgUtf16LE[0] && pb[1] == k_rgUtf16LE[1]) || + (pb[0] == k_rgUtf16BE[0] && pb[1] == k_rgUtf16BE[1]); - if (pb[0] == k_rgUtf8Bom[0] && pb[1] == k_rgUtf8Bom[1] && pb[2] == k_rgUtf8Bom[2]) - { - fHasBom = true; - strRaw.erase (0, 3); - } + CBREx (!fUtf16, HRESULT_FROM_WIN32 (ERROR_UNSUPPORTED_TYPE)); } - // - // Convert UTF-8 to wstring lines - // - - if (!strRaw.empty()) + if (strBody.size() >= 3 && pb[0] == k_rgUtf8Bom[0] && pb[1] == k_rgUtf8Bom[1] && pb[2] == k_rgUtf8Bom[2]) { - int cchNeeded = MultiByteToWideChar (CP_UTF8, 0, strRaw.data(), static_cast(strRaw.size()), nullptr, 0); + fHasBom = true; + strBody.erase (0, 3); + } - CBRAEx (cchNeeded > 0, HRESULT_FROM_WIN32 (GetLastError())); + BAIL_OUT_IF (strBody.empty(), S_OK); - wstring strWide (static_cast(cchNeeded), L'\0'); + cchNeeded = MultiByteToWideChar (CP_UTF8, 0, strBody.data(), static_cast(strBody.size()), nullptr, 0); + CBRAEx (cchNeeded > 0, HRESULT_FROM_WIN32 (GetLastError())); - MultiByteToWideChar (CP_UTF8, 0, strRaw.data(), static_cast(strRaw.size()), strWide.data(), cchNeeded); + strWide.resize (static_cast(cchNeeded), L'\0'); - // - // Split into lines (handle \r\n, \n, \r) - // + MultiByteToWideChar (CP_UTF8, 0, strBody.data(), static_cast(strBody.size()), strWide.data(), cchNeeded); - size_t pos = 0; + SplitIntoLines (strWide, rgLines); - while (pos < strWide.size()) - { - size_t end = strWide.find_first_of (L"\r\n", pos); +Error: + return hr; +} - if (end == wstring::npos) - { - rgLines.push_back (strWide.substr (pos)); - break; - } - rgLines.push_back (strWide.substr (pos, end - pos)); - if (end + 1 < strWide.size() && strWide[end] == L'\r' && strWide[end + 1] == L'\n') - { - pos = end + 2; - } - else - { - pos = end + 1; - } + + +//////////////////////////////////////////////////////////////////////////////// +// +// CProfileFileManager::SplitIntoLines +// +// Splits on \r\n, \n, or \r. +// +//////////////////////////////////////////////////////////////////////////////// + +void CProfileFileManager::SplitIntoLines (const wstring & strWide, vector & rgLines) +{ + size_t pos = 0; + size_t end = 0; + + + + while (pos < strWide.size()) + { + end = strWide.find_first_of (L"\r\n", pos); + + if (end == wstring::npos) + { + rgLines.push_back (strWide.substr (pos)); + break; } - } -Error: - return hr; + rgLines.push_back (strWide.substr (pos, end - pos)); + + if (end + 1 < strWide.size() && strWide[end] == L'\r' && strWide[end + 1] == L'\n') + { + pos = end + 2; + } + else + { + pos = end + 1; + } + } } @@ -299,7 +353,7 @@ HRESULT CProfileFileManager::WriteProfileFile (const wstring & strPath, const ve wstring strContent; int cbNeeded = 0; string strUtf8; - FILE * pf = nullptr; + string strPayload; @@ -352,21 +406,15 @@ HRESULT CProfileFileManager::WriteProfileFile (const wstring & strPath, const ve // Write to file // - _wfopen_s (&pf, strPath.c_str(), L"wb"); - CBRAEx (pf != nullptr, HRESULT_FROM_WIN32 (ERROR_ACCESS_DENIED)); - if (fPreserveBom) { - fwrite (k_rgUtf8Bom, 1, sizeof (k_rgUtf8Bom), pf); + strPayload.assign (reinterpret_cast(k_rgUtf8Bom), sizeof (k_rgUtf8Bom)); } - if (!strUtf8.empty()) - { - fwrite (strUtf8.data(), 1, strUtf8.size(), pf); - } + strPayload += strUtf8; - fclose (pf); - pf = nullptr; + hr = m_fileAccess.WriteAllBytes (strPath, strPayload); + CHR (hr); Error: return hr; diff --git a/TCDirCore/ProfileFileManager.h b/TCDirCore/ProfileFileManager.h index ccf1641..cc714e1 100644 --- a/TCDirCore/ProfileFileManager.h +++ b/TCDirCore/ProfileFileManager.h @@ -1,6 +1,7 @@ #pragma once #include "ProfilePathResolver.h" +#include "ProfileFileAccess.h" @@ -40,6 +41,9 @@ struct SAliasBlock class CProfileFileManager { public: + CProfileFileManager(); + explicit CProfileFileManager (IProfileFileAccess & fileAccess); + HRESULT ReadProfileFile (const wstring & strPath, vector & rgLines, bool & fHasBom); HRESULT FindAliasBlock (const vector & rgLines, SAliasBlock & block); HRESULT WriteProfileFile (const wstring & strPath, const vector & rgLines, bool fPreserveBom); @@ -47,4 +51,10 @@ class CProfileFileManager void ReplaceAliasBlock (vector & rgLines, const SAliasBlock & block, const vector & rgNewBlock); void AppendAliasBlock (vector & rgLines, const vector & rgNewBlock); void RemoveAliasBlock (vector & rgLines, const SAliasBlock & block); + +private: + HRESULT ParseProfileBytes (const string & strRaw, vector & rgLines, bool & fHasBom); + void SplitIntoLines (const wstring & strWide, vector & rgLines); + + IProfileFileAccess & m_fileAccess; }; diff --git a/TCDirCore/TCDirCore.vcxproj b/TCDirCore/TCDirCore.vcxproj index 3dd7602..fcb4894 100644 --- a/TCDirCore/TCDirCore.vcxproj +++ b/TCDirCore/TCDirCore.vcxproj @@ -244,6 +244,7 @@ + diff --git a/UnitTest/ProfileFileManagerTests.cpp b/UnitTest/ProfileFileManagerTests.cpp index 5051003..60cd91d 100644 --- a/UnitTest/ProfileFileManagerTests.cpp +++ b/UnitTest/ProfileFileManagerTests.cpp @@ -16,6 +16,55 @@ using namespace Microsoft::VisualStudio::CppUnitTestFramework; namespace UnitTest { + //////////////////////////////////////////////////////////////////////////// + // + // CProfileFileAccessMock + // + // In-memory stand-in for the profile file I/O, so read/write failure + // modes can be exercised without touching the real file system. + // + //////////////////////////////////////////////////////////////////////////// + + class CProfileFileAccessMock : public IProfileFileAccess + { + public: + + HRESULT ReadAllBytes (const wstring & strPath, string & strBytes) override + { + m_strReadPath = strPath; + + strBytes.clear(); + + if (SUCCEEDED (m_hrRead)) + { + strBytes = m_strContent; + } + + return m_hrRead; + } + + HRESULT WriteAllBytes (const wstring & strPath, const string & strBytes) override + { + m_strWritePath = strPath; + m_strWritten = strBytes; + m_cWriteCalls += 1; + + return m_hrWrite; + } + + HRESULT m_hrRead = S_OK; + HRESULT m_hrWrite = S_OK; + string m_strContent; + string m_strWritten; + wstring m_strReadPath; + wstring m_strWritePath; + int m_cWriteCalls = 0; + }; + + + + + TEST_CLASS(ProfileFileManagerTests) { public: @@ -163,5 +212,230 @@ namespace UnitTest Assert::AreEqual (L"# before", rgLines[0].c_str()); Assert::AreEqual (L"# after", rgLines[1].c_str()); } + + + + + // + // A failed read must never surface as success with no content. That is + // the shape that let a whole profile get replaced by just an alias + // block: read silently returned nothing, and the caller wrote it back. + // + + TEST_METHOD(ReadProfileFile_ReadFails_ReturnsFailureAndNoLines) + { + CProfileFileAccessMock mock; + vector rgLines; + bool fHasBom = false; + HRESULT hr = S_OK; + + mock.m_hrRead = HRESULT_FROM_WIN32 (ERROR_READ_FAULT); + + CProfileFileManager fileMgr (mock); + + + + hr = fileMgr.ReadProfileFile (L"C:\\test\\profile.ps1", rgLines, fHasBom); + + Assert::IsTrue (FAILED (hr)); + Assert::IsTrue (rgLines.empty()); + } + + + + + // + // The caller's vector must not keep stale content after a failed read. + // + + TEST_METHOD(ReadProfileFile_ReadFails_ClearsCallerLines) + { + CProfileFileAccessMock mock; + vector rgLines = { L"stale one", L"stale two" }; + bool fHasBom = false; + HRESULT hr = S_OK; + + mock.m_hrRead = HRESULT_FROM_WIN32 (ERROR_READ_FAULT); + + CProfileFileManager fileMgr (mock); + + + + hr = fileMgr.ReadProfileFile (L"C:\\test\\profile.ps1", rgLines, fHasBom); + + Assert::IsTrue (FAILED (hr)); + Assert::IsTrue (rgLines.empty()); + } + + + + + // + // A genuinely empty file is still a success — the fix must not turn + // "nothing to read" into an error. + // + + TEST_METHOD(ReadProfileFile_EmptyFile_SucceedsWithNoLines) + { + CProfileFileAccessMock mock; + vector rgLines; + bool fHasBom = false; + HRESULT hr = S_OK; + + CProfileFileManager fileMgr (mock); + + + + hr = fileMgr.ReadProfileFile (L"C:\\test\\profile.ps1", rgLines, fHasBom); + + Assert::IsTrue (SUCCEEDED (hr)); + Assert::IsTrue (rgLines.empty()); + Assert::IsFalse (fHasBom); + } + + + + + TEST_METHOD(ReadProfileFile_ContentSplitsIntoLines) + { + CProfileFileAccessMock mock; + vector rgLines; + bool fHasBom = false; + HRESULT hr = S_OK; + + mock.m_strContent = "first\r\nsecond\nthird\rfourth"; + + CProfileFileManager fileMgr (mock); + + + + hr = fileMgr.ReadProfileFile (L"C:\\test\\profile.ps1", rgLines, fHasBom); + + Assert::IsTrue (SUCCEEDED (hr)); + Assert::AreEqual (4u, static_cast(rgLines.size())); + Assert::AreEqual (L"first", rgLines[0].c_str()); + Assert::AreEqual (L"second", rgLines[1].c_str()); + Assert::AreEqual (L"third", rgLines[2].c_str()); + Assert::AreEqual (L"fourth", rgLines[3].c_str()); + } + + + + + TEST_METHOD(ReadProfileFile_Utf8Bom_StrippedAndFlagged) + { + CProfileFileAccessMock mock; + vector rgLines; + bool fHasBom = false; + HRESULT hr = S_OK; + string strWithBom; + + strWithBom.push_back (static_cast(0xEF)); + strWithBom.push_back (static_cast(0xBB)); + strWithBom.push_back (static_cast(0xBF)); + strWithBom += "content"; + + mock.m_strContent = strWithBom; + + CProfileFileManager fileMgr (mock); + + + + hr = fileMgr.ReadProfileFile (L"C:\\test\\profile.ps1", rgLines, fHasBom); + + Assert::IsTrue (SUCCEEDED (hr)); + Assert::IsTrue (fHasBom); + Assert::AreEqual (1u, static_cast(rgLines.size())); + Assert::AreEqual (L"content", rgLines[0].c_str()); + } + + + + + TEST_METHOD(ReadProfileFile_Utf16_IsRejected) + { + CProfileFileAccessMock mock; + vector rgLines; + bool fHasBom = false; + HRESULT hr = S_OK; + string strUtf16; + + strUtf16.push_back (static_cast(0xFF)); + strUtf16.push_back (static_cast(0xFE)); + strUtf16 += "content"; + + mock.m_strContent = strUtf16; + + CProfileFileManager fileMgr (mock); + + + + hr = fileMgr.ReadProfileFile (L"C:\\test\\profile.ps1", rgLines, fHasBom); + + Assert::IsTrue (FAILED (hr)); + Assert::IsTrue (rgLines.empty()); + } + + + + + TEST_METHOD(WriteProfileFile_JoinsLinesWithCrLfAndTrailingNewline) + { + CProfileFileAccessMock mock; + vector rgLines = { L"alpha", L"beta" }; + HRESULT hr = S_OK; + + CProfileFileManager fileMgr (mock); + + + + hr = fileMgr.WriteProfileFile (L"C:\\test\\profile.ps1", rgLines, false); + + Assert::IsTrue (SUCCEEDED (hr)); + Assert::AreEqual (1, mock.m_cWriteCalls); + Assert::AreEqual ("alpha\r\nbeta\r\n", mock.m_strWritten.c_str()); + } + + + + + TEST_METHOD(WriteProfileFile_PreserveBom_EmitsBomBytes) + { + CProfileFileAccessMock mock; + vector rgLines = { L"alpha" }; + HRESULT hr = S_OK; + + CProfileFileManager fileMgr (mock); + + + + hr = fileMgr.WriteProfileFile (L"C:\\test\\profile.ps1", rgLines, true); + + Assert::IsTrue (SUCCEEDED (hr)); + Assert::IsTrue (mock.m_strWritten.size() > 3); + Assert::AreEqual (static_cast(0xEF), static_cast(mock.m_strWritten[0])); + Assert::AreEqual (static_cast(0xBB), static_cast(mock.m_strWritten[1])); + Assert::AreEqual (static_cast(0xBF), static_cast(mock.m_strWritten[2])); + } + + + + + TEST_METHOD(WriteProfileFile_WriteFails_PropagatesFailure) + { + CProfileFileAccessMock mock; + vector rgLines = { L"alpha" }; + HRESULT hr = S_OK; + + mock.m_hrWrite = HRESULT_FROM_WIN32 (ERROR_WRITE_FAULT); + + CProfileFileManager fileMgr (mock); + + + + hr = fileMgr.WriteProfileFile (L"C:\\test\\profile.ps1", rgLines, false); + + Assert::IsTrue (FAILED (hr)); + } }; }