michael@0: /* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
michael@0: /* vim: se cin sw=2 ts=2 et : */
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 "nsDownloadScanner.h"
michael@0: #include <comcat.h>
michael@0: #include <process.h>
michael@0: #include "nsDownloadManager.h"
michael@0: #include "nsIXULAppInfo.h"
michael@0: #include "nsXULAppAPI.h"
michael@0: #include "nsIPrefService.h"
michael@0: #include "nsNetUtil.h"
michael@0: #include "nsDeque.h"
michael@0: #include "nsIFileURL.h"
michael@0: #include "nsIPrefBranch.h"
michael@0: #include "nsXPCOMCIDInternal.h"
michael@0: 
michael@0: /**
michael@0:  * Code overview
michael@0:  *
michael@0:  * Download scanner attempts to make use of one of two different virus
michael@0:  * scanning interfaces available on Windows - IOfficeAntiVirus (Windows
michael@0:  * 95/NT 4 and IE 5) and IAttachmentExecute (XPSP2 and up).  The latter
michael@0:  * interface supports calling IOfficeAntiVirus internally, while also
michael@0:  * adding support for XPSP2+ ADS forks which define security related
michael@0:  * prompting on downloaded content.  
michael@0:  *
michael@0:  * Both interfaces are synchronous and can take a while, so it is not a
michael@0:  * good idea to call either from the main thread. Some antivirus scanners can
michael@0:  * take a long time to scan or the call might block while the scanner shows
michael@0:  * its UI so if the user were to download many files that finished around the
michael@0:  * same time, they would have to wait a while if the scanning were done on
michael@0:  * exactly one other thread. Since the overhead of creating a thread is
michael@0:  * relatively small compared to the time it takes to download a file and scan
michael@0:  * it, a new thread is spawned for each download that is to be scanned. Since
michael@0:  * most of the mozilla codebase is not threadsafe, all the information needed
michael@0:  * for the scanner is gathered in the main thread in nsDownloadScanner::Scan::Start.
michael@0:  * The only function of nsDownloadScanner::Scan which is invoked on another
michael@0:  * thread is DoScan.
michael@0:  *
michael@0:  * Watchdog overview
michael@0:  *
michael@0:  * The watchdog is used internally by the scanner. It maintains a queue of
michael@0:  * current download scans. In a separate thread, it dequeues the oldest scan
michael@0:  * and waits on that scan's thread with a timeout given by WATCHDOG_TIMEOUT
michael@0:  * (default is 30 seconds). If the wait times out, then the watchdog notifies
michael@0:  * the Scan that it has timed out. If the scan really has timed out, then the
michael@0:  * Scan object will dispatch its run method to the main thread; this will
michael@0:  * release the watchdog thread's addref on the Scan. If it has not timed out
michael@0:  * (i.e. the Scan just finished in time), then the watchdog dispatches a 
michael@0:  * ReleaseDispatcher to release its ref of the Scan on the main thread.
michael@0:  *
michael@0:  * In order to minimize execution time, there are two events used to notify the
michael@0:  * watchdog thread of a non-empty queue and a quit event. Every blocking wait 
michael@0:  * that the watchdog thread does waits on the quit event; this lets the thread
michael@0:  * quickly exit when shutting down. Also, the download scan queue will be empty
michael@0:  * most of the time; rather than use a spin loop, a simple event is triggered
michael@0:  * by the main thread when a new scan is added to an empty queue. When the
michael@0:  * watchdog thread knows that it has run out of elements in the queue, it will
michael@0:  * wait on the new item event.
michael@0:  *
michael@0:  * Memory/resource leaks due to timeout:
michael@0:  * In the event of a timeout, the thread must remain alive; terminating it may
michael@0:  * very well cause the antivirus scanner to crash or be put into an
michael@0:  * inconsistent state; COM resources may also not be cleaned up. The downside
michael@0:  * is that we need to leave the thread running; suspending it may lead to a
michael@0:  * deadlock. Because the scan call may be ongoing, it may be dependent on the
michael@0:  * memory referenced by the MSOAVINFO structure, so we cannot free mName, mPath
michael@0:  * or mOrigin; this means that we cannot free the Scan object since doing so
michael@0:  * will deallocate that memory. Note that mDownload is set to null upon timeout
michael@0:  * or completion, so the download itself is never leaked. If the scan does
michael@0:  * eventually complete, then the all the memory and resources will be freed.
michael@0:  * It is possible, however extremely rare, that in the event of a timeout, the
michael@0:  * mStateSync critical section will leak its event; this will happen only if
michael@0:  * the scanning thread, watchdog thread or main thread try to enter the
michael@0:  * critical section when one of the others is already in it.
michael@0:  *
michael@0:  * Reasoning for CheckAndSetState - there exists a race condition between the time when
michael@0:  * either the timeout or normal scan sets the state and when Scan::Run is
michael@0:  * executed on the main thread. Ex: mStatus could be set by Scan::DoScan* which
michael@0:  * then queues a dispatch on the main thread. Before that dispatch is executed,
michael@0:  * the timeout code fires and sets mStatus to AVSCAN_TIMEDOUT which then queues
michael@0:  * its dispatch to the main thread (the same function as DoScan*). Both
michael@0:  * dispatches run and both try to set the download state to AVSCAN_TIMEDOUT
michael@0:  * which is incorrect.
michael@0:  *
michael@0:  * There are 5 possible outcomes of the virus scan:
michael@0:  *    AVSCAN_GOOD     => the file is clean
michael@0:  *    AVSCAN_BAD      => the file has a virus
michael@0:  *    AVSCAN_UGLY     => the file had a virus, but it was cleaned
michael@0:  *    AVSCAN_FAILED   => something else went wrong with the virus scanner.
michael@0:  *    AVSCAN_TIMEDOUT => the scan (thread setup + execution) took too long
michael@0:  *
michael@0:  * Both the good and ugly states leave the user with a benign file, so they
michael@0:  * transition to the finished state. Bad files are sent to the blocked state.
michael@0:  * The failed and timedout states transition to finished downloads.
michael@0:  *
michael@0:  * Possible Future enhancements:
michael@0:  *  * Create an interface for scanning files in general
michael@0:  *  * Make this a service
michael@0:  *  * Get antivirus scanner status via WMI/registry
michael@0:  */
michael@0: 
michael@0: // IAttachementExecute supports user definable settings for certain
michael@0: // security related prompts. This defines a general GUID for use in
michael@0: // all projects. Individual projects can define an individual guid
michael@0: // if they want to.
michael@0: #ifndef MOZ_VIRUS_SCANNER_PROMPT_GUID
michael@0: #define MOZ_VIRUS_SCANNER_PROMPT_GUID \
michael@0:   { 0xb50563d1, 0x16b6, 0x43c2, { 0xa6, 0x6a, 0xfa, 0xe6, 0xd2, 0x11, 0xf2, \
michael@0:   0xea } }
michael@0: #endif
michael@0: static const GUID GUID_MozillaVirusScannerPromptGeneric =
michael@0:   MOZ_VIRUS_SCANNER_PROMPT_GUID;
michael@0: 
michael@0: // Initial timeout is 30 seconds
michael@0: #define WATCHDOG_TIMEOUT (30*PR_USEC_PER_SEC)
michael@0: 
michael@0: // Maximum length for URI's passed into IAE
michael@0: #define MAX_IAEURILENGTH 1683
michael@0: 
michael@0: class nsDownloadScannerWatchdog 
michael@0: {
michael@0:   typedef nsDownloadScanner::Scan Scan;
michael@0: public:
michael@0:   nsDownloadScannerWatchdog();
michael@0:   ~nsDownloadScannerWatchdog();
michael@0: 
michael@0:   nsresult Init();
michael@0:   nsresult Shutdown();
michael@0: 
michael@0:   void Watch(Scan *scan);
michael@0: private:
michael@0:   static unsigned int __stdcall WatchdogThread(void *p);
michael@0:   CRITICAL_SECTION mQueueSync;
michael@0:   nsDeque mScanQueue;
michael@0:   HANDLE mThread;
michael@0:   HANDLE mNewItemEvent;
michael@0:   HANDLE mQuitEvent;
michael@0: };
michael@0: 
michael@0: nsDownloadScanner::nsDownloadScanner() :
michael@0:   mAESExists(false)
michael@0: {
michael@0: }
michael@0:  
michael@0: // This destructor appeases the compiler; it would otherwise complain about an
michael@0: // incomplete type for nsDownloadWatchdog in the instantiation of
michael@0: // nsAutoPtr::~nsAutoPtr
michael@0: // Plus, it's a handy location to call nsDownloadScannerWatchdog::Shutdown from
michael@0: nsDownloadScanner::~nsDownloadScanner() {
michael@0:   if (mWatchdog)
michael@0:     (void)mWatchdog->Shutdown();
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsDownloadScanner::Init()
michael@0: {
michael@0:   // This CoInitialize/CoUninitialize pattern seems to be common in the Mozilla
michael@0:   // codebase. All other COM calls/objects are made on different threads.
michael@0:   nsresult rv = NS_OK;
michael@0:   CoInitialize(nullptr);
michael@0: 
michael@0:   if (!IsAESAvailable()) {
michael@0:     CoUninitialize();
michael@0:     return NS_ERROR_NOT_AVAILABLE;
michael@0:   }
michael@0: 
michael@0:   mAESExists = true;
michael@0: 
michael@0:   // Initialize scanning
michael@0:   mWatchdog = new nsDownloadScannerWatchdog();
michael@0:   if (mWatchdog) {
michael@0:     rv = mWatchdog->Init();
michael@0:     if (FAILED(rv))
michael@0:       mWatchdog = nullptr;
michael@0:   } else {
michael@0:     rv = NS_ERROR_OUT_OF_MEMORY;
michael@0:   }
michael@0:   
michael@0:   if (NS_FAILED(rv))
michael@0:     return rv;
michael@0: 
michael@0:   return rv;
michael@0: }
michael@0: 
michael@0: bool
michael@0: nsDownloadScanner::IsAESAvailable()
michael@0: {
michael@0:   // Try to instantiate IAE to see if it's available.    
michael@0:   nsRefPtr<IAttachmentExecute> ae;
michael@0:   HRESULT hr;
michael@0:   hr = CoCreateInstance(CLSID_AttachmentServices, nullptr, CLSCTX_INPROC,
michael@0:                         IID_IAttachmentExecute, getter_AddRefs(ae));
michael@0:   if (FAILED(hr)) {
michael@0:     NS_WARNING("Could not instantiate attachment execution service\n");
michael@0:     return false;
michael@0:   }
michael@0:   return true;
michael@0: }
michael@0: 
michael@0: // If IAttachementExecute is available, use the CheckPolicy call to find out
michael@0: // if this download should be prevented due to Security Zone Policy settings.
michael@0: AVCheckPolicyState
michael@0: nsDownloadScanner::CheckPolicy(nsIURI *aSource, nsIURI *aTarget)
michael@0: {
michael@0:   nsresult rv;
michael@0: 
michael@0:   if (!mAESExists || !aSource || !aTarget)
michael@0:     return AVPOLICY_DOWNLOAD;
michael@0: 
michael@0:   nsAutoCString source;
michael@0:   rv = aSource->GetSpec(source);
michael@0:   if (NS_FAILED(rv))
michael@0:     return AVPOLICY_DOWNLOAD;
michael@0: 
michael@0:   nsCOMPtr<nsIFileURL> fileUrl(do_QueryInterface(aTarget));
michael@0:   if (!fileUrl)
michael@0:     return AVPOLICY_DOWNLOAD;
michael@0: 
michael@0:   nsCOMPtr<nsIFile> theFile;
michael@0:   nsAutoString aFileName;
michael@0:   if (NS_FAILED(fileUrl->GetFile(getter_AddRefs(theFile))) ||
michael@0:       NS_FAILED(theFile->GetLeafName(aFileName)))
michael@0:     return AVPOLICY_DOWNLOAD;
michael@0: 
michael@0:   // IAttachementExecute prohibits src data: schemes by default but we
michael@0:   // support them. If this is a data src, skip off doing a policy check.
michael@0:   // (The file will still be scanned once it lands on the local system.)
michael@0:   bool isDataScheme(false);
michael@0:   nsCOMPtr<nsIURI> innerURI = NS_GetInnermostURI(aSource);
michael@0:   if (innerURI)
michael@0:     (void)innerURI->SchemeIs("data", &isDataScheme);
michael@0:   if (isDataScheme)
michael@0:     return AVPOLICY_DOWNLOAD;
michael@0: 
michael@0:   nsRefPtr<IAttachmentExecute> ae;
michael@0:   HRESULT hr;
michael@0:   hr = CoCreateInstance(CLSID_AttachmentServices, nullptr, CLSCTX_INPROC,
michael@0:                         IID_IAttachmentExecute, getter_AddRefs(ae));
michael@0:   if (FAILED(hr))
michael@0:     return AVPOLICY_DOWNLOAD;
michael@0: 
michael@0:   (void)ae->SetClientGuid(GUID_MozillaVirusScannerPromptGeneric);
michael@0:   (void)ae->SetSource(NS_ConvertUTF8toUTF16(source).get());
michael@0:   (void)ae->SetFileName(aFileName.get());
michael@0: 
michael@0:   // Any failure means the file download/exec will be blocked by the system.
michael@0:   // S_OK or S_FALSE imply it's ok.
michael@0:   hr = ae->CheckPolicy();
michael@0: 
michael@0:   if (hr == S_OK)
michael@0:     return AVPOLICY_DOWNLOAD;
michael@0: 
michael@0:   if (hr == S_FALSE)
michael@0:     return AVPOLICY_PROMPT;
michael@0: 
michael@0:   if (hr == E_INVALIDARG)
michael@0:     return AVPOLICY_PROMPT;
michael@0: 
michael@0:   return AVPOLICY_BLOCKED;
michael@0: }
michael@0: 
michael@0: #ifndef THREAD_MODE_BACKGROUND_BEGIN
michael@0: #define THREAD_MODE_BACKGROUND_BEGIN 0x00010000
michael@0: #endif
michael@0: 
michael@0: #ifndef THREAD_MODE_BACKGROUND_END
michael@0: #define THREAD_MODE_BACKGROUND_END 0x00020000
michael@0: #endif
michael@0: 
michael@0: unsigned int __stdcall
michael@0: nsDownloadScanner::ScannerThreadFunction(void *p)
michael@0: {
michael@0:   HANDLE currentThread = GetCurrentThread();
michael@0:   NS_ASSERTION(!NS_IsMainThread(), "Antivirus scan should not be run on the main thread");
michael@0:   nsDownloadScanner::Scan *scan = static_cast<nsDownloadScanner::Scan*>(p);
michael@0:   if (!SetThreadPriority(currentThread, THREAD_MODE_BACKGROUND_BEGIN))
michael@0:     (void)SetThreadPriority(currentThread, THREAD_PRIORITY_IDLE);
michael@0:   scan->DoScan();
michael@0:   (void)SetThreadPriority(currentThread, THREAD_MODE_BACKGROUND_END);
michael@0:   _endthreadex(0);
michael@0:   return 0;
michael@0: }
michael@0: 
michael@0: // The sole purpose of this class is to release an object on the main thread
michael@0: // It assumes that its creator will addref it and it will release itself on
michael@0: // the main thread too
michael@0: class ReleaseDispatcher : public nsRunnable {
michael@0: public:
michael@0:   ReleaseDispatcher(nsISupports *ptr)
michael@0:     : mPtr(ptr) {}
michael@0:   NS_IMETHOD Run();
michael@0: private:
michael@0:   nsISupports *mPtr;
michael@0: };
michael@0: 
michael@0: nsresult ReleaseDispatcher::Run() {
michael@0:   NS_ASSERTION(NS_IsMainThread(), "Antivirus scan release dispatch should be run on the main thread");
michael@0:   NS_RELEASE(mPtr);
michael@0:   NS_RELEASE_THIS();
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsDownloadScanner::Scan::Scan(nsDownloadScanner *scanner, nsDownload *download)
michael@0:   : mDLScanner(scanner), mThread(nullptr), 
michael@0:     mDownload(download), mStatus(AVSCAN_NOTSTARTED),
michael@0:     mSkipSource(false)
michael@0: {
michael@0:   InitializeCriticalSection(&mStateSync);
michael@0: }
michael@0: 
michael@0: nsDownloadScanner::Scan::~Scan() {
michael@0:   DeleteCriticalSection(&mStateSync);
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsDownloadScanner::Scan::Start()
michael@0: {
michael@0:   mStartTime = PR_Now();
michael@0: 
michael@0:   mThread = (HANDLE)_beginthreadex(nullptr, 0, ScannerThreadFunction,
michael@0:       this, CREATE_SUSPENDED, nullptr);
michael@0:   if (!mThread)
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0: 
michael@0:   nsresult rv = NS_OK;
michael@0: 
michael@0:   // Get the path to the file on disk
michael@0:   nsCOMPtr<nsIFile> file;
michael@0:   rv = mDownload->GetTargetFile(getter_AddRefs(file));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0:   rv = file->GetPath(mPath);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   // Grab the app name
michael@0:   nsCOMPtr<nsIXULAppInfo> appinfo =
michael@0:     do_GetService(XULAPPINFO_SERVICE_CONTRACTID, &rv);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   nsAutoCString name;
michael@0:   rv = appinfo->GetName(name);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0:   CopyUTF8toUTF16(name, mName);
michael@0: 
michael@0:   // Get the origin
michael@0:   nsCOMPtr<nsIURI> uri;
michael@0:   rv = mDownload->GetSource(getter_AddRefs(uri));
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   nsAutoCString origin;
michael@0:   rv = uri->GetSpec(origin);
michael@0:   NS_ENSURE_SUCCESS(rv, rv);
michael@0: 
michael@0:   // Certain virus interfaces do not like extremely long uris.
michael@0:   // Chop off the path and cgi data and just pass the base domain. 
michael@0:   if (origin.Length() > MAX_IAEURILENGTH) {
michael@0:     rv = uri->GetPrePath(origin);
michael@0:     NS_ENSURE_SUCCESS(rv, rv);
michael@0:   }
michael@0: 
michael@0:   CopyUTF8toUTF16(origin, mOrigin);
michael@0: 
michael@0:   // We count https/ftp/http as an http download
michael@0:   bool isHttp(false), isFtp(false), isHttps(false);
michael@0:   nsCOMPtr<nsIURI> innerURI = NS_GetInnermostURI(uri);
michael@0:   if (!innerURI) innerURI = uri;
michael@0:   (void)innerURI->SchemeIs("http", &isHttp);
michael@0:   (void)innerURI->SchemeIs("ftp", &isFtp);
michael@0:   (void)innerURI->SchemeIs("https", &isHttps);
michael@0:   mIsHttpDownload = isHttp || isFtp || isHttps;
michael@0: 
michael@0:   // IAttachementExecute prohibits src data: schemes by default but we
michael@0:   // support them. Mark the download if it's a data scheme, so we
michael@0:   // can skip off supplying the src to IAttachementExecute when we scan
michael@0:   // the resulting file.
michael@0:   (void)innerURI->SchemeIs("data", &mSkipSource);
michael@0: 
michael@0:   // ResumeThread returns the previous suspend count
michael@0:   if (1 != ::ResumeThread(mThread)) {
michael@0:     CloseHandle(mThread);
michael@0:     return NS_ERROR_UNEXPECTED;
michael@0:   }
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsDownloadScanner::Scan::Run()
michael@0: {
michael@0:   NS_ASSERTION(NS_IsMainThread(), "Antivirus scan dispatch should be run on the main thread");
michael@0: 
michael@0:   // Cleanup our thread
michael@0:   if (mStatus != AVSCAN_TIMEDOUT)
michael@0:     WaitForSingleObject(mThread, INFINITE);
michael@0:   CloseHandle(mThread);
michael@0: 
michael@0:   DownloadState downloadState = 0;
michael@0:   EnterCriticalSection(&mStateSync);
michael@0:   switch (mStatus) {
michael@0:     case AVSCAN_BAD:
michael@0:       downloadState = nsIDownloadManager::DOWNLOAD_DIRTY;
michael@0:       break;
michael@0:     default:
michael@0:     case AVSCAN_FAILED:
michael@0:     case AVSCAN_GOOD:
michael@0:     case AVSCAN_UGLY:
michael@0:     case AVSCAN_TIMEDOUT:
michael@0:       downloadState = nsIDownloadManager::DOWNLOAD_FINISHED;
michael@0:       break;
michael@0:   }
michael@0:   LeaveCriticalSection(&mStateSync);
michael@0:   // Download will be null if we already timed out
michael@0:   if (mDownload)
michael@0:     (void)mDownload->SetState(downloadState);
michael@0: 
michael@0:   // Clean up some other variables
michael@0:   // In the event of a timeout, our destructor won't be called
michael@0:   mDownload = nullptr;
michael@0: 
michael@0:   NS_RELEASE_THIS();
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: static DWORD
michael@0: ExceptionFilterFunction(DWORD exceptionCode) {
michael@0:   switch(exceptionCode) {
michael@0:     case EXCEPTION_ACCESS_VIOLATION:
michael@0:     case EXCEPTION_ILLEGAL_INSTRUCTION:
michael@0:     case EXCEPTION_IN_PAGE_ERROR:
michael@0:     case EXCEPTION_PRIV_INSTRUCTION:
michael@0:     case EXCEPTION_STACK_OVERFLOW:
michael@0:       return EXCEPTION_EXECUTE_HANDLER;
michael@0:     default:
michael@0:       return EXCEPTION_CONTINUE_SEARCH;
michael@0:   }
michael@0: }
michael@0: 
michael@0: bool
michael@0: nsDownloadScanner::Scan::DoScanAES()
michael@0: {
michael@0:   // This warning is for the destructor of ae which will not be invoked in the
michael@0:   // event of a win32 exception
michael@0: #pragma warning(disable: 4509)
michael@0:   HRESULT hr;
michael@0:   nsRefPtr<IAttachmentExecute> ae;
michael@0:   MOZ_SEH_TRY {
michael@0:     hr = CoCreateInstance(CLSID_AttachmentServices, nullptr, CLSCTX_ALL,
michael@0:                           IID_IAttachmentExecute, getter_AddRefs(ae));
michael@0:   } MOZ_SEH_EXCEPT(ExceptionFilterFunction(GetExceptionCode())) {
michael@0:     return CheckAndSetState(AVSCAN_NOTSTARTED,AVSCAN_FAILED);
michael@0:   }
michael@0: 
michael@0:   // If we (somehow) already timed out, then don't bother scanning
michael@0:   if (CheckAndSetState(AVSCAN_SCANNING, AVSCAN_NOTSTARTED)) {
michael@0:     AVScanState newState;
michael@0:     if (SUCCEEDED(hr)) {
michael@0:       bool gotException = false;
michael@0:       MOZ_SEH_TRY {
michael@0:         (void)ae->SetClientGuid(GUID_MozillaVirusScannerPromptGeneric);
michael@0:         (void)ae->SetLocalPath(mPath.get());
michael@0:         // Provide the src for everything but data: schemes.
michael@0:         if (!mSkipSource)
michael@0:           (void)ae->SetSource(mOrigin.get());
michael@0: 
michael@0:         // Save() will invoke the scanner
michael@0:         hr = ae->Save();
michael@0:       } MOZ_SEH_EXCEPT(ExceptionFilterFunction(GetExceptionCode())) {
michael@0:         gotException = true;
michael@0:       }
michael@0: 
michael@0:       MOZ_SEH_TRY {
michael@0:         ae = nullptr;
michael@0:       } MOZ_SEH_EXCEPT(ExceptionFilterFunction(GetExceptionCode())) {
michael@0:         gotException = true;
michael@0:       }
michael@0: 
michael@0:       if(gotException) {
michael@0:         newState = AVSCAN_FAILED;
michael@0:       }
michael@0:       else if (SUCCEEDED(hr)) { // Passed the scan
michael@0:         newState = AVSCAN_GOOD;
michael@0:       }
michael@0:       else if (HRESULT_CODE(hr) == ERROR_FILE_NOT_FOUND) {
michael@0:         NS_WARNING("Downloaded file disappeared before it could be scanned");
michael@0:         newState = AVSCAN_FAILED;
michael@0:       }
michael@0:       else if (hr == E_INVALIDARG) {
michael@0:         NS_WARNING("IAttachementExecute returned invalid argument error");
michael@0:         newState = AVSCAN_FAILED;
michael@0:       }
michael@0:       else { 
michael@0:         newState = AVSCAN_UGLY;
michael@0:       }
michael@0:     }
michael@0:     else {
michael@0:       newState = AVSCAN_FAILED;
michael@0:     }
michael@0:     return CheckAndSetState(newState, AVSCAN_SCANNING);
michael@0:   }
michael@0:   return false;
michael@0: }
michael@0: #pragma warning(default: 4509)
michael@0: 
michael@0: void
michael@0: nsDownloadScanner::Scan::DoScan()
michael@0: {
michael@0:   CoInitialize(nullptr);
michael@0: 
michael@0:   if (DoScanAES()) {
michael@0:     // We need to do a few more things on the main thread
michael@0:     NS_DispatchToMainThread(this);
michael@0:   } else {
michael@0:     // We timed out, so just release
michael@0:     ReleaseDispatcher* releaser = new ReleaseDispatcher(this);
michael@0:     if(releaser) {
michael@0:       NS_ADDREF(releaser);
michael@0:       NS_DispatchToMainThread(releaser);
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   MOZ_SEH_TRY {
michael@0:     CoUninitialize();
michael@0:   } MOZ_SEH_EXCEPT(ExceptionFilterFunction(GetExceptionCode())) {
michael@0:     // Not much we can do at this point...
michael@0:   }
michael@0: }
michael@0: 
michael@0: HANDLE
michael@0: nsDownloadScanner::Scan::GetWaitableThreadHandle() const
michael@0: {
michael@0:   HANDLE targetHandle = INVALID_HANDLE_VALUE;
michael@0:   (void)DuplicateHandle(GetCurrentProcess(), mThread,
michael@0:                         GetCurrentProcess(), &targetHandle,
michael@0:                         SYNCHRONIZE, // Only allow clients to wait on this handle
michael@0:                         FALSE, // cannot be inherited by child processes
michael@0:                         0);
michael@0:   return targetHandle;
michael@0: }
michael@0: 
michael@0: bool
michael@0: nsDownloadScanner::Scan::NotifyTimeout()
michael@0: {
michael@0:   bool didTimeout = CheckAndSetState(AVSCAN_TIMEDOUT, AVSCAN_SCANNING) ||
michael@0:                       CheckAndSetState(AVSCAN_TIMEDOUT, AVSCAN_NOTSTARTED);
michael@0:   if (didTimeout) {
michael@0:     // We need to do a few more things on the main thread
michael@0:     NS_DispatchToMainThread(this);
michael@0:   }
michael@0:   return didTimeout;
michael@0: }
michael@0: 
michael@0: bool
michael@0: nsDownloadScanner::Scan::CheckAndSetState(AVScanState newState, AVScanState expectedState) {
michael@0:   bool gotExpectedState = false;
michael@0:   EnterCriticalSection(&mStateSync);
michael@0:   if(gotExpectedState = (mStatus == expectedState))
michael@0:     mStatus = newState;
michael@0:   LeaveCriticalSection(&mStateSync);
michael@0:   return gotExpectedState;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsDownloadScanner::ScanDownload(nsDownload *download)
michael@0: {
michael@0:   if (!mAESExists)
michael@0:     return NS_ERROR_NOT_AVAILABLE;
michael@0: 
michael@0:   // No ref ptr, see comment below
michael@0:   Scan *scan = new Scan(this, download);
michael@0:   if (!scan)
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0: 
michael@0:   NS_ADDREF(scan);
michael@0: 
michael@0:   nsresult rv = scan->Start();
michael@0: 
michael@0:   // Note that we only release upon error. On success, the scan is passed off
michael@0:   // to a new thread. It is eventually released in Scan::Run on the main thread.
michael@0:   if (NS_FAILED(rv))
michael@0:     NS_RELEASE(scan);
michael@0:   else
michael@0:     // Notify the watchdog
michael@0:     mWatchdog->Watch(scan);
michael@0: 
michael@0:   return rv;
michael@0: }
michael@0: 
michael@0: nsDownloadScannerWatchdog::nsDownloadScannerWatchdog() 
michael@0:   : mNewItemEvent(nullptr), mQuitEvent(nullptr) {
michael@0:   InitializeCriticalSection(&mQueueSync);
michael@0: }
michael@0: nsDownloadScannerWatchdog::~nsDownloadScannerWatchdog() {
michael@0:   DeleteCriticalSection(&mQueueSync);
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsDownloadScannerWatchdog::Init() {
michael@0:   // Both events are auto-reset
michael@0:   mNewItemEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr);
michael@0:   if (INVALID_HANDLE_VALUE == mNewItemEvent)
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0:   mQuitEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr);
michael@0:   if (INVALID_HANDLE_VALUE == mQuitEvent) {
michael@0:     (void)CloseHandle(mNewItemEvent);
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0:   }
michael@0: 
michael@0:   // This thread is always running, however it will be asleep
michael@0:   // for most of the dlmgr's lifetime
michael@0:   mThread = (HANDLE)_beginthreadex(nullptr, 0, WatchdogThread,
michael@0:                                    this, 0, nullptr);
michael@0:   if (!mThread) {
michael@0:     (void)CloseHandle(mNewItemEvent);
michael@0:     (void)CloseHandle(mQuitEvent);
michael@0:     return NS_ERROR_OUT_OF_MEMORY;
michael@0:   }
michael@0: 
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: nsresult
michael@0: nsDownloadScannerWatchdog::Shutdown() {
michael@0:   // Tell the watchdog thread to quite
michael@0:   (void)SetEvent(mQuitEvent);
michael@0:   (void)WaitForSingleObject(mThread, INFINITE);
michael@0:   (void)CloseHandle(mThread);
michael@0:   // Manually clear and release the queued scans
michael@0:   while (mScanQueue.GetSize() != 0) {
michael@0:     Scan *scan = reinterpret_cast<Scan*>(mScanQueue.Pop());
michael@0:     NS_RELEASE(scan);
michael@0:   }
michael@0:   (void)CloseHandle(mNewItemEvent);
michael@0:   (void)CloseHandle(mQuitEvent);
michael@0:   return NS_OK;
michael@0: }
michael@0: 
michael@0: void
michael@0: nsDownloadScannerWatchdog::Watch(Scan *scan) {
michael@0:   bool wasEmpty;
michael@0:   // Note that there is no release in this method
michael@0:   // The scan will be released by the watchdog ALWAYS on the main thread
michael@0:   // when either the watchdog thread processes the scan or the watchdog
michael@0:   // is shut down
michael@0:   NS_ADDREF(scan);
michael@0:   EnterCriticalSection(&mQueueSync);
michael@0:   wasEmpty = mScanQueue.GetSize()==0;
michael@0:   mScanQueue.Push(scan);
michael@0:   LeaveCriticalSection(&mQueueSync);
michael@0:   // If the queue was empty, then the watchdog thread is/will be asleep
michael@0:   if (wasEmpty)
michael@0:     (void)SetEvent(mNewItemEvent);
michael@0: }
michael@0: 
michael@0: unsigned int
michael@0: __stdcall
michael@0: nsDownloadScannerWatchdog::WatchdogThread(void *p) {
michael@0:   NS_ASSERTION(!NS_IsMainThread(), "Antivirus scan watchdog should not be run on the main thread");
michael@0:   nsDownloadScannerWatchdog *watchdog = (nsDownloadScannerWatchdog*)p;
michael@0:   HANDLE waitHandles[3] = {watchdog->mNewItemEvent, watchdog->mQuitEvent, INVALID_HANDLE_VALUE};
michael@0:   DWORD waitStatus;
michael@0:   DWORD queueItemsLeft = 0;
michael@0:   // Loop until quit event or error
michael@0:   while (0 != queueItemsLeft ||
michael@0:          (WAIT_OBJECT_0 + 1) !=
michael@0:            (waitStatus =
michael@0:               WaitForMultipleObjects(2, waitHandles, FALSE, INFINITE)) &&
michael@0:          waitStatus != WAIT_FAILED) {
michael@0:     Scan *scan = nullptr;
michael@0:     PRTime startTime, expectedEndTime, now;
michael@0:     DWORD waitTime;
michael@0: 
michael@0:     // Pop scan from queue
michael@0:     EnterCriticalSection(&watchdog->mQueueSync);
michael@0:     scan = reinterpret_cast<Scan*>(watchdog->mScanQueue.Pop());
michael@0:     queueItemsLeft = watchdog->mScanQueue.GetSize();
michael@0:     LeaveCriticalSection(&watchdog->mQueueSync);
michael@0: 
michael@0:     // Calculate expected end time
michael@0:     startTime = scan->GetStartTime();
michael@0:     expectedEndTime = WATCHDOG_TIMEOUT + startTime;
michael@0:     now = PR_Now();
michael@0:     // PRTime is not guaranteed to be a signed integral type (afaik), but
michael@0:     // currently it is
michael@0:     if (now > expectedEndTime) {
michael@0:       waitTime = 0;
michael@0:     } else {
michael@0:       // This is a positive value, and we know that it will not overflow
michael@0:       // (bounded by WATCHDOG_TIMEOUT)
michael@0:       // waitTime is in milliseconds, nspr uses microseconds
michael@0:       waitTime = static_cast<DWORD>((expectedEndTime - now)/PR_USEC_PER_MSEC);
michael@0:     }
michael@0:     HANDLE hThread = waitHandles[2] = scan->GetWaitableThreadHandle();
michael@0: 
michael@0:     // Wait for the thread (obj 1) or quit event (obj 0)
michael@0:     waitStatus = WaitForMultipleObjects(2, (waitHandles+1), FALSE, waitTime);
michael@0:     CloseHandle(hThread);
michael@0: 
michael@0:     ReleaseDispatcher* releaser = new ReleaseDispatcher(scan);
michael@0:     if(!releaser)
michael@0:       continue;
michael@0:     NS_ADDREF(releaser);
michael@0:     // Got quit event or error
michael@0:     if (waitStatus == WAIT_FAILED || waitStatus == WAIT_OBJECT_0) {
michael@0:       NS_DispatchToMainThread(releaser);
michael@0:       break;
michael@0:     // Thread exited normally
michael@0:     } else if (waitStatus == (WAIT_OBJECT_0+1)) {
michael@0:       NS_DispatchToMainThread(releaser);
michael@0:       continue;
michael@0:     // Timeout case
michael@0:     } else { 
michael@0:       NS_ASSERTION(waitStatus == WAIT_TIMEOUT, "Unexpected wait status in dlmgr watchdog thread");
michael@0:       if (!scan->NotifyTimeout()) {
michael@0:         // If we didn't time out, then release the thread
michael@0:         NS_DispatchToMainThread(releaser);
michael@0:       } else {
michael@0:         // NotifyTimeout did a dispatch which will release the scan, so we
michael@0:         // don't need to release the scan
michael@0:         NS_RELEASE(releaser);
michael@0:       }
michael@0:     }
michael@0:   }
michael@0:   _endthreadex(0);
michael@0:   return 0;
michael@0: }