file_sink.cpp (1687B)
1 /* This Source Code Form is subject to the terms of the Mozilla Public 2 * License, v. 2.0. If a copy of the MPL was not distributed with this 3 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ 4 5 #include "file_sink.h" 6 #include <iostream> 7 #include <string> 8 #include <windows.h> 9 10 // Open the download receiver 11 bool FileSink::open(std::wstring& filename) { 12 mFilename.assign(filename); 13 // Note: this only succeeds if the filename does not exist. 14 fileHandle.own(CreateFileW(filename.c_str(), GENERIC_WRITE, 15 FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr, 16 CREATE_NEW, FILE_ATTRIBUTE_NORMAL, nullptr)); 17 if (fileHandle.get() == INVALID_HANDLE_VALUE) { 18 return false; 19 } else { 20 return true; 21 } 22 } 23 24 // Send data to the download receiver 25 bool FileSink::accept(char* buf, int bytesToWrite) { 26 while (bytesToWrite > 0) { 27 DWORD bytesWritten = 0; 28 if (!WriteFile(fileHandle, buf, bytesToWrite, &bytesWritten, nullptr)) { 29 std::cout << "Failed to write! " << GetLastError() << std::endl; 30 return false; // Some kind of error happened 31 } 32 bytesToWrite -= bytesWritten; 33 } 34 return true; 35 } 36 37 bool FileSink::freeze() { 38 // Exchange our write-only handle for a read-only one. This allows us to 39 // prevent renaming or deleting the file under our nose, while also being 40 // able to actually run it. 41 HANDLE readOnlyHandle = CreateFileW( 42 mFilename.c_str(), GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 43 nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); 44 if (readOnlyHandle == INVALID_HANDLE_VALUE) { 45 return false; 46 } 47 48 fileHandle.own(readOnlyHandle); 49 return true; 50 }