michael@0: /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
michael@0: /* vim: set ts=8 sts=2 et sw=2 tw=80: */
michael@0: /* This Source Code Form is subject to the terms of the Mozilla Public
michael@0:  * License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0:  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
michael@0: 
michael@0: #include "TestHarness.h"
michael@0: 
michael@0: #include <windows.h>
michael@0: #include <winnetwk.h>
michael@0: 
michael@0: #include "mozilla/FileUtilsWin.h"
michael@0: #include "mozilla/DebugOnly.h"
michael@0: #include "nsCRTGlue.h"
michael@0: 
michael@0: class DriveMapping
michael@0: {
michael@0: public:
michael@0:   DriveMapping(const nsAString& aRemoteUNCPath);
michael@0:   ~DriveMapping();
michael@0: 
michael@0:   bool
michael@0:   Init();
michael@0:   bool
michael@0:   ChangeDriveLetter();
michael@0:   wchar_t
michael@0:   GetDriveLetter() { return mDriveLetter; }
michael@0: 
michael@0: private:
michael@0:   bool
michael@0:   DoMapping();
michael@0:   void
michael@0:   Disconnect(wchar_t aDriveLetter);
michael@0: 
michael@0:   wchar_t   mDriveLetter;
michael@0:   nsString  mRemoteUNCPath;
michael@0: };
michael@0: 
michael@0: DriveMapping::DriveMapping(const nsAString& aRemoteUNCPath)
michael@0:   : mDriveLetter(0)
michael@0:   , mRemoteUNCPath(aRemoteUNCPath)
michael@0: {
michael@0: }
michael@0: 
michael@0: bool
michael@0: DriveMapping::Init()
michael@0: {
michael@0:   if (mDriveLetter) {
michael@0:     return false;
michael@0:   }
michael@0:   return DoMapping();
michael@0: }
michael@0: 
michael@0: bool
michael@0: DriveMapping::DoMapping()
michael@0: {
michael@0:   wchar_t drvTemplate[] = L" :";
michael@0:   NETRESOURCEW netRes = {0};
michael@0:   netRes.dwType = RESOURCETYPE_DISK;
michael@0:   netRes.lpLocalName = drvTemplate;
michael@0:   netRes.lpRemoteName = reinterpret_cast<wchar_t*>(mRemoteUNCPath.BeginWriting());
michael@0:   wchar_t driveLetter = L'D';
michael@0:   DWORD result = NO_ERROR;
michael@0:   do {
michael@0:     drvTemplate[0] = driveLetter;
michael@0:     result = WNetAddConnection2W(&netRes, nullptr, nullptr, CONNECT_TEMPORARY);
michael@0:   } while (result == ERROR_ALREADY_ASSIGNED && ++driveLetter <= L'Z');
michael@0:   if (result != NO_ERROR) {
michael@0:     return false;
michael@0:   }
michael@0:   mDriveLetter = driveLetter;
michael@0:   return true;
michael@0: }
michael@0: 
michael@0: bool
michael@0: DriveMapping::ChangeDriveLetter()
michael@0: {
michael@0:   wchar_t prevDriveLetter = mDriveLetter;
michael@0:   bool result = DoMapping();
michael@0:   MOZ_ASSERT(mDriveLetter != prevDriveLetter);
michael@0:   if (result && prevDriveLetter) {
michael@0:     Disconnect(prevDriveLetter);
michael@0:   }
michael@0:   return result;
michael@0: }
michael@0: 
michael@0: void
michael@0: DriveMapping::Disconnect(wchar_t aDriveLetter)
michael@0: {
michael@0:   wchar_t drvTemplate[] = {aDriveLetter, L':', L'\0'};
michael@0:   mozilla::DebugOnly<DWORD> result = WNetCancelConnection2W(drvTemplate, 0, TRUE);
michael@0:   MOZ_ASSERT(result == NO_ERROR);
michael@0: }
michael@0: 
michael@0: DriveMapping::~DriveMapping()
michael@0: {
michael@0:   if (mDriveLetter) {
michael@0:     Disconnect(mDriveLetter);
michael@0:   }
michael@0: }
michael@0: 
michael@0: bool
michael@0: DriveToNtPath(const wchar_t aDriveLetter, nsAString& aNtPath)
michael@0: {
michael@0:   const wchar_t drvTpl[] = {aDriveLetter, L':', L'\0'};
michael@0:   aNtPath.SetLength(MAX_PATH);
michael@0:   DWORD pathLen;
michael@0:   while (true) {
michael@0:     pathLen = QueryDosDeviceW(drvTpl, reinterpret_cast<wchar_t*>(aNtPath.BeginWriting()), aNtPath.Length());
michael@0:     if (pathLen || GetLastError() != ERROR_INSUFFICIENT_BUFFER) {
michael@0:       break;
michael@0:     }
michael@0:     aNtPath.SetLength(aNtPath.Length() * 2);
michael@0:   }
michael@0:   if (!pathLen) {
michael@0:     return false;
michael@0:   }
michael@0:   // aNtPath contains embedded NULLs, so we need to figure out the real length
michael@0:   // via wcslen.
michael@0:   aNtPath.SetLength(NS_strlen(aNtPath.BeginReading()));
michael@0:   return true;
michael@0: }
michael@0: 
michael@0: bool
michael@0: TestNtPathToDosPath(const wchar_t* aNtPath,
michael@0:                     const wchar_t* aExpectedDosPath)
michael@0: {
michael@0:   nsAutoString output;
michael@0:   bool result = mozilla::NtPathToDosPath(nsDependentString(aNtPath), output);
michael@0:   return result && output == aExpectedDosPath;
michael@0: }
michael@0: 
michael@0: int main(int argc, char* argv[])
michael@0: {
michael@0:   ScopedXPCOM xpcom("NtPathToDosPath");
michael@0:   if (xpcom.failed()) {
michael@0:     fail("XPCOM Startup");
michael@0:     return 1;
michael@0:   }
michael@0:   nsAutoString cDrive;
michael@0:   if (!DriveToNtPath(L'C', cDrive)) {
michael@0:     fail("Querying for this machine's C:");
michael@0:     return 1;
michael@0:   }
michael@0: 
michael@0:   int result = 0;
michael@0: 
michael@0:   // empty string
michael@0:   if (!TestNtPathToDosPath(L"", L"")) {
michael@0:     fail("Empty string");
michael@0:     result = 1;
michael@0:   }
michael@0:   // non-existent device, must fail
michael@0:   if (TestNtPathToDosPath(L"\\Device\\ThisDeviceDoesNotExist\\Foo", nullptr)) {
michael@0:     fail("Non-existent device");
michael@0:     result = 1;
michael@0:   }
michael@0:   // base case
michael@0:   nsAutoString testPath(cDrive);
michael@0:   testPath.Append(L"\\Foo");
michael@0:   if (!TestNtPathToDosPath(testPath.get(), L"C:\\Foo")) {
michael@0:     fail("Base case");
michael@0:     result = 1;
michael@0:   }
michael@0:   // drive letters as symbolic links (NtCreateFile uses these)
michael@0:   if (!TestNtPathToDosPath(L"\\??\\C:\\Foo", L"C:\\Foo")) {
michael@0:     fail("Path specified as symbolic link");
michael@0:     result = 1;
michael@0:   }
michael@0:   // other symbolic links (should fail)
michael@0:   if (TestNtPathToDosPath(L"\\??\\MountPointManager", nullptr)) {
michael@0:     fail("Other symbolic link");
michael@0:     result = 1;
michael@0:   }
michael@0:   // socket (should fail)
michael@0:   if (TestNtPathToDosPath(L"\\Device\\Afd\\Endpoint", nullptr)) {
michael@0:     fail("Socket");
michael@0:     result = 1;
michael@0:   }
michael@0:   // currently UNC paths that are not mapped to drive letters are unsupported,
michael@0:   // so this should fail
michael@0:   if (TestNtPathToDosPath(L"\\Device\\Mup\\127.0.0.1\\C$", nullptr)) {
michael@0:     fail("Unmapped UNC path");
michael@0:     result = 1;
michael@0:   }
michael@0:   DriveMapping drvMapping(NS_LITERAL_STRING("\\\\127.0.0.1\\C$"));
michael@0:   // Only run these tests if we were able to map; some machines don't have perms
michael@0:   if (drvMapping.Init()) {
michael@0:     wchar_t expected[] = L" :\\";
michael@0:     expected[0] = drvMapping.GetDriveLetter();
michael@0:     nsAutoString networkPath;
michael@0:     if (!DriveToNtPath(drvMapping.GetDriveLetter(), networkPath)) {
michael@0:       fail("Querying network drive");
michael@0:       return 1;
michael@0:     }
michael@0:     networkPath += MOZ_UTF16("\\");
michael@0:     if (!TestNtPathToDosPath(networkPath.get(), expected)) {
michael@0:       fail("Mapped UNC path");
michael@0:       result = 1;
michael@0:     }
michael@0:     // NtPathToDosPath must correctly handle paths whose drive letter mapping has
michael@0:     // changed. We need to test this because the APIs called by NtPathToDosPath
michael@0:     // return different info if this has happened.
michael@0:     if (!drvMapping.ChangeDriveLetter()) {
michael@0:       fail("Change drive letter");
michael@0:       return 1;
michael@0:     }
michael@0:     expected[0] = drvMapping.GetDriveLetter();
michael@0:     if (!DriveToNtPath(drvMapping.GetDriveLetter(), networkPath)) {
michael@0:       fail("Querying second network drive");
michael@0:       return 1;
michael@0:     }
michael@0:     networkPath += MOZ_UTF16("\\");
michael@0:     if (!TestNtPathToDosPath(networkPath.get(), expected)) {
michael@0:       fail("Re-mapped UNC path");
michael@0:       result = 1;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   return result;
michael@0: }
michael@0: