dom/webbrowserpersist/nsWebBrowserPersist.cpp

   1 /* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*-
   2  * This Source Code Form is subject to the terms of the Mozilla Public
   3  * License, v. 2.0. If a copy of the MPL was not distributed with this
   4  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
   5
   6 #include "mozilla/ArrayUtils.h"
   7 #include "mozilla/TextUtils.h"
   8
   9 #include "nspr.h"
  10
  11 #include "nsIFileStreams.h"  // New Necko file streams
  12 #include <algorithm>
  13
  14 #include "nsNetCID.h"
  15 #include "nsNetUtil.h"
  16 #include "nsIClassOfService.h"
  17 #include "nsIInterfaceRequestorUtils.h"
  18 #include "nsIPrivateBrowsingChannel.h"
  19 #include "nsComponentManagerUtils.h"
  20 #include "nsIStorageStream.h"
  21 #include "nsISeekableStream.h"
  22 #include "nsIHttpChannel.h"
  23 #include "nsIEncodedChannel.h"
  24 #include "nsIUploadChannel.h"
  25 #include "nsICacheInfoChannel.h"
  26 #include "nsIFileChannel.h"
  27 #include "nsEscape.h"
  28 #include "nsIStringEnumerator.h"
  29 #include "nsStreamUtils.h"
  30
  31 #include "nsCExternalHandlerService.h"
  32
  33 #include "nsIURL.h"
  34 #include "nsIFileURL.h"
  35 #include "nsIWebProgressListener.h"
  36 #include "nsIAuthPrompt.h"
  37 #include "nsIPrompt.h"
  38 #include "nsIThreadRetargetableRequest.h"
  39 #include "nsContentUtils.h"
  40
  41 #include "nsIStringBundle.h"
  42 #include "nsIProtocolHandler.h"
  43
  44 #include "nsWebBrowserPersist.h"
  45 #include "WebBrowserPersistLocalDocument.h"
  46
  47 #include "nsIContent.h"
  48 #include "nsIMIMEInfo.h"
  49 #include "mozilla/dom/Document.h"
  50 #include "mozilla/net/CookieJarSettings.h"
  51 #include "mozilla/Mutex.h"
  52 #include "mozilla/Printf.h"
  53 #include "ReferrerInfo.h"
  54 #include "nsIURIMutator.h"
  55 #include "mozilla/WebBrowserPersistDocumentParent.h"
  56 #include "mozilla/dom/CanonicalBrowsingContext.h"
  57 #include "mozilla/dom/WindowGlobalParent.h"
  58 #include "mozilla/dom/ContentParent.h"
  59 #include "mozilla/dom/PContentParent.h"
  60 #include "mozilla/dom/BrowserParent.h"
  61 #include "nsIDocumentEncoder.h"
  62
  63 using namespace mozilla;
  64 using namespace mozilla::dom;
  65
  66 // Buffer file writes in 32kb chunks
  67 #define BUFFERED_OUTPUT_SIZE (1024 * 32)
  68
  69 struct nsWebBrowserPersist::WalkData {
  70   nsCOMPtr<nsIWebBrowserPersistDocument> mDocument;
  71   nsCOMPtr<nsIURI> mFile;
  72   nsCOMPtr<nsIURI> mDataPath;
  73 };
  74
  75 // Information about a DOM document
  76 struct nsWebBrowserPersist::DocData {
  77   nsCOMPtr<nsIURI> mBaseURI;
  78   nsCOMPtr<nsIWebBrowserPersistDocument> mDocument;
  79   nsCOMPtr<nsIURI> mFile;
  80   nsCString mCharset;
  81 };
  82
  83 // Information about a URI
  84 struct nsWebBrowserPersist::URIData {
  85   bool mNeedsPersisting;
  86   bool mSaved;
  87   bool mIsSubFrame;
  88   bool mDataPathIsRelative;
  89   bool mNeedsFixup;
  90   nsString mFilename;
  91   nsString mSubFrameExt;
  92   nsCOMPtr<nsIURI> mFile;
  93   nsCOMPtr<nsIURI> mDataPath;
  94   nsCOMPtr<nsIURI> mRelativeDocumentURI;
  95   nsCOMPtr<nsIPrincipal> mTriggeringPrincipal;
  96   nsCOMPtr<nsICookieJarSettings> mCookieJarSettings;
  97   nsContentPolicyType mContentPolicyType;
  98   nsCString mRelativePathToData;
  99   nsCString mCharset;
 100
 101   nsresult GetLocalURI(nsIURI* targetBaseURI, nsCString& aSpecOut);
 102 };
 103
 104 // Information about the output stream
 105 // Note that this data structure (and the map that nsWebBrowserPersist keeps,
 106 // where these are values) is used from two threads: the main thread,
 107 // and the background task thread.
 108 // The background thread only writes to mStream (from OnDataAvailable), and
 109 // this access is guarded using mStreamMutex. It reads the mFile member, which
 110 // is only written to on the main thread when the object is constructed and
 111 // from OnStartRequest (if mCalcFileExt), both guaranteed to happen before
 112 // OnDataAvailable is fired.
 113 // The main thread gets OnStartRequest, OnStopRequest, and progress sink events,
 114 // and accesses the other members.
 115 struct nsWebBrowserPersist::OutputData {
 116   nsCOMPtr<nsIURI> mFile;
 117   nsCOMPtr<nsIURI> mOriginalLocation;
 118   nsCOMPtr<nsIOutputStream> mStream;
 119   Mutex mStreamMutex MOZ_UNANNOTATED;
 120   int64_t mSelfProgress;
 121   int64_t mSelfProgressMax;
 122   bool mCalcFileExt;
 123
 124   OutputData(nsIURI* aFile, nsIURI* aOriginalLocation, bool aCalcFileExt)
 125       : mFile(aFile),
 126         mOriginalLocation(aOriginalLocation),
 127         mStreamMutex("nsWebBrowserPersist::OutputData::mStreamMutex"),
 128         mSelfProgress(0),
 129         mSelfProgressMax(10000),
 130         mCalcFileExt(aCalcFileExt) {}
 131   ~OutputData() {
 132     // Gaining this lock in the destructor is pretty icky. It should be OK
 133     // because the only other place we lock the mutex is in OnDataAvailable,
 134     // which will never itself cause the OutputData instance to be
 135     // destroyed.
 136     MutexAutoLock lock(mStreamMutex);
 137     if (mStream) {
 138       mStream->Close();
 139     }
 140   }
 141 };
 142
 143 struct nsWebBrowserPersist::UploadData {
 144   nsCOMPtr<nsIURI> mFile;
 145   int64_t mSelfProgress;
 146   int64_t mSelfProgressMax;
 147
 148   explicit UploadData(nsIURI* aFile)
 149       : mFile(aFile), mSelfProgress(0), mSelfProgressMax(10000) {}
 150 };
 151
 152 struct nsWebBrowserPersist::CleanupData {
 153   nsCOMPtr<nsIFile> mFile;
 154   // Snapshot of what the file actually is at the time of creation so that if
 155   // it transmutes into something else later on it can be ignored. For example,
 156   // catch files that turn into dirs or vice versa.
 157   bool mIsDirectory;
 158 };
 159
 160 class nsWebBrowserPersist::OnWalk final
 161     : public nsIWebBrowserPersistResourceVisitor {
 162  public:
 163   OnWalk(nsWebBrowserPersist* aParent, nsIURI* aFile, nsIFile* aDataPath)
 164       : mParent(aParent),
 165         mFile(aFile),
 166         mDataPath(aDataPath),
 167         mPendingDocuments(1),
 168         mStatus(NS_OK) {}
 169
 170   NS_DECL_NSIWEBBROWSERPERSISTRESOURCEVISITOR
 171   NS_DECL_ISUPPORTS
 172  private:
 173   RefPtr<nsWebBrowserPersist> mParent;
 174   nsCOMPtr<nsIURI> mFile;
 175   nsCOMPtr<nsIFile> mDataPath;
 176
 177   uint32_t mPendingDocuments;
 178   nsresult mStatus;
 179
 180   virtual ~OnWalk() = default;
 181 };
 182
 183 NS_IMPL_ISUPPORTS(nsWebBrowserPersist::OnWalk,
 184                   nsIWebBrowserPersistResourceVisitor)
 185
 186 class nsWebBrowserPersist::OnRemoteWalk final
 187     : public nsIWebBrowserPersistDocumentReceiver {
 188  public:
 189   OnRemoteWalk(nsIWebBrowserPersistResourceVisitor* aVisitor,
 190                nsIWebBrowserPersistDocument* aDocument)
 191       : mVisitor(aVisitor), mDocument(aDocument) {}
 192
 193   NS_DECL_NSIWEBBROWSERPERSISTDOCUMENTRECEIVER
 194   NS_DECL_ISUPPORTS
 195  private:
 196   nsCOMPtr<nsIWebBrowserPersistResourceVisitor> mVisitor;
 197   nsCOMPtr<nsIWebBrowserPersistDocument> mDocument;
 198
 199   virtual ~OnRemoteWalk() = default;
 200 };
 201
 202 NS_IMPL_ISUPPORTS(nsWebBrowserPersist::OnRemoteWalk,
 203                   nsIWebBrowserPersistDocumentReceiver)
 204
 205 class nsWebBrowserPersist::OnWrite final
 206     : public nsIWebBrowserPersistWriteCompletion {
 207  public:
 208   OnWrite(nsWebBrowserPersist* aParent, nsIURI* aFile, nsIFile* aLocalFile)
 209       : mParent(aParent), mFile(aFile), mLocalFile(aLocalFile) {}
 210
 211   NS_DECL_NSIWEBBROWSERPERSISTWRITECOMPLETION
 212   NS_DECL_ISUPPORTS
 213  private:
 214   RefPtr<nsWebBrowserPersist> mParent;
 215   nsCOMPtr<nsIURI> mFile;
 216   nsCOMPtr<nsIFile> mLocalFile;
 217
 218   virtual ~OnWrite() = default;
 219 };
 220
 221 NS_IMPL_ISUPPORTS(nsWebBrowserPersist::OnWrite,
 222                   nsIWebBrowserPersistWriteCompletion)
 223
 224 class nsWebBrowserPersist::FlatURIMap final
 225     : public nsIWebBrowserPersistURIMap {
 226  public:
 227   explicit FlatURIMap(const nsACString& aTargetBase)
 228       : mTargetBase(aTargetBase) {}
 229
 230   void Add(const nsACString& aMapFrom, const nsACString& aMapTo) {
 231     mMapFrom.AppendElement(aMapFrom);
 232     mMapTo.AppendElement(aMapTo);
 233   }
 234
 235   NS_DECL_NSIWEBBROWSERPERSISTURIMAP
 236   NS_DECL_ISUPPORTS
 237
 238  private:
 239   nsTArray<nsCString> mMapFrom;
 240   nsTArray<nsCString> mMapTo;
 241   nsCString mTargetBase;
 242
 243   virtual ~FlatURIMap() = default;
 244 };
 245
 246 NS_IMPL_ISUPPORTS(nsWebBrowserPersist::FlatURIMap, nsIWebBrowserPersistURIMap)
 247
 248 NS_IMETHODIMP
 249 nsWebBrowserPersist::FlatURIMap::GetNumMappedURIs(uint32_t* aNum) {
 250   MOZ_ASSERT(mMapFrom.Length() == mMapTo.Length());
 251   *aNum = mMapTo.Length();
 252   return NS_OK;
 253 }
 254
 255 NS_IMETHODIMP
 256 nsWebBrowserPersist::FlatURIMap::GetTargetBaseURI(nsACString& aTargetBase) {
 257   aTargetBase = mTargetBase;
 258   return NS_OK;
 259 }
 260
 261 NS_IMETHODIMP
 262 nsWebBrowserPersist::FlatURIMap::GetURIMapping(uint32_t aIndex,
 263                                                nsACString& aMapFrom,
 264                                                nsACString& aMapTo) {
 265   MOZ_ASSERT(mMapFrom.Length() == mMapTo.Length());
 266   if (aIndex >= mMapTo.Length()) {
 267     return NS_ERROR_INVALID_ARG;
 268   }
 269   aMapFrom = mMapFrom[aIndex];
 270   aMapTo = mMapTo[aIndex];
 271   return NS_OK;
 272 }
 273
 274 // Maximum file length constant. The max file name length is
 275 // volume / server dependent but it is difficult to obtain
 276 // that information. Instead this constant is a reasonable value that
 277 // modern systems should able to cope with.
 278 const uint32_t kDefaultMaxFilenameLength = 64;
 279
 280 // Default flags for persistence
 281 const uint32_t kDefaultPersistFlags =
 282     nsIWebBrowserPersist::PERSIST_FLAGS_NO_CONVERSION |
 283     nsIWebBrowserPersist::PERSIST_FLAGS_REPLACE_EXISTING_FILES;
 284
 285 // String bundle where error messages come from
 286 const char* kWebBrowserPersistStringBundle =
 287     "chrome://global/locale/nsWebBrowserPersist.properties";
 288
 289 nsWebBrowserPersist::nsWebBrowserPersist()
 290     : mCurrentDataPathIsRelative(false),
 291       mCurrentThingsToPersist(0),
 292       mOutputMapMutex("nsWebBrowserPersist::mOutputMapMutex"),
 293       mFirstAndOnlyUse(true),
 294       mSavingDocument(false),
 295       mCancel(false),
 296       mEndCalled(false),
 297       mCompleted(false),
 298       mStartSaving(false),
 299       mReplaceExisting(true),
 300       mSerializingOutput(false),
 301       mIsPrivate(false),
 302       mPersistFlags(kDefaultPersistFlags),
 303       mPersistResult(NS_OK),
 304       mTotalCurrentProgress(0),
 305       mTotalMaxProgress(0),
 306       mWrapColumn(72),
 307       mEncodingFlags(0) {}
 308
 309 nsWebBrowserPersist::~nsWebBrowserPersist() { Cleanup(); }
 310
 311 //*****************************************************************************
 312 // nsWebBrowserPersist::nsISupports
 313 //*****************************************************************************
 314
 315 NS_IMPL_ADDREF(nsWebBrowserPersist)
 316 NS_IMPL_RELEASE(nsWebBrowserPersist)
 317
 318 NS_INTERFACE_MAP_BEGIN(nsWebBrowserPersist)
 319   NS_INTERFACE_MAP_ENTRY_AMBIGUOUS(nsISupports, nsIWebBrowserPersist)
 320   NS_INTERFACE_MAP_ENTRY(nsIWebBrowserPersist)
 321   NS_INTERFACE_MAP_ENTRY(nsICancelable)
 322   NS_INTERFACE_MAP_ENTRY(nsIInterfaceRequestor)
 323   NS_INTERFACE_MAP_ENTRY(nsISupportsWeakReference)
 324   NS_INTERFACE_MAP_ENTRY(nsIStreamListener)
 325   NS_INTERFACE_MAP_ENTRY(nsIThreadRetargetableStreamListener)
 326   NS_INTERFACE_MAP_ENTRY(nsIRequestObserver)
 327   NS_INTERFACE_MAP_ENTRY(nsIProgressEventSink)
 328 NS_INTERFACE_MAP_END
 329
 330 //*****************************************************************************
 331 // nsWebBrowserPersist::nsIInterfaceRequestor
 332 //*****************************************************************************
 333
 334 NS_IMETHODIMP nsWebBrowserPersist::GetInterface(const nsIID& aIID,
 335                                                 void** aIFace) {
 336   NS_ENSURE_ARG_POINTER(aIFace);
 337
 338   *aIFace = nullptr;
 339
 340   nsresult rv = QueryInterface(aIID, aIFace);
 341   if (NS_SUCCEEDED(rv)) {
 342     return rv;
 343   }
 344
 345   if (mProgressListener && (aIID.Equals(NS_GET_IID(nsIAuthPrompt)) ||
 346                             aIID.Equals(NS_GET_IID(nsIPrompt)))) {
 347     mProgressListener->QueryInterface(aIID, aIFace);
 348     if (*aIFace) return NS_OK;
 349   }
 350
 351   nsCOMPtr<nsIInterfaceRequestor> req = do_QueryInterface(mProgressListener);
 352   if (req) {
 353     return req->GetInterface(aIID, aIFace);
 354   }
 355
 356   return NS_ERROR_NO_INTERFACE;
 357 }
 358
 359 //*****************************************************************************
 360 // nsWebBrowserPersist::nsIWebBrowserPersist
 361 //*****************************************************************************
 362
 363 NS_IMETHODIMP nsWebBrowserPersist::GetPersistFlags(uint32_t* aPersistFlags) {
 364   NS_ENSURE_ARG_POINTER(aPersistFlags);
 365   *aPersistFlags = mPersistFlags;
 366   return NS_OK;
 367 }
 368 NS_IMETHODIMP nsWebBrowserPersist::SetPersistFlags(uint32_t aPersistFlags) {
 369   mPersistFlags = aPersistFlags;
 370   mReplaceExisting = (mPersistFlags & PERSIST_FLAGS_REPLACE_EXISTING_FILES);
 371   mSerializingOutput = (mPersistFlags & PERSIST_FLAGS_SERIALIZE_OUTPUT);
 372   return NS_OK;
 373 }
 374
 375 NS_IMETHODIMP nsWebBrowserPersist::GetCurrentState(uint32_t* aCurrentState) {
 376   NS_ENSURE_ARG_POINTER(aCurrentState);
 377   if (mCompleted) {
 378     *aCurrentState = PERSIST_STATE_FINISHED;
 379   } else if (mFirstAndOnlyUse) {
 380     *aCurrentState = PERSIST_STATE_SAVING;
 381   } else {
 382     *aCurrentState = PERSIST_STATE_READY;
 383   }
 384   return NS_OK;
 385 }
 386
 387 NS_IMETHODIMP nsWebBrowserPersist::GetResult(nsresult* aResult) {
 388   NS_ENSURE_ARG_POINTER(aResult);
 389   *aResult = mPersistResult;
 390   return NS_OK;
 391 }
 392
 393 NS_IMETHODIMP nsWebBrowserPersist::GetProgressListener(
 394     nsIWebProgressListener** aProgressListener) {
 395   NS_ENSURE_ARG_POINTER(aProgressListener);
 396   *aProgressListener = mProgressListener;
 397   NS_IF_ADDREF(*aProgressListener);
 398   return NS_OK;
 399 }
 400
 401 NS_IMETHODIMP nsWebBrowserPersist::SetProgressListener(
 402     nsIWebProgressListener* aProgressListener) {
 403   mProgressListener = aProgressListener;
 404   mProgressListener2 = do_QueryInterface(aProgressListener);
 405   mEventSink = do_GetInterface(aProgressListener);
 406   return NS_OK;
 407 }
 408
 409 NS_IMETHODIMP nsWebBrowserPersist::SaveURI(
 410     nsIURI* aURI, nsIPrincipal* aPrincipal, uint32_t aCacheKey,
 411     nsIReferrerInfo* aReferrerInfo, nsICookieJarSettings* aCookieJarSettings,
 412     nsIInputStream* aPostData, const char* aExtraHeaders, nsISupports* aFile,
 413     nsContentPolicyType aContentPolicy, bool aIsPrivate) {
 414   NS_ENSURE_TRUE(mFirstAndOnlyUse, NS_ERROR_FAILURE);
 415   mFirstAndOnlyUse = false;  // Stop people from reusing this object!
 416
 417   nsCOMPtr<nsIURI> fileAsURI;
 418   nsresult rv;
 419   rv = GetValidURIFromObject(aFile, getter_AddRefs(fileAsURI));
 420   NS_ENSURE_SUCCESS(rv, NS_ERROR_INVALID_ARG);
 421
 422   // SaveURIInternal doesn't like broken uris.
 423   mPersistFlags |= PERSIST_FLAGS_FAIL_ON_BROKEN_LINKS;
 424   rv = SaveURIInternal(aURI, aPrincipal, aContentPolicy, aCacheKey,
 425                        aReferrerInfo, aCookieJarSettings, aPostData,
 426                        aExtraHeaders, fileAsURI, false, aIsPrivate);
 427   return NS_FAILED(rv) ? rv : NS_OK;
 428 }
 429
 430 NS_IMETHODIMP nsWebBrowserPersist::SaveChannel(nsIChannel* aChannel,
 431                                                nsISupports* aFile) {
 432   NS_ENSURE_TRUE(mFirstAndOnlyUse, NS_ERROR_FAILURE);
 433   mFirstAndOnlyUse = false;  // Stop people from reusing this object!
 434
 435   nsCOMPtr<nsIURI> fileAsURI;
 436   nsresult rv;
 437   rv = GetValidURIFromObject(aFile, getter_AddRefs(fileAsURI));
 438   NS_ENSURE_SUCCESS(rv, NS_ERROR_INVALID_ARG);
 439
 440   rv = aChannel->GetURI(getter_AddRefs(mURI));
 441   NS_ENSURE_SUCCESS(rv, rv);
 442
 443   // SaveChannelInternal doesn't like broken uris.
 444   mPersistFlags |= PERSIST_FLAGS_FAIL_ON_BROKEN_LINKS;
 445   rv = SaveChannelInternal(aChannel, fileAsURI, false);
 446   return NS_FAILED(rv) ? rv : NS_OK;
 447 }
 448
 449 NS_IMETHODIMP nsWebBrowserPersist::SaveDocument(nsISupports* aDocument,
 450                                                 nsISupports* aFile,
 451                                                 nsISupports* aDataPath,
 452                                                 const char* aOutputContentType,
 453                                                 uint32_t aEncodingFlags,
 454                                                 uint32_t aWrapColumn) {
 455   NS_ENSURE_TRUE(mFirstAndOnlyUse, NS_ERROR_FAILURE);
 456   mFirstAndOnlyUse = false;  // Stop people from reusing this object!
 457
 458   // We need a STATE_IS_NETWORK start/stop pair to bracket the
 459   // notification callbacks.  For a whole document we generate those
 460   // here and in EndDownload(), but for the single-request methods
 461   // that's done in On{Start,Stop}Request instead.
 462   mSavingDocument = true;
 463
 464   NS_ENSURE_ARG_POINTER(aDocument);
 465   NS_ENSURE_ARG_POINTER(aFile);
 466
 467   nsCOMPtr<nsIURI> fileAsURI;
 468   nsCOMPtr<nsIURI> datapathAsURI;
 469   nsresult rv;
 470
 471   rv = GetValidURIFromObject(aFile, getter_AddRefs(fileAsURI));
 472   NS_ENSURE_SUCCESS(rv, NS_ERROR_INVALID_ARG);
 473   if (aDataPath) {
 474     rv = GetValidURIFromObject(aDataPath, getter_AddRefs(datapathAsURI));
 475     NS_ENSURE_SUCCESS(rv, NS_ERROR_INVALID_ARG);
 476   }
 477
 478   mWrapColumn = aWrapColumn;
 479   mEncodingFlags = aEncodingFlags;
 480
 481   if (aOutputContentType) {
 482     mContentType.AssignASCII(aOutputContentType);
 483   }
 484
 485   // State start notification
 486   if (mProgressListener) {
 487     mProgressListener->OnStateChange(
 488         nullptr, nullptr,
 489         nsIWebProgressListener::STATE_START |
 490             nsIWebProgressListener::STATE_IS_NETWORK,
 491         NS_OK);
 492   }
 493
 494   nsCOMPtr<nsIWebBrowserPersistDocument> doc = do_QueryInterface(aDocument);
 495   if (!doc) {
 496     nsCOMPtr<Document> localDoc = do_QueryInterface(aDocument);
 497     if (localDoc) {
 498       doc = new mozilla::WebBrowserPersistLocalDocument(localDoc);
 499     } else {
 500       rv = NS_ERROR_NO_INTERFACE;
 501     }
 502   }
 503
 504   bool closed = false;
 505   if (doc && NS_SUCCEEDED(doc->GetIsClosed(&closed)) && !closed) {
 506     rv = SaveDocumentInternal(doc, fileAsURI, datapathAsURI);
 507   }
 508
 509   if (NS_FAILED(rv) || closed) {
 510     SendErrorStatusChange(true, rv, nullptr, mURI);
 511     EndDownload(rv);
 512   }
 513   return rv;
 514 }
 515
 516 NS_IMETHODIMP nsWebBrowserPersist::Cancel(nsresult aReason) {
 517   // No point cancelling if we're already complete.
 518   if (mEndCalled) {
 519     return NS_OK;
 520   }
 521   mCancel = true;
 522   EndDownload(aReason);
 523   return NS_OK;
 524 }
 525
 526 NS_IMETHODIMP nsWebBrowserPersist::CancelSave() {
 527   return Cancel(NS_BINDING_ABORTED);
 528 }
 529
 530 nsresult nsWebBrowserPersist::StartUpload(nsIStorageStream* storStream,
 531                                           nsIURI* aDestinationURI,
 532                                           const nsACString& aContentType) {
 533   // setup the upload channel if the destination is not local
 534   nsCOMPtr<nsIInputStream> inputstream;
 535   nsresult rv = storStream->NewInputStream(0, getter_AddRefs(inputstream));
 536   NS_ENSURE_TRUE(inputstream, NS_ERROR_FAILURE);
 537   NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
 538   return StartUpload(inputstream, aDestinationURI, aContentType);
 539 }
 540
 541 nsresult nsWebBrowserPersist::StartUpload(nsIInputStream* aInputStream,
 542                                           nsIURI* aDestinationURI,
 543                                           const nsACString& aContentType) {
 544   nsCOMPtr<nsIChannel> destChannel;
 545   CreateChannelFromURI(aDestinationURI, getter_AddRefs(destChannel));
 546   nsCOMPtr<nsIUploadChannel> uploadChannel(do_QueryInterface(destChannel));
 547   NS_ENSURE_TRUE(uploadChannel, NS_ERROR_FAILURE);
 548
 549   // Set the upload stream
 550   // NOTE: ALL data must be available in "inputstream"
 551   nsresult rv = uploadChannel->SetUploadStream(aInputStream, aContentType, -1);
 552   NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
 553   rv = destChannel->AsyncOpen(this);
 554   NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
 555
 556   // add this to the upload list
 557   nsCOMPtr<nsISupports> keyPtr = do_QueryInterface(destChannel);
 558   mUploadList.InsertOrUpdate(keyPtr, MakeUnique<UploadData>(aDestinationURI));
 559
 560   return NS_OK;
 561 }
 562
 563 void nsWebBrowserPersist::SerializeNextFile() {
 564   nsresult rv = NS_OK;
 565   MOZ_ASSERT(mWalkStack.Length() == 0);
 566
 567   // First, handle gathered URIs.
 568   // This is potentially O(n^2), when taking into account the
 569   // number of times this method is called.  If it becomes a
 570   // bottleneck, the count of not-yet-persisted URIs could be
 571   // maintained separately, and we can skip iterating mURIMap if there are none.
 572
 573   // Persist each file in the uri map. The document(s)
 574   // will be saved after the last one of these is saved.
 575   for (const auto& entry : mURIMap) {
 576     URIData* data = entry.GetWeak();
 577
 578     if (!data->mNeedsPersisting || data->mSaved) {
 579       continue;
 580     }
 581
 582     // Create a URI from the key.
 583     nsCOMPtr<nsIURI> uri;
 584     rv = NS_NewURI(getter_AddRefs(uri), entry.GetKey(), data->mCharset.get());
 585     if (NS_WARN_IF(NS_FAILED(rv))) {
 586       break;
 587     }
 588
 589     // Make a URI to save the data to.
 590     nsCOMPtr<nsIURI> fileAsURI = data->mDataPath;
 591     rv = AppendPathToURI(fileAsURI, data->mFilename, fileAsURI);
 592     if (NS_WARN_IF(NS_FAILED(rv))) {
 593       break;
 594     }
 595
 596     rv = SaveURIInternal(uri, data->mTriggeringPrincipal,
 597                          data->mContentPolicyType, 0, nullptr,
 598                          data->mCookieJarSettings, nullptr, nullptr, fileAsURI,
 599                          true, mIsPrivate);
 600     // If SaveURIInternal fails, then it will have called EndDownload,
 601     // which means that |data| is no longer valid memory. We MUST bail.
 602     if (NS_WARN_IF(NS_FAILED(rv))) {
 603       break;
 604     }
 605
 606     if (rv == NS_OK) {
 607       // URIData.mFile will be updated to point to the correct
 608       // URI object when it is fixed up with the right file extension
 609       // in OnStartRequest
 610       data->mFile = fileAsURI;
 611       data->mSaved = true;
 612     } else {
 613       data->mNeedsFixup = false;
 614     }
 615
 616     if (mSerializingOutput) {
 617       break;
 618     }
 619   }
 620
 621   // If there are downloads happening, wait until they're done; the
 622   // OnStopRequest handler will call this method again.
 623   if (mOutputMap.Count() > 0) {
 624     return;
 625   }
 626
 627   // If serializing, also wait until last upload is done.
 628   if (mSerializingOutput && mUploadList.Count() > 0) {
 629     return;
 630   }
 631
 632   // If there are also no more documents, then we're done.
 633   if (mDocList.Length() == 0) {
 634     // ...or not quite done, if there are still uploads.
 635     if (mUploadList.Count() > 0) {
 636       return;
 637     }
 638     // Finish and clean things up.  Defer this because the caller
 639     // may have been expecting to use the listeners that that
 640     // method will clear.
 641     NS_DispatchToCurrentThread(
 642         NewRunnableMethod("nsWebBrowserPersist::FinishDownload", this,
 643                           &nsWebBrowserPersist::FinishDownload));
 644     return;
 645   }
 646
 647   // There are no URIs to save, so just save the next document.
 648   mStartSaving = true;
 649   mozilla::UniquePtr<DocData> docData(mDocList.ElementAt(0));
 650   mDocList.RemoveElementAt(0);  // O(n^2) but probably doesn't matter.
 651   MOZ_ASSERT(docData);
 652   if (!docData) {
 653     EndDownload(NS_ERROR_FAILURE);
 654     return;
 655   }
 656
 657   mCurrentBaseURI = docData->mBaseURI;
 658   mCurrentCharset = docData->mCharset;
 659   mTargetBaseURI = docData->mFile;
 660
 661   // Save the document, fixing it up with the new URIs as we do
 662
 663   nsAutoCString targetBaseSpec;
 664   if (mTargetBaseURI) {
 665     rv = mTargetBaseURI->GetSpec(targetBaseSpec);
 666     if (NS_FAILED(rv)) {
 667       SendErrorStatusChange(true, rv, nullptr, nullptr);
 668       EndDownload(rv);
 669       return;
 670     }
 671   }
 672
 673   // mFlatURIMap must be rebuilt each time through SerializeNextFile, as
 674   // mTargetBaseURI is used to create the relative URLs and will be different
 675   // with each serialized document.
 676   RefPtr<FlatURIMap> flatMap = new FlatURIMap(targetBaseSpec);
 677   for (const auto& uriEntry : mURIMap) {
 678     nsAutoCString mapTo;
 679     nsresult rv = uriEntry.GetWeak()->GetLocalURI(mTargetBaseURI, mapTo);
 680     if (NS_SUCCEEDED(rv) || !mapTo.IsVoid()) {
 681       flatMap->Add(uriEntry.GetKey(), mapTo);
 682     }
 683   }
 684   mFlatURIMap = std::move(flatMap);
 685
 686   nsCOMPtr<nsIFile> localFile;
 687   GetLocalFileFromURI(docData->mFile, getter_AddRefs(localFile));
 688   if (localFile) {
 689     // if we're not replacing an existing file but the file
 690     // exists, something is wrong
 691     bool fileExists = false;
 692     rv = localFile->Exists(&fileExists);
 693     if (NS_SUCCEEDED(rv) && !mReplaceExisting && fileExists) {
 694       rv = NS_ERROR_FILE_ALREADY_EXISTS;
 695     }
 696     if (NS_FAILED(rv)) {
 697       SendErrorStatusChange(false, rv, nullptr, docData->mFile);
 698       EndDownload(rv);
 699       return;
 700     }
 701   }
 702   nsCOMPtr<nsIOutputStream> outputStream;
 703   rv = MakeOutputStream(docData->mFile, getter_AddRefs(outputStream));
 704   if (NS_SUCCEEDED(rv) && !outputStream) {
 705     rv = NS_ERROR_FAILURE;
 706   }
 707   if (NS_FAILED(rv)) {
 708     SendErrorStatusChange(false, rv, nullptr, docData->mFile);
 709     EndDownload(rv);
 710     return;
 711   }
 712
 713   RefPtr<OnWrite> finish = new OnWrite(this, docData->mFile, localFile);
 714   rv = docData->mDocument->WriteContent(outputStream, mFlatURIMap,
 715                                         NS_ConvertUTF16toUTF8(mContentType),
 716                                         mEncodingFlags, mWrapColumn, finish);
 717   if (NS_FAILED(rv)) {
 718     SendErrorStatusChange(false, rv, nullptr, docData->mFile);
 719     EndDownload(rv);
 720   }
 721 }
 722
 723 NS_IMETHODIMP
 724 nsWebBrowserPersist::OnWrite::OnFinish(nsIWebBrowserPersistDocument* aDoc,
 725                                        nsIOutputStream* aStream,
 726                                        const nsACString& aContentType,
 727                                        nsresult aStatus) {
 728   nsresult rv = aStatus;
 729
 730   if (NS_FAILED(rv)) {
 731     mParent->SendErrorStatusChange(false, rv, nullptr, mFile);
 732     mParent->EndDownload(rv);
 733     return NS_OK;
 734   }
 735   if (!mLocalFile) {
 736     nsCOMPtr<nsIStorageStream> storStream(do_QueryInterface(aStream));
 737     if (storStream) {
 738       aStream->Close();
 739       rv = mParent->StartUpload(storStream, mFile, aContentType);
 740       if (NS_FAILED(rv)) {
 741         mParent->SendErrorStatusChange(false, rv, nullptr, mFile);
 742         mParent->EndDownload(rv);
 743       }
 744       // Either we failed and we're done, or we're uploading and
 745       // the OnStopRequest callback is responsible for the next
 746       // SerializeNextFile().
 747       return NS_OK;
 748     }
 749   }
 750   NS_DispatchToCurrentThread(
 751       NewRunnableMethod("nsWebBrowserPersist::SerializeNextFile", mParent,
 752                         &nsWebBrowserPersist::SerializeNextFile));
 753   return NS_OK;
 754 }
 755
 756 //*****************************************************************************
 757 // nsWebBrowserPersist::nsIRequestObserver
 758 //*****************************************************************************
 759
 760 NS_IMETHODIMP nsWebBrowserPersist::OnStartRequest(nsIRequest* request) {
 761   if (mProgressListener) {
 762     uint32_t stateFlags = nsIWebProgressListener::STATE_START |
 763                           nsIWebProgressListener::STATE_IS_REQUEST;
 764     if (!mSavingDocument) {
 765       stateFlags |= nsIWebProgressListener::STATE_IS_NETWORK;
 766     }
 767     mProgressListener->OnStateChange(nullptr, request, stateFlags, NS_OK);
 768   }
 769
 770   nsCOMPtr<nsIChannel> channel = do_QueryInterface(request);
 771   NS_ENSURE_TRUE(channel, NS_ERROR_FAILURE);
 772
 773   nsCOMPtr<nsISupports> keyPtr = do_QueryInterface(request);
 774   OutputData* data = mOutputMap.Get(keyPtr);
 775
 776   // NOTE: This code uses the channel as a hash key so it will not
 777   //       recognize redirected channels because the key is not the same.
 778   //       When that happens we remove and add the data entry to use the
 779   //       new channel as the hash key.
 780   if (!data) {
 781     UploadData* upData = mUploadList.Get(keyPtr);
 782     if (!upData) {
 783       // Redirect? Try and fixup the output table
 784       nsresult rv = FixRedirectedChannelEntry(channel);
 785       NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
 786
 787       // Should be able to find the data after fixup unless redirects
 788       // are disabled.
 789       data = mOutputMap.Get(keyPtr);
 790       if (!data) {
 791         return NS_ERROR_FAILURE;
 792       }
 793     }
 794   }
 795
 796   if (data && data->mFile) {
 797     nsCOMPtr<nsIThreadRetargetableRequest> r = do_QueryInterface(request);
 798     // Determine if we're uploading. Only use OMT onDataAvailable if not.
 799     nsCOMPtr<nsIFile> localFile;
 800     GetLocalFileFromURI(data->mFile, getter_AddRefs(localFile));
 801     if (r && localFile) {
 802       if (!mBackgroundQueue) {
 803         NS_CreateBackgroundTaskQueue("WebBrowserPersist",
 804                                      getter_AddRefs(mBackgroundQueue));
 805       }
 806       if (mBackgroundQueue) {
 807         r->RetargetDeliveryTo(mBackgroundQueue);
 808       }
 809     }
 810
 811     // If PERSIST_FLAGS_AUTODETECT_APPLY_CONVERSION is set in mPersistFlags,
 812     // try to determine whether this channel needs to apply Content-Encoding
 813     // conversions.
 814     NS_ASSERTION(
 815         !((mPersistFlags & PERSIST_FLAGS_AUTODETECT_APPLY_CONVERSION) &&
 816           (mPersistFlags & PERSIST_FLAGS_NO_CONVERSION)),
 817         "Conflict in persist flags: both AUTODETECT and NO_CONVERSION set");
 818     if (mPersistFlags & PERSIST_FLAGS_AUTODETECT_APPLY_CONVERSION)
 819       SetApplyConversionIfNeeded(channel);
 820
 821     if (data->mCalcFileExt &&
 822         !(mPersistFlags & PERSIST_FLAGS_DONT_CHANGE_FILENAMES)) {
 823       nsCOMPtr<nsIURI> uriWithExt;
 824       // this is the first point at which the server can tell us the mimetype
 825       nsresult rv = CalculateAndAppendFileExt(
 826           data->mFile, channel, data->mOriginalLocation, uriWithExt);
 827       if (NS_SUCCEEDED(rv)) {
 828         data->mFile = uriWithExt;
 829       }
 830
 831       // now make filename conformant and unique
 832       nsCOMPtr<nsIURI> uniqueFilenameURI;
 833       rv = CalculateUniqueFilename(data->mFile, uniqueFilenameURI);
 834       if (NS_SUCCEEDED(rv)) {
 835         data->mFile = uniqueFilenameURI;
 836       }
 837
 838       // The URIData entry is pointing to the old unfixed URI, so we need
 839       // to update it.
 840       nsCOMPtr<nsIURI> chanURI;
 841       rv = channel->GetOriginalURI(getter_AddRefs(chanURI));
 842       if (NS_SUCCEEDED(rv)) {
 843         nsAutoCString spec;
 844         chanURI->GetSpec(spec);
 845         URIData* uridata;
 846         if (mURIMap.Get(spec, &uridata)) {
 847           uridata->mFile = data->mFile;
 848         }
 849       }
 850     }
 851
 852     // compare uris and bail before we add to output map if they are equal
 853     bool isEqual = false;
 854     if (NS_SUCCEEDED(data->mFile->Equals(data->mOriginalLocation, &isEqual)) &&
 855         isEqual) {
 856       {
 857         MutexAutoLock lock(mOutputMapMutex);
 858         // remove from output map
 859         mOutputMap.Remove(keyPtr);
 860       }
 861
 862       // cancel; we don't need to know any more
 863       // stop request will get called
 864       request->Cancel(NS_BINDING_ABORTED);
 865     }
 866   }
 867
 868   return NS_OK;
 869 }
 870
 871 NS_IMETHODIMP nsWebBrowserPersist::OnStopRequest(nsIRequest* request,
 872                                                  nsresult status) {
 873   nsCOMPtr<nsISupports> keyPtr = do_QueryInterface(request);
 874   OutputData* data = mOutputMap.Get(keyPtr);
 875   if (data) {
 876     if (NS_SUCCEEDED(mPersistResult) && NS_FAILED(status)) {
 877       SendErrorStatusChange(true, status, request, data->mFile);
 878     }
 879
 880     // If there is a stream ref and we weren't canceled,
 881     // close it away from the main thread.
 882     // We don't do this when there's an error/cancelation,
 883     // because our consumer may try to delete the file, which will error
 884     // if we're still holding on to it, so we have to close it pronto.
 885     {
 886       MutexAutoLock lock(data->mStreamMutex);
 887       if (data->mStream && NS_SUCCEEDED(status) && !mCancel) {
 888         if (!mBackgroundQueue) {
 889           nsresult rv = NS_CreateBackgroundTaskQueue(
 890               "WebBrowserPersist", getter_AddRefs(mBackgroundQueue));
 891           if (NS_FAILED(rv)) {
 892             return rv;
 893           }
 894         }
 895         // Now steal the stream ref and close it away from the main thread,
 896         // keeping the promise around so we don't finish before all files
 897         // are flushed and closed.
 898         mFileClosePromises.AppendElement(InvokeAsync(
 899             mBackgroundQueue, __func__, [stream = std::move(data->mStream)]() {
 900               nsresult rv = stream->Close();
 901               // We don't care if closing failed; we don't care in the
 902               // destructor either...
 903               return ClosePromise::CreateAndResolve(rv, __func__);
 904             }));
 905       }
 906     }
 907     MutexAutoLock lock(mOutputMapMutex);
 908     mOutputMap.Remove(keyPtr);
 909   } else {
 910     // if we didn't find the data in mOutputMap, try mUploadList
 911     UploadData* upData = mUploadList.Get(keyPtr);
 912     if (upData) {
 913       mUploadList.Remove(keyPtr);
 914     }
 915   }
 916
 917   // Do more work.
 918   SerializeNextFile();
 919
 920   if (mProgressListener) {
 921     uint32_t stateFlags = nsIWebProgressListener::STATE_STOP |
 922                           nsIWebProgressListener::STATE_IS_REQUEST;
 923     if (!mSavingDocument) {
 924       stateFlags |= nsIWebProgressListener::STATE_IS_NETWORK;
 925     }
 926     mProgressListener->OnStateChange(nullptr, request, stateFlags, status);
 927   }
 928
 929   return NS_OK;
 930 }
 931
 932 //*****************************************************************************
 933 // nsWebBrowserPersist::nsIStreamListener
 934 //*****************************************************************************
 935
 936 // Note: this is supposed to (but not guaranteed to) fire on a background
 937 // thread when used to save to local disk (channels not using local files will
 938 // use the main thread).
 939 // (Read) Access to mOutputMap is guarded via mOutputMapMutex.
 940 // Access to individual OutputData::mStream is guarded via its mStreamMutex.
 941 // mCancel is atomic, as is mPersistFlags (accessed via MakeOutputStream).
 942 // If you end up touching this method and needing other member access, bear
 943 // this in mind.
 944 NS_IMETHODIMP
 945 nsWebBrowserPersist::OnDataAvailable(nsIRequest* request,
 946                                      nsIInputStream* aIStream, uint64_t aOffset,
 947                                      uint32_t aLength) {
 948   // MOZ_ASSERT(!NS_IsMainThread()); // no guarantees, but it's likely.
 949
 950   bool cancel = mCancel;
 951   if (!cancel) {
 952     nsresult rv = NS_OK;
 953     uint32_t bytesRemaining = aLength;
 954
 955     nsCOMPtr<nsIChannel> channel = do_QueryInterface(request);
 956     NS_ENSURE_TRUE(channel, NS_ERROR_FAILURE);
 957
 958     MutexAutoLock lock(mOutputMapMutex);
 959     nsCOMPtr<nsISupports> keyPtr = do_QueryInterface(request);
 960     OutputData* data = mOutputMap.Get(keyPtr);
 961     if (!data) {
 962       // might be uploadData; consume necko's buffer and bail...
 963       uint32_t n;
 964       return aIStream->ReadSegments(NS_DiscardSegment, nullptr, aLength, &n);
 965     }
 966
 967     bool readError = true;
 968
 969     MutexAutoLock streamLock(data->mStreamMutex);
 970     // Make the output stream
 971     if (!data->mStream) {
 972       rv = MakeOutputStream(data->mFile, getter_AddRefs(data->mStream));
 973       if (NS_FAILED(rv)) {
 974         readError = false;
 975         cancel = true;
 976       }
 977     }
 978
 979     // Read data from the input and write to the output
 980     char buffer[8192];
 981     uint32_t bytesRead;
 982     while (!cancel && bytesRemaining) {
 983       readError = true;
 984       rv = aIStream->Read(buffer,
 985                           std::min(uint32_t(sizeof(buffer)), bytesRemaining),
 986                           &bytesRead);
 987       if (NS_SUCCEEDED(rv)) {
 988         readError = false;
 989         // Write out the data until something goes wrong, or, it is
 990         // all written.  We loop because for some errors (e.g., disk
 991         // full), we get NS_OK with some bytes written, then an error.
 992         // So, we want to write again in that case to get the actual
 993         // error code.
 994         const char* bufPtr = buffer;  // Where to write from.
 995         while (NS_SUCCEEDED(rv) && bytesRead) {
 996           uint32_t bytesWritten = 0;
 997           rv = data->mStream->Write(bufPtr, bytesRead, &bytesWritten);
 998           if (NS_SUCCEEDED(rv)) {
 999             bytesRead -= bytesWritten;
1000             bufPtr += bytesWritten;
1001             bytesRemaining -= bytesWritten;
1002             // Force an error if (for some reason) we get NS_OK but
1003             // no bytes written.
1004             if (!bytesWritten) {
1005               rv = NS_ERROR_FAILURE;
1006               cancel = true;
1007             }
1008           } else {
1009             // Disaster - can't write out the bytes - disk full / permission?
1010             cancel = true;
1011           }
1012         }
1013       } else {
1014         // Disaster - can't read the bytes - broken link / file error?
1015         cancel = true;
1016       }
1017     }
1018
1019     int64_t channelContentLength = -1;
1020     if (!cancel &&
1021         NS_SUCCEEDED(channel->GetContentLength(&channelContentLength))) {
1022       // if we get -1 at this point, we didn't get content-length header
1023       // assume that we got all of the data and push what we have;
1024       // that's the best we can do now
1025       if ((-1 == channelContentLength) ||
1026           ((channelContentLength - (aOffset + aLength)) == 0)) {
1027         NS_WARNING_ASSERTION(
1028             channelContentLength != -1,
1029             "nsWebBrowserPersist::OnDataAvailable() no content length "
1030             "header, pushing what we have");
1031         // we're done with this pass; see if we need to do upload
1032         nsAutoCString contentType;
1033         channel->GetContentType(contentType);
1034         // if we don't have the right type of output stream then it's a local
1035         // file
1036         nsCOMPtr<nsIStorageStream> storStream(do_QueryInterface(data->mStream));
1037         if (storStream) {
1038           data->mStream->Close();
1039           data->mStream =
1040               nullptr;  // null out stream so we don't close it later
1041           MOZ_ASSERT(NS_IsMainThread(),
1042                      "Uploads should be on the main thread.");
1043           rv = StartUpload(storStream, data->mFile, contentType);
1044           if (NS_FAILED(rv)) {
1045             readError = false;
1046             cancel = true;
1047           }
1048         }
1049       }
1050     }
1051
1052     // Notify listener if an error occurred.
1053     if (cancel) {
1054       RefPtr<nsIRequest> req = readError ? request : nullptr;
1055       nsCOMPtr<nsIURI> file = data->mFile;
1056       RefPtr<Runnable> errorOnMainThread = NS_NewRunnableFunction(
1057           "nsWebBrowserPersist::SendErrorStatusChange",
1058           [self = RefPtr{this}, req, file, readError, rv]() {
1059             self->SendErrorStatusChange(readError, rv, req, file);
1060           });
1061       NS_DispatchToMainThread(errorOnMainThread);
1062
1063       // And end the download on the main thread.
1064       nsCOMPtr<nsIRunnable> endOnMainThread = NewRunnableMethod<nsresult>(
1065           "nsWebBrowserPersist::EndDownload", this,
1066           &nsWebBrowserPersist::EndDownload, NS_BINDING_ABORTED);
1067       NS_DispatchToMainThread(endOnMainThread);
1068     }
1069   }
1070
1071   return cancel ? NS_BINDING_ABORTED : NS_OK;
1072 }
1073
1074 //*****************************************************************************
1075 // nsWebBrowserPersist::nsIThreadRetargetableStreamListener
1076 //*****************************************************************************
1077
1078 NS_IMETHODIMP nsWebBrowserPersist::CheckListenerChain() { return NS_OK; }
1079
1080 NS_IMETHODIMP
1081 nsWebBrowserPersist::OnDataFinished(nsresult) { return NS_OK; }
1082
1083 //*****************************************************************************
1084 // nsWebBrowserPersist::nsIProgressEventSink
1085 //*****************************************************************************
1086
1087 NS_IMETHODIMP nsWebBrowserPersist::OnProgress(nsIRequest* request,
1088                                               int64_t aProgress,
1089                                               int64_t aProgressMax) {
1090   if (!mProgressListener) {
1091     return NS_OK;
1092   }
1093
1094   // Store the progress of this request
1095   nsCOMPtr<nsISupports> keyPtr = do_QueryInterface(request);
1096   OutputData* data = mOutputMap.Get(keyPtr);
1097   if (data) {
1098     data->mSelfProgress = aProgress;
1099     data->mSelfProgressMax = aProgressMax;
1100   } else {
1101     UploadData* upData = mUploadList.Get(keyPtr);
1102     if (upData) {
1103       upData->mSelfProgress = aProgress;
1104       upData->mSelfProgressMax = aProgressMax;
1105     }
1106   }
1107
1108   // Notify listener of total progress
1109   CalcTotalProgress();
1110   if (mProgressListener2) {
1111     mProgressListener2->OnProgressChange64(nullptr, request, aProgress,
1112                                            aProgressMax, mTotalCurrentProgress,
1113                                            mTotalMaxProgress);
1114   } else {
1115     // have to truncate 64-bit to 32bit
1116     mProgressListener->OnProgressChange(
1117         nullptr, request, uint64_t(aProgress), uint64_t(aProgressMax),
1118         mTotalCurrentProgress, mTotalMaxProgress);
1119   }
1120
1121   // If our progress listener implements nsIProgressEventSink,
1122   // forward the notification
1123   if (mEventSink) {
1124     mEventSink->OnProgress(request, aProgress, aProgressMax);
1125   }
1126
1127   return NS_OK;
1128 }
1129
1130 NS_IMETHODIMP nsWebBrowserPersist::OnStatus(nsIRequest* request,
1131                                             nsresult status,
1132                                             const char16_t* statusArg) {
1133   if (mProgressListener) {
1134     // We need to filter out non-error error codes.
1135     // Is the only NS_SUCCEEDED value NS_OK?
1136     switch (status) {
1137       case NS_NET_STATUS_RESOLVING_HOST:
1138       case NS_NET_STATUS_RESOLVED_HOST:
1139       case NS_NET_STATUS_CONNECTING_TO:
1140       case NS_NET_STATUS_CONNECTED_TO:
1141       case NS_NET_STATUS_TLS_HANDSHAKE_STARTING:
1142       case NS_NET_STATUS_TLS_HANDSHAKE_ENDED:
1143       case NS_NET_STATUS_SENDING_TO:
1144       case NS_NET_STATUS_RECEIVING_FROM:
1145       case NS_NET_STATUS_WAITING_FOR:
1146       case NS_NET_STATUS_READING:
1147       case NS_NET_STATUS_WRITING:
1148         break;
1149
1150       default:
1151         // Pass other notifications (for legitimate errors) along.
1152         mProgressListener->OnStatusChange(nullptr, request, status, statusArg);
1153         break;
1154     }
1155   }
1156
1157   // If our progress listener implements nsIProgressEventSink,
1158   // forward the notification
1159   if (mEventSink) {
1160     mEventSink->OnStatus(request, status, statusArg);
1161   }
1162
1163   return NS_OK;
1164 }
1165
1166 //*****************************************************************************
1167 // nsWebBrowserPersist private methods
1168 //*****************************************************************************
1169
1170 // Convert error info into proper message text and send OnStatusChange
1171 // notification to the web progress listener.
1172 nsresult nsWebBrowserPersist::SendErrorStatusChange(bool aIsReadError,
1173                                                     nsresult aResult,
1174                                                     nsIRequest* aRequest,
1175                                                     nsIURI* aURI) {
1176   NS_ENSURE_ARG_POINTER(aURI);
1177
1178   if (!mProgressListener) {
1179     // Do nothing
1180     return NS_OK;
1181   }
1182
1183   // Get the file path or spec from the supplied URI
1184   nsCOMPtr<nsIFile> file;
1185   GetLocalFileFromURI(aURI, getter_AddRefs(file));
1186   AutoTArray<nsString, 1> strings;
1187   nsresult rv;
1188   if (file) {
1189     file->GetPath(*strings.AppendElement());
1190   } else {
1191     nsAutoCString fileurl;
1192     rv = aURI->GetSpec(fileurl);
1193     NS_ENSURE_SUCCESS(rv, rv);
1194     CopyUTF8toUTF16(fileurl, *strings.AppendElement());
1195   }
1196
1197   const char* msgId;
1198   switch (aResult) {
1199     case NS_ERROR_FILE_NAME_TOO_LONG:
1200       // File name too long.
1201       msgId = "fileNameTooLongError";
1202       break;
1203     case NS_ERROR_FILE_ALREADY_EXISTS:
1204       // File exists with same name as directory.
1205       msgId = "fileAlreadyExistsError";
1206       break;
1207     case NS_ERROR_FILE_NO_DEVICE_SPACE:
1208       // Out of space on target volume.
1209       msgId = "diskFull";
1210       break;
1211
1212     case NS_ERROR_FILE_READ_ONLY:
1213       // Attempt to write to read/only file.
1214       msgId = "readOnly";
1215       break;
1216
1217     case NS_ERROR_FILE_ACCESS_DENIED:
1218       // Attempt to write without sufficient permissions.
1219       msgId = "accessError";
1220       break;
1221
1222     default:
1223       // Generic read/write error message.
1224       if (aIsReadError)
1225         msgId = "readError";
1226       else
1227         msgId = "writeError";
1228       break;
1229   }
1230   // Get properties file bundle and extract status string.
1231   nsCOMPtr<nsIStringBundleService> s =
1232       do_GetService(NS_STRINGBUNDLE_CONTRACTID, &rv);
1233   NS_ENSURE_TRUE(NS_SUCCEEDED(rv) && s, NS_ERROR_FAILURE);
1234
1235   nsCOMPtr<nsIStringBundle> bundle;
1236   rv = s->CreateBundle(kWebBrowserPersistStringBundle, getter_AddRefs(bundle));
1237   NS_ENSURE_TRUE(NS_SUCCEEDED(rv) && bundle, NS_ERROR_FAILURE);
1238
1239   nsAutoString msgText;
1240   rv = bundle->FormatStringFromName(msgId, strings, msgText);
1241   NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
1242
1243   mProgressListener->OnStatusChange(nullptr, aRequest, aResult, msgText.get());
1244
1245   return NS_OK;
1246 }
1247
1248 nsresult nsWebBrowserPersist::GetValidURIFromObject(nsISupports* aObject,
1249                                                     nsIURI** aURI) const {
1250   NS_ENSURE_ARG_POINTER(aObject);
1251   NS_ENSURE_ARG_POINTER(aURI);
1252
1253   nsCOMPtr<nsIFile> objAsFile = do_QueryInterface(aObject);
1254   if (objAsFile) {
1255     return NS_NewFileURI(aURI, objAsFile);
1256   }
1257   nsCOMPtr<nsIURI> objAsURI = do_QueryInterface(aObject);
1258   if (objAsURI) {
1259     *aURI = objAsURI;
1260     NS_ADDREF(*aURI);
1261     return NS_OK;
1262   }
1263
1264   return NS_ERROR_FAILURE;
1265 }
1266
1267 /* static */
1268 nsresult nsWebBrowserPersist::GetLocalFileFromURI(nsIURI* aURI,
1269                                                   nsIFile** aLocalFile) {
1270   nsresult rv;
1271
1272   nsCOMPtr<nsIFileURL> fileURL = do_QueryInterface(aURI, &rv);
1273   if (NS_FAILED(rv)) return rv;
1274
1275   nsCOMPtr<nsIFile> file;
1276   rv = fileURL->GetFile(getter_AddRefs(file));
1277   if (NS_FAILED(rv)) {
1278     return rv;
1279   }
1280
1281   file.forget(aLocalFile);
1282   return NS_OK;
1283 }
1284
1285 /* static */
1286 nsresult nsWebBrowserPersist::AppendPathToURI(nsIURI* aURI,
1287                                               const nsAString& aPath,
1288                                               nsCOMPtr<nsIURI>& aOutURI) {
1289   NS_ENSURE_ARG_POINTER(aURI);
1290
1291   nsAutoCString newPath;
1292   nsresult rv = aURI->GetPathQueryRef(newPath);
1293   NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
1294
1295   // Append a forward slash if necessary
1296   int32_t len = newPath.Length();
1297   if (len > 0 && newPath.CharAt(len - 1) != '/') {
1298     newPath.Append('/');
1299   }
1300
1301   // Store the path back on the URI
1302   AppendUTF16toUTF8(aPath, newPath);
1303
1304   return NS_MutateURI(aURI).SetPathQueryRef(newPath).Finalize(aOutURI);
1305 }
1306
1307 nsresult nsWebBrowserPersist::SaveURIInternal(
1308     nsIURI* aURI, nsIPrincipal* aTriggeringPrincipal,
1309     nsContentPolicyType aContentPolicyType, uint32_t aCacheKey,
1310     nsIReferrerInfo* aReferrerInfo, nsICookieJarSettings* aCookieJarSettings,
1311     nsIInputStream* aPostData, const char* aExtraHeaders, nsIURI* aFile,
1312     bool aCalcFileExt, bool aIsPrivate) {
1313   NS_ENSURE_ARG_POINTER(aURI);
1314   NS_ENSURE_ARG_POINTER(aFile);
1315   NS_ENSURE_ARG_POINTER(aTriggeringPrincipal);
1316
1317   nsresult rv = NS_OK;
1318
1319   mURI = aURI;
1320
1321   nsLoadFlags loadFlags = nsIRequest::LOAD_NORMAL;
1322   if (mPersistFlags & PERSIST_FLAGS_BYPASS_CACHE) {
1323     loadFlags |= nsIRequest::LOAD_BYPASS_CACHE;
1324   } else if (mPersistFlags & PERSIST_FLAGS_FROM_CACHE) {
1325     loadFlags |= nsIRequest::LOAD_FROM_CACHE;
1326   }
1327
1328   // If there is no cookieJarSetting given, we need to create a new
1329   // cookieJarSettings for this download in order to send cookies based on the
1330   // current state of the prefs/permissions.
1331   nsCOMPtr<nsICookieJarSettings> cookieJarSettings = aCookieJarSettings;
1332   if (!cookieJarSettings) {
1333     // Although the variable is called 'triggering principal', it is used as the
1334     // loading principal in the download channel, so we treat it as a loading
1335     // principal also.
1336     bool shouldResistFingerprinting =
1337         nsContentUtils::ShouldResistFingerprinting_dangerous(
1338             aTriggeringPrincipal,
1339             "We are creating a new CookieJar Settings, so none exists "
1340             "currently. Although the variable is called 'triggering principal',"
1341             "it is used as the loading principal in the download channel, so we"
1342             "treat it as a loading principal also.",
1343             RFPTarget::IsAlwaysEnabledForPrecompute);
1344     cookieJarSettings =
1345         aIsPrivate
1346             ? net::CookieJarSettings::Create(net::CookieJarSettings::ePrivate,
1347                                              shouldResistFingerprinting)
1348             : net::CookieJarSettings::Create(net::CookieJarSettings::eRegular,
1349                                              shouldResistFingerprinting);
1350   }
1351
1352   // Open a channel to the URI
1353   nsCOMPtr<nsIChannel> inputChannel;
1354   rv = NS_NewChannel(getter_AddRefs(inputChannel), aURI, aTriggeringPrincipal,
1355                      nsILoadInfo::SEC_ALLOW_CROSS_ORIGIN_SEC_CONTEXT_IS_NULL,
1356                      aContentPolicyType, cookieJarSettings,
1357                      nullptr,  // aPerformanceStorage
1358                      nullptr,  // aLoadGroup
1359                      static_cast<nsIInterfaceRequestor*>(this), loadFlags);
1360
1361   nsCOMPtr<nsIPrivateBrowsingChannel> pbChannel =
1362       do_QueryInterface(inputChannel);
1363   if (pbChannel) {
1364     pbChannel->SetPrivate(aIsPrivate);
1365   }
1366
1367   if (NS_FAILED(rv) || inputChannel == nullptr) {
1368     EndDownload(NS_ERROR_FAILURE);
1369     return NS_ERROR_FAILURE;
1370   }
1371
1372   // Disable content conversion
1373   if (mPersistFlags & PERSIST_FLAGS_NO_CONVERSION) {
1374     nsCOMPtr<nsIEncodedChannel> encodedChannel(do_QueryInterface(inputChannel));
1375     if (encodedChannel) {
1376       encodedChannel->SetApplyConversion(false);
1377     }
1378   }
1379
1380   nsCOMPtr<nsILoadInfo> loadInfo = inputChannel->LoadInfo();
1381   loadInfo->SetIsUserTriggeredSave(true);
1382
1383   // Set the referrer, post data and headers if any
1384   nsCOMPtr<nsIHttpChannel> httpChannel(do_QueryInterface(inputChannel));
1385   if (httpChannel) {
1386     if (aReferrerInfo) {
1387       DebugOnly<nsresult> success = httpChannel->SetReferrerInfo(aReferrerInfo);
1388       MOZ_ASSERT(NS_SUCCEEDED(success));
1389     }
1390
1391     // Post data
1392     if (aPostData) {
1393       nsCOMPtr<nsISeekableStream> stream(do_QueryInterface(aPostData));
1394       if (stream) {
1395         // Rewind the postdata stream
1396         stream->Seek(nsISeekableStream::NS_SEEK_SET, 0);
1397         nsCOMPtr<nsIUploadChannel> uploadChannel(
1398             do_QueryInterface(httpChannel));
1399         NS_ASSERTION(uploadChannel, "http must support nsIUploadChannel");
1400         // Attach the postdata to the http channel
1401         uploadChannel->SetUploadStream(aPostData, ""_ns, -1);
1402       }
1403     }
1404
1405     // Cache key
1406     nsCOMPtr<nsICacheInfoChannel> cacheChannel(do_QueryInterface(httpChannel));
1407     if (cacheChannel && aCacheKey != 0) {
1408       cacheChannel->SetCacheKey(aCacheKey);
1409     }
1410
1411     // Headers
1412     if (aExtraHeaders) {
1413       nsAutoCString oneHeader;
1414       nsAutoCString headerName;
1415       nsAutoCString headerValue;
1416       int32_t crlf = 0;
1417       int32_t colon = 0;
1418       const char* kWhitespace = "\b\t\r\n ";
1419       nsAutoCString extraHeaders(aExtraHeaders);
1420       while (true) {
1421         crlf = extraHeaders.Find("\r\n");
1422         if (crlf == -1) break;
1423         extraHeaders.Mid(oneHeader, 0, crlf);
1424         extraHeaders.Cut(0, crlf + 2);
1425         colon = oneHeader.Find(":");
1426         if (colon == -1) break;  // Should have a colon
1427         oneHeader.Left(headerName, colon);
1428         colon++;
1429         oneHeader.Mid(headerValue, colon, oneHeader.Length() - colon);
1430         headerName.Trim(kWhitespace);
1431         headerValue.Trim(kWhitespace);
1432         // Add the header (merging if required)
1433         rv = httpChannel->SetRequestHeader(headerName, headerValue, true);
1434         if (NS_FAILED(rv)) {
1435           EndDownload(NS_ERROR_FAILURE);
1436           return NS_ERROR_FAILURE;
1437         }
1438       }
1439     }
1440   }
1441   return SaveChannelInternal(inputChannel, aFile, aCalcFileExt);
1442 }
1443
1444 nsresult nsWebBrowserPersist::SaveChannelInternal(nsIChannel* aChannel,
1445                                                   nsIURI* aFile,
1446                                                   bool aCalcFileExt) {
1447   NS_ENSURE_ARG_POINTER(aChannel);
1448   NS_ENSURE_ARG_POINTER(aFile);
1449
1450   // The default behaviour of SaveChannelInternal is to download the source
1451   // into a storage stream and upload that to the target. MakeOutputStream
1452   // special-cases a file target and creates a file output stream directly.
1453   // We want to special-case a file source and create a file input stream,
1454   // but we don't need to do this in the case of a file target.
1455   nsCOMPtr<nsIFileChannel> fc(do_QueryInterface(aChannel));
1456   nsCOMPtr<nsIFileURL> fu(do_QueryInterface(aFile));
1457
1458   if (fc && !fu) {
1459     nsCOMPtr<nsIInputStream> fileInputStream, bufferedInputStream;
1460     nsresult rv = aChannel->Open(getter_AddRefs(fileInputStream));
1461     NS_ENSURE_SUCCESS(rv, rv);
1462     rv = NS_NewBufferedInputStream(getter_AddRefs(bufferedInputStream),
1463                                    fileInputStream.forget(),
1464                                    BUFFERED_OUTPUT_SIZE);
1465     NS_ENSURE_SUCCESS(rv, rv);
1466     nsAutoCString contentType;
1467     aChannel->GetContentType(contentType);
1468     return StartUpload(bufferedInputStream, aFile, contentType);
1469   }
1470
1471   // Mark save channel as throttleable.
1472   nsCOMPtr<nsIClassOfService> cos(do_QueryInterface(aChannel));
1473   if (cos) {
1474     cos->AddClassFlags(nsIClassOfService::Throttleable);
1475   }
1476
1477   // Read from the input channel
1478   nsresult rv = aChannel->AsyncOpen(this);
1479   if (rv == NS_ERROR_NO_CONTENT) {
1480     // Assume this is a protocol such as mailto: which does not feed out
1481     // data and just ignore it.
1482     return NS_SUCCESS_DONT_FIXUP;
1483   }
1484
1485   if (NS_FAILED(rv)) {
1486     // Opening failed, but do we care?
1487     if (mPersistFlags & PERSIST_FLAGS_FAIL_ON_BROKEN_LINKS) {
1488       SendErrorStatusChange(true, rv, aChannel, aFile);
1489       EndDownload(NS_ERROR_FAILURE);
1490       return NS_ERROR_FAILURE;
1491     }
1492     return NS_SUCCESS_DONT_FIXUP;
1493   }
1494
1495   MutexAutoLock lock(mOutputMapMutex);
1496   // Add the output transport to the output map with the channel as the key
1497   nsCOMPtr<nsISupports> keyPtr = do_QueryInterface(aChannel);
1498   mOutputMap.InsertOrUpdate(keyPtr,
1499                             MakeUnique<OutputData>(aFile, mURI, aCalcFileExt));
1500
1501   return NS_OK;
1502 }
1503
1504 nsresult nsWebBrowserPersist::GetExtensionForContentType(
1505     const char16_t* aContentType, char16_t** aExt) {
1506   NS_ENSURE_ARG_POINTER(aContentType);
1507   NS_ENSURE_ARG_POINTER(aExt);
1508
1509   *aExt = nullptr;
1510
1511   nsresult rv;
1512   if (!mMIMEService) {
1513     mMIMEService = do_GetService(NS_MIMESERVICE_CONTRACTID, &rv);
1514     NS_ENSURE_TRUE(mMIMEService, NS_ERROR_FAILURE);
1515   }
1516
1517   nsAutoCString contentType;
1518   LossyCopyUTF16toASCII(MakeStringSpan(aContentType), contentType);
1519   nsAutoCString ext;
1520   rv = mMIMEService->GetPrimaryExtension(contentType, ""_ns, ext);
1521   if (NS_SUCCEEDED(rv)) {
1522     *aExt = UTF8ToNewUnicode(ext);
1523     NS_ENSURE_TRUE(*aExt, NS_ERROR_OUT_OF_MEMORY);
1524     return NS_OK;
1525   }
1526
1527   return NS_ERROR_FAILURE;
1528 }
1529
1530 nsresult nsWebBrowserPersist::SaveDocumentDeferred(
1531     mozilla::UniquePtr<WalkData>&& aData) {
1532   nsresult rv =
1533       SaveDocumentInternal(aData->mDocument, aData->mFile, aData->mDataPath);
1534   if (NS_FAILED(rv)) {
1535     SendErrorStatusChange(true, rv, nullptr, mURI);
1536     EndDownload(rv);
1537   }
1538   return rv;
1539 }
1540
1541 nsresult nsWebBrowserPersist::SaveDocumentInternal(
1542     nsIWebBrowserPersistDocument* aDocument, nsIURI* aFile, nsIURI* aDataPath) {
1543   mURI = nullptr;
1544   NS_ENSURE_ARG_POINTER(aDocument);
1545   NS_ENSURE_ARG_POINTER(aFile);
1546
1547   nsresult rv = aDocument->SetPersistFlags(mPersistFlags);
1548   NS_ENSURE_SUCCESS(rv, rv);
1549
1550   rv = aDocument->GetIsPrivate(&mIsPrivate);
1551   NS_ENSURE_SUCCESS(rv, rv);
1552
1553   // See if we can get the local file representation of this URI
1554   nsCOMPtr<nsIFile> localFile;
1555   rv = GetLocalFileFromURI(aFile, getter_AddRefs(localFile));
1556
1557   nsCOMPtr<nsIFile> localDataPath;
1558   if (NS_SUCCEEDED(rv) && aDataPath) {
1559     // See if we can get the local file representation of this URI
1560     rv = GetLocalFileFromURI(aDataPath, getter_AddRefs(localDataPath));
1561     NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
1562   }
1563
1564   // Persist the main document
1565   rv = aDocument->GetCharacterSet(mCurrentCharset);
1566   NS_ENSURE_SUCCESS(rv, rv);
1567   nsAutoCString uriSpec;
1568   rv = aDocument->GetDocumentURI(uriSpec);
1569   NS_ENSURE_SUCCESS(rv, rv);
1570   rv = NS_NewURI(getter_AddRefs(mURI), uriSpec, mCurrentCharset.get());
1571   NS_ENSURE_SUCCESS(rv, rv);
1572   rv = aDocument->GetBaseURI(uriSpec);
1573   NS_ENSURE_SUCCESS(rv, rv);
1574   rv = NS_NewURI(getter_AddRefs(mCurrentBaseURI), uriSpec,
1575                  mCurrentCharset.get());
1576   NS_ENSURE_SUCCESS(rv, rv);
1577
1578   // Does the caller want to fixup the referenced URIs and save those too?
1579   if (aDataPath) {
1580     // Basic steps are these.
1581     //
1582     // 1. Iterate through the document (and subdocuments) building a list
1583     //    of unique URIs.
1584     // 2. For each URI create an OutputData entry and open a channel to save
1585     //    it. As each URI is saved, discover the mime type and fix up the
1586     //    local filename with the correct extension.
1587     // 3. Store the document in a list and wait for URI persistence to finish
1588     // 4. After URI persistence completes save the list of documents,
1589     //    fixing it up as it goes out to file.
1590
1591     mCurrentDataPathIsRelative = false;
1592     mCurrentDataPath = aDataPath;
1593     mCurrentRelativePathToData = "";
1594     mCurrentThingsToPersist = 0;
1595     mTargetBaseURI = aFile;
1596
1597     // Determine if the specified data path is relative to the
1598     // specified file, (e.g. c:\docs\htmldata is relative to
1599     // c:\docs\myfile.htm, but not to d:\foo\data.
1600
1601     // Starting with the data dir work back through its parents
1602     // checking if one of them matches the base directory.
1603
1604     if (localDataPath && localFile) {
1605       nsCOMPtr<nsIFile> baseDir;
1606       localFile->GetParent(getter_AddRefs(baseDir));
1607
1608       nsAutoCString relativePathToData;
1609       nsCOMPtr<nsIFile> dataDirParent;
1610       dataDirParent = localDataPath;
1611       while (dataDirParent) {
1612         bool sameDir = false;
1613         dataDirParent->Equals(baseDir, &sameDir);
1614         if (sameDir) {
1615           mCurrentRelativePathToData = relativePathToData;
1616           mCurrentDataPathIsRelative = true;
1617           break;
1618         }
1619
1620         nsAutoString dirName;
1621         dataDirParent->GetLeafName(dirName);
1622
1623         nsAutoCString newRelativePathToData;
1624         newRelativePathToData =
1625             NS_ConvertUTF16toUTF8(dirName) + "/"_ns + relativePathToData;
1626         relativePathToData = newRelativePathToData;
1627
1628         nsCOMPtr<nsIFile> newDataDirParent;
1629         rv = dataDirParent->GetParent(getter_AddRefs(newDataDirParent));
1630         dataDirParent = newDataDirParent;
1631       }
1632     } else {
1633       // generate a relative path if possible
1634       nsCOMPtr<nsIURL> pathToBaseURL(do_QueryInterface(aFile));
1635       if (pathToBaseURL) {
1636         nsAutoCString relativePath;  // nsACString
1637         if (NS_SUCCEEDED(
1638                 pathToBaseURL->GetRelativeSpec(aDataPath, relativePath))) {
1639           mCurrentDataPathIsRelative = true;
1640           mCurrentRelativePathToData = relativePath;
1641         }
1642       }
1643     }
1644
1645     // Store the document in a list so when URI persistence is done and the
1646     // filenames of saved URIs are known, the documents can be fixed up and
1647     // saved
1648
1649     auto* docData = new DocData;
1650     docData->mBaseURI = mCurrentBaseURI;
1651     docData->mCharset = mCurrentCharset;
1652     docData->mDocument = aDocument;
1653     docData->mFile = aFile;
1654     mDocList.AppendElement(docData);
1655
1656     // Walk the DOM gathering a list of externally referenced URIs in the uri
1657     // map
1658     nsCOMPtr<nsIWebBrowserPersistResourceVisitor> visit =
1659         new OnWalk(this, aFile, localDataPath);
1660     return aDocument->ReadResources(visit);
1661   } else {
1662     auto* docData = new DocData;
1663     docData->mBaseURI = mCurrentBaseURI;
1664     docData->mCharset = mCurrentCharset;
1665     docData->mDocument = aDocument;
1666     docData->mFile = aFile;
1667     mDocList.AppendElement(docData);
1668
1669     // Not walking DOMs, so go directly to serialization.
1670     SerializeNextFile();
1671     return NS_OK;
1672   }
1673 }
1674
1675 NS_IMETHODIMP
1676 nsWebBrowserPersist::OnWalk::VisitResource(
1677     nsIWebBrowserPersistDocument* aDoc, const nsACString& aURI,
1678     nsContentPolicyType aContentPolicyType) {
1679   return mParent->StoreURI(aURI, aDoc, aContentPolicyType);
1680 }
1681
1682 NS_IMETHODIMP
1683 nsWebBrowserPersist::OnWalk::VisitDocument(
1684     nsIWebBrowserPersistDocument* aDoc, nsIWebBrowserPersistDocument* aSubDoc) {
1685   URIData* data = nullptr;
1686   nsAutoCString uriSpec;
1687   nsresult rv = aSubDoc->GetDocumentURI(uriSpec);
1688   NS_ENSURE_SUCCESS(rv, rv);
1689   rv = mParent->StoreURI(uriSpec, aDoc, nsIContentPolicy::TYPE_SUBDOCUMENT,
1690                          false, &data);
1691   NS_ENSURE_SUCCESS(rv, rv);
1692   if (!data) {
1693     // If the URI scheme isn't persistable, then don't persist.
1694     return NS_OK;
1695   }
1696   data->mIsSubFrame = true;
1697   return mParent->SaveSubframeContent(aSubDoc, aDoc, uriSpec, data);
1698 }
1699
1700 NS_IMETHODIMP
1701 nsWebBrowserPersist::OnWalk::VisitBrowsingContext(
1702     nsIWebBrowserPersistDocument* aDoc, BrowsingContext* aContext) {
1703   RefPtr<dom::CanonicalBrowsingContext> context = aContext->Canonical();
1704
1705   if (NS_WARN_IF(!context->GetCurrentWindowGlobal())) {
1706     EndVisit(nullptr, NS_ERROR_FAILURE);
1707     return NS_ERROR_FAILURE;
1708   }
1709
1710   RefPtr<WebBrowserPersistDocumentParent> actor(
1711       new WebBrowserPersistDocumentParent());
1712
1713   nsCOMPtr<nsIWebBrowserPersistDocumentReceiver> receiver =
1714       new OnRemoteWalk(this, aDoc);
1715   actor->SetOnReady(receiver);
1716
1717   RefPtr<dom::BrowserParent> browserParent =
1718       context->GetCurrentWindowGlobal()->GetBrowserParent();
1719
1720   bool ok =
1721       context->GetContentParent()->SendPWebBrowserPersistDocumentConstructor(
1722           actor, browserParent, context);
1723
1724   if (NS_WARN_IF(!ok)) {
1725     // (The actor will be destroyed on constructor failure.)
1726     EndVisit(nullptr, NS_ERROR_FAILURE);
1727     return NS_ERROR_FAILURE;
1728   }
1729
1730   ++mPendingDocuments;
1731
1732   return NS_OK;
1733 }
1734
1735 NS_IMETHODIMP
1736 nsWebBrowserPersist::OnWalk::EndVisit(nsIWebBrowserPersistDocument* aDoc,
1737                                       nsresult aStatus) {
1738   if (NS_FAILED(mStatus)) {
1739     return mStatus;
1740   }
1741
1742   if (NS_FAILED(aStatus)) {
1743     mStatus = aStatus;
1744     mParent->SendErrorStatusChange(true, aStatus, nullptr, mFile);
1745     mParent->EndDownload(aStatus);
1746     return aStatus;
1747   }
1748
1749   if (--mPendingDocuments) {
1750     // We're not done yet, wait for more.
1751     return NS_OK;
1752   }
1753
1754   mParent->FinishSaveDocumentInternal(mFile, mDataPath);
1755   return NS_OK;
1756 }
1757
1758 NS_IMETHODIMP
1759 nsWebBrowserPersist::OnRemoteWalk::OnDocumentReady(
1760     nsIWebBrowserPersistDocument* aSubDocument) {
1761   mVisitor->VisitDocument(mDocument, aSubDocument);
1762   mVisitor->EndVisit(mDocument, NS_OK);
1763   return NS_OK;
1764 }
1765
1766 NS_IMETHODIMP
1767 nsWebBrowserPersist::OnRemoteWalk::OnError(nsresult aFailure) {
1768   mVisitor->EndVisit(nullptr, aFailure);
1769   return NS_OK;
1770 }
1771
1772 void nsWebBrowserPersist::FinishSaveDocumentInternal(nsIURI* aFile,
1773                                                      nsIFile* aDataPath) {
1774   // If there are things to persist, create a directory to hold them
1775   if (mCurrentThingsToPersist > 0) {
1776     if (aDataPath) {
1777       bool exists = false;
1778       bool haveDir = false;
1779
1780       aDataPath->Exists(&exists);
1781       if (exists) {
1782         aDataPath->IsDirectory(&haveDir);
1783       }
1784       if (!haveDir) {
1785         nsresult rv = aDataPath->Create(nsIFile::DIRECTORY_TYPE, 0755);
1786         if (NS_SUCCEEDED(rv)) {
1787           haveDir = true;
1788         } else {
1789           SendErrorStatusChange(false, rv, nullptr, aFile);
1790         }
1791       }
1792       if (!haveDir) {
1793         EndDownload(NS_ERROR_FAILURE);
1794         return;
1795       }
1796       if (mPersistFlags & PERSIST_FLAGS_CLEANUP_ON_FAILURE) {
1797         // Add to list of things to delete later if all goes wrong
1798         auto* cleanupData = new CleanupData;
1799         cleanupData->mFile = aDataPath;
1800         cleanupData->mIsDirectory = true;
1801         mCleanupList.AppendElement(cleanupData);
1802       }
1803     }
1804   }
1805
1806   if (mWalkStack.Length() > 0) {
1807     mozilla::UniquePtr<WalkData> toWalk = mWalkStack.PopLastElement();
1808     // Bounce this off the event loop to avoid stack overflow.
1809     using WalkStorage = StoreCopyPassByRRef<decltype(toWalk)>;
1810     auto saveMethod = &nsWebBrowserPersist::SaveDocumentDeferred;
1811     nsCOMPtr<nsIRunnable> saveLater = NewRunnableMethod<WalkStorage>(
1812         "nsWebBrowserPersist::FinishSaveDocumentInternal", this, saveMethod,
1813         std::move(toWalk));
1814     NS_DispatchToCurrentThread(saveLater);
1815   } else {
1816     // Done walking DOMs; on to the serialization phase.
1817     SerializeNextFile();
1818   }
1819 }
1820
1821 void nsWebBrowserPersist::Cleanup() {
1822   mURIMap.Clear();
1823   nsClassHashtable<nsISupportsHashKey, OutputData> outputMapCopy;
1824   {
1825     MutexAutoLock lock(mOutputMapMutex);
1826     mOutputMap.SwapElements(outputMapCopy);
1827   }
1828   for (const auto& key : outputMapCopy.Keys()) {
1829     nsCOMPtr<nsIChannel> channel = do_QueryInterface(key);
1830     if (channel) {
1831       channel->Cancel(NS_BINDING_ABORTED);
1832     }
1833   }
1834   outputMapCopy.Clear();
1835
1836   for (const auto& key : mUploadList.Keys()) {
1837     nsCOMPtr<nsIChannel> channel = do_QueryInterface(key);
1838     if (channel) {
1839       channel->Cancel(NS_BINDING_ABORTED);
1840     }
1841   }
1842   mUploadList.Clear();
1843
1844   uint32_t i;
1845   for (i = 0; i < mDocList.Length(); i++) {
1846     DocData* docData = mDocList.ElementAt(i);
1847     delete docData;
1848   }
1849   mDocList.Clear();
1850
1851   for (i = 0; i < mCleanupList.Length(); i++) {
1852     CleanupData* cleanupData = mCleanupList.ElementAt(i);
1853     delete cleanupData;
1854   }
1855   mCleanupList.Clear();
1856
1857   mFilenameList.Clear();
1858 }
1859
1860 void nsWebBrowserPersist::CleanupLocalFiles() {
1861   // Two passes, the first pass cleans up files, the second pass tests
1862   // for and then deletes empty directories. Directories that are not
1863   // empty after the first pass must contain files from something else
1864   // and are not deleted.
1865   int pass;
1866   for (pass = 0; pass < 2; pass++) {
1867     uint32_t i;
1868     for (i = 0; i < mCleanupList.Length(); i++) {
1869       CleanupData* cleanupData = mCleanupList.ElementAt(i);
1870       nsCOMPtr<nsIFile> file = cleanupData->mFile;
1871
1872       // Test if the dir / file exists (something in an earlier loop
1873       // may have already removed it)
1874       bool exists = false;
1875       file->Exists(&exists);
1876       if (!exists) continue;
1877
1878       // Test if the file has changed in between creation and deletion
1879       // in some way that means it should be ignored
1880       bool isDirectory = false;
1881       file->IsDirectory(&isDirectory);
1882       if (isDirectory != cleanupData->mIsDirectory)
1883         continue;  // A file has become a dir or vice versa !
1884
1885       if (pass == 0 && !isDirectory) {
1886         file->Remove(false);
1887       } else if (pass == 1 && isDirectory)  // Directory
1888       {
1889         // Directories are more complicated. Enumerate through
1890         // children looking for files. Any files created by the
1891         // persist object would have been deleted by the first
1892         // pass so if there are any there at this stage, the dir
1893         // cannot be deleted because it has someone else's files
1894         // in it. Empty child dirs are deleted but they must be
1895         // recursed through to ensure they are actually empty.
1896
1897         bool isEmptyDirectory = true;
1898         nsCOMArray<nsIDirectoryEnumerator> dirStack;
1899         int32_t stackSize = 0;
1900
1901         // Push the top level enum onto the stack
1902         nsCOMPtr<nsIDirectoryEnumerator> pos;
1903         if (NS_SUCCEEDED(file->GetDirectoryEntries(getter_AddRefs(pos))))
1904           dirStack.AppendObject(pos);
1905
1906         while (isEmptyDirectory && (stackSize = dirStack.Count())) {
1907           // Pop the last element
1908           nsCOMPtr<nsIDirectoryEnumerator> curPos;
1909           curPos = dirStack[stackSize - 1];
1910           dirStack.RemoveObjectAt(stackSize - 1);
1911
1912           nsCOMPtr<nsIFile> child;
1913           if (NS_FAILED(curPos->GetNextFile(getter_AddRefs(child))) || !child) {
1914             continue;
1915           }
1916
1917           bool childIsSymlink = false;
1918           child->IsSymlink(&childIsSymlink);
1919           bool childIsDir = false;
1920           child->IsDirectory(&childIsDir);
1921           if (!childIsDir || childIsSymlink) {
1922             // Some kind of file or symlink which means dir
1923             // is not empty so just drop out.
1924             isEmptyDirectory = false;
1925             break;
1926           }
1927           // Push parent enumerator followed by child enumerator
1928           nsCOMPtr<nsIDirectoryEnumerator> childPos;
1929           child->GetDirectoryEntries(getter_AddRefs(childPos));
1930           dirStack.AppendObject(curPos);
1931           if (childPos) dirStack.AppendObject(childPos);
1932         }
1933         dirStack.Clear();
1934
1935         // If after all that walking the dir is deemed empty, delete it
1936         if (isEmptyDirectory) {
1937           file->Remove(true);
1938         }
1939       }
1940     }
1941   }
1942 }
1943
1944 nsresult nsWebBrowserPersist::CalculateUniqueFilename(
1945     nsIURI* aURI, nsCOMPtr<nsIURI>& aOutURI) {
1946   nsCOMPtr<nsIURL> url(do_QueryInterface(aURI));
1947   NS_ENSURE_TRUE(url, NS_ERROR_FAILURE);
1948
1949   bool nameHasChanged = false;
1950   nsresult rv;
1951
1952   // Get the old filename
1953   nsAutoCString filename;
1954   rv = url->GetFileName(filename);
1955   NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
1956   nsAutoCString directory;
1957   rv = url->GetDirectory(directory);
1958   NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
1959
1960   // Split the filename into a base and an extension.
1961   // e.g. "foo.html" becomes "foo" & ".html"
1962   //
1963   // The nsIURL methods GetFileBaseName & GetFileExtension don't
1964   // preserve the dot whereas this code does to save some effort
1965   // later when everything is put back together.
1966   int32_t lastDot = filename.RFind(".");
1967   nsAutoCString base;
1968   nsAutoCString ext;
1969   if (lastDot >= 0) {
1970     filename.Mid(base, 0, lastDot);
1971     filename.Mid(ext, lastDot, filename.Length() - lastDot);  // includes dot
1972   } else {
1973     // filename contains no dot
1974     base = filename;
1975   }
1976
1977   // Test if the filename is longer than allowed by the OS
1978   int32_t needToChop = filename.Length() - kDefaultMaxFilenameLength;
1979   if (needToChop > 0) {
1980     // Truncate the base first and then the ext if necessary
1981     if (base.Length() > (uint32_t)needToChop) {
1982       base.Truncate(base.Length() - needToChop);
1983     } else {
1984       needToChop -= base.Length() - 1;
1985       base.Truncate(1);
1986       if (ext.Length() > (uint32_t)needToChop) {
1987         ext.Truncate(ext.Length() - needToChop);
1988       } else {
1989         ext.Truncate(0);
1990       }
1991       // If kDefaultMaxFilenameLength were 1 we'd be in trouble here,
1992       // but that won't happen because it will be set to a sensible
1993       // value.
1994     }
1995
1996     filename.Assign(base);
1997     filename.Append(ext);
1998     nameHasChanged = true;
1999   }
2000
2001   // Ensure the filename is unique
2002   // Create a filename if it's empty, or if the filename / datapath is
2003   // already taken by another URI and create an alternate name.
2004
2005   if (base.IsEmpty() || !mFilenameList.IsEmpty()) {
2006     nsAutoCString tmpPath;
2007     nsAutoCString tmpBase;
2008     uint32_t duplicateCounter = 1;
2009     while (true) {
2010       // Make a file name,
2011       // Foo become foo_001, foo_002, etc.
2012       // Empty files become _001, _002 etc.
2013
2014       if (base.IsEmpty() || duplicateCounter > 1) {
2015         SmprintfPointer tmp = mozilla::Smprintf("_%03d", duplicateCounter);
2016         NS_ENSURE_TRUE(tmp, NS_ERROR_OUT_OF_MEMORY);
2017         if (filename.Length() < kDefaultMaxFilenameLength - 4) {
2018           tmpBase = base;
2019         } else {
2020           base.Mid(tmpBase, 0, base.Length() - 4);
2021         }
2022         tmpBase.Append(tmp.get());
2023       } else {
2024         tmpBase = base;
2025       }
2026
2027       tmpPath.Assign(directory);
2028       tmpPath.Append(tmpBase);
2029       tmpPath.Append(ext);
2030
2031       // Test if the name is a duplicate
2032       if (!mFilenameList.Contains(tmpPath)) {
2033         if (!base.Equals(tmpBase)) {
2034           filename.Assign(tmpBase);
2035           filename.Append(ext);
2036           nameHasChanged = true;
2037         }
2038         break;
2039       }
2040       duplicateCounter++;
2041     }
2042   }
2043
2044   // Add name to list of those already used
2045   nsAutoCString newFilepath(directory);
2046   newFilepath.Append(filename);
2047   mFilenameList.AppendElement(newFilepath);
2048
2049   // Update the uri accordingly if the filename actually changed
2050   if (nameHasChanged) {
2051     // Final sanity test
2052     if (filename.Length() > kDefaultMaxFilenameLength) {
2053       NS_WARNING(
2054           "Filename wasn't truncated less than the max file length - how can "
2055           "that be?");
2056       return NS_ERROR_FAILURE;
2057     }
2058
2059     nsCOMPtr<nsIFile> localFile;
2060     GetLocalFileFromURI(aURI, getter_AddRefs(localFile));
2061
2062     if (localFile) {
2063       nsAutoString filenameAsUnichar;
2064       CopyASCIItoUTF16(filename, filenameAsUnichar);
2065       localFile->SetLeafName(filenameAsUnichar);
2066
2067       // Resync the URI with the file after the extension has been appended
2068       return NS_MutateURI(aURI)
2069           .Apply(&nsIFileURLMutator::SetFile, localFile)
2070           .Finalize(aOutURI);
2071     }
2072     return NS_MutateURI(url)
2073         .Apply(&nsIURLMutator::SetFileName, filename, nullptr)
2074         .Finalize(aOutURI);
2075   }
2076
2077   // TODO (:valentin) This method should always clone aURI
2078   aOutURI = aURI;
2079   return NS_OK;
2080 }
2081
2082 nsresult nsWebBrowserPersist::MakeFilenameFromURI(nsIURI* aURI,
2083                                                   nsString& aFilename) {
2084   // Try to get filename from the URI.
2085   nsAutoString fileName;
2086
2087   // Get a suggested file name from the URL but strip it of characters
2088   // likely to cause the name to be illegal.
2089
2090   nsCOMPtr<nsIURL> url(do_QueryInterface(aURI));
2091   if (url) {
2092     nsAutoCString nameFromURL;
2093     url->GetFileName(nameFromURL);
2094     if (mPersistFlags & PERSIST_FLAGS_DONT_CHANGE_FILENAMES) {
2095       CopyASCIItoUTF16(NS_UnescapeURL(nameFromURL), fileName);
2096       aFilename = fileName;
2097       return NS_OK;
2098     }
2099     if (!nameFromURL.IsEmpty()) {
2100       // Unescape the file name (GetFileName escapes it)
2101       NS_UnescapeURL(nameFromURL);
2102       uint32_t nameLength = 0;
2103       const char* p = nameFromURL.get();
2104       for (; *p && *p != ';' && *p != '?' && *p != '#' && *p != '.'; p++) {
2105         if (IsAsciiAlpha(*p) || IsAsciiDigit(*p) || *p == '.' || *p == '-' ||
2106             *p == '_' || (*p == ' ')) {
2107           fileName.Append(char16_t(*p));
2108           if (++nameLength == kDefaultMaxFilenameLength) {
2109             // Note:
2110             // There is no point going any further since it will be
2111             // truncated in CalculateUniqueFilename anyway.
2112             // More importantly, certain implementations of
2113             // nsIFile (e.g. the Mac impl) might truncate
2114             // names in undesirable ways, such as truncating from
2115             // the middle, inserting ellipsis and so on.
2116             break;
2117           }
2118         }
2119       }
2120     }
2121   }
2122
2123   // Empty filenames can confuse the local file object later
2124   // when it attempts to set the leaf name in CalculateUniqueFilename
2125   // for duplicates and ends up replacing the parent dir. To avoid
2126   // the problem, all filenames are made at least one character long.
2127   if (fileName.IsEmpty()) {
2128     fileName.Append(char16_t('a'));  // 'a' is for arbitrary
2129   }
2130
2131   aFilename = fileName;
2132   return NS_OK;
2133 }
2134
2135 nsresult nsWebBrowserPersist::CalculateAndAppendFileExt(
2136     nsIURI* aURI, nsIChannel* aChannel, nsIURI* aOriginalURIWithExtension,
2137     nsCOMPtr<nsIURI>& aOutURI) {
2138   nsresult rv = NS_OK;
2139
2140   if (!mMIMEService) {
2141     mMIMEService = do_GetService(NS_MIMESERVICE_CONTRACTID, &rv);
2142     NS_ENSURE_TRUE(mMIMEService, NS_ERROR_FAILURE);
2143   }
2144
2145   nsAutoCString contentType;
2146
2147   // Get the content type from the channel
2148   aChannel->GetContentType(contentType);
2149
2150   // Get the content type from the MIME service
2151   if (contentType.IsEmpty()) {
2152     nsCOMPtr<nsIURI> uri;
2153     aChannel->GetOriginalURI(getter_AddRefs(uri));
2154     mMIMEService->GetTypeFromURI(uri, contentType);
2155   }
2156
2157   // Validate the filename
2158   if (!contentType.IsEmpty()) {
2159     nsAutoString newFileName;
2160     if (NS_SUCCEEDED(mMIMEService->GetValidFileName(
2161             aChannel, contentType, aOriginalURIWithExtension,
2162             nsIMIMEService::VALIDATE_DEFAULT, newFileName))) {
2163       nsCOMPtr<nsIFile> localFile;
2164       GetLocalFileFromURI(aURI, getter_AddRefs(localFile));
2165       if (localFile) {
2166         localFile->SetLeafName(newFileName);
2167
2168         // Resync the URI with the file after the extension has been appended
2169         return NS_MutateURI(aURI)
2170             .Apply(&nsIFileURLMutator::SetFile, localFile)
2171             .Finalize(aOutURI);
2172       }
2173       return NS_MutateURI(aURI)
2174           .Apply(&nsIURLMutator::SetFileName,
2175                  NS_ConvertUTF16toUTF8(newFileName), nullptr)
2176           .Finalize(aOutURI);
2177     }
2178   }
2179
2180   // TODO (:valentin) This method should always clone aURI
2181   aOutURI = aURI;
2182   return NS_OK;
2183 }
2184
2185 // Note: the MakeOutputStream helpers can be called from a background thread.
2186 nsresult nsWebBrowserPersist::MakeOutputStream(
2187     nsIURI* aURI, nsIOutputStream** aOutputStream) {
2188   nsresult rv;
2189
2190   nsCOMPtr<nsIFile> localFile;
2191   GetLocalFileFromURI(aURI, getter_AddRefs(localFile));
2192   if (localFile)
2193     rv = MakeOutputStreamFromFile(localFile, aOutputStream);
2194   else
2195     rv = MakeOutputStreamFromURI(aURI, aOutputStream);
2196
2197   return rv;
2198 }
2199
2200 nsresult nsWebBrowserPersist::MakeOutputStreamFromFile(
2201     nsIFile* aFile, nsIOutputStream** aOutputStream) {
2202   nsresult rv = NS_OK;
2203
2204   nsCOMPtr<nsIFileOutputStream> fileOutputStream =
2205       do_CreateInstance(NS_LOCALFILEOUTPUTSTREAM_CONTRACTID, &rv);
2206   NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
2207
2208   // XXX brade:  get the right flags here!
2209   int32_t ioFlags = -1;
2210   if (mPersistFlags & nsIWebBrowserPersist::PERSIST_FLAGS_APPEND_TO_FILE)
2211     ioFlags = PR_APPEND | PR_CREATE_FILE | PR_WRONLY;
2212   rv = fileOutputStream->Init(aFile, ioFlags, -1, 0);
2213   NS_ENSURE_SUCCESS(rv, rv);
2214
2215   rv = NS_NewBufferedOutputStream(aOutputStream, fileOutputStream.forget(),
2216                                   BUFFERED_OUTPUT_SIZE);
2217   NS_ENSURE_SUCCESS(rv, rv);
2218
2219   if (mPersistFlags & PERSIST_FLAGS_CLEANUP_ON_FAILURE) {
2220     // Add to cleanup list in event of failure
2221     auto* cleanupData = new CleanupData;
2222     cleanupData->mFile = aFile;
2223     cleanupData->mIsDirectory = false;
2224     if (NS_IsMainThread()) {
2225       mCleanupList.AppendElement(cleanupData);
2226     } else {
2227       // If we're on a background thread, add the cleanup back on the main
2228       // thread.
2229       RefPtr<Runnable> addCleanup = NS_NewRunnableFunction(
2230           "nsWebBrowserPersist::AddCleanupToList",
2231           [self = RefPtr{this}, cleanup = std::move(cleanupData)]() {
2232             self->mCleanupList.AppendElement(cleanup);
2233           });
2234       NS_DispatchToMainThread(addCleanup);
2235     }
2236   }
2237
2238   return NS_OK;
2239 }
2240
2241 nsresult nsWebBrowserPersist::MakeOutputStreamFromURI(
2242     nsIURI* aURI, nsIOutputStream** aOutputStream) {
2243   uint32_t segsize = 8192;
2244   uint32_t maxsize = uint32_t(-1);
2245   nsCOMPtr<nsIStorageStream> storStream;
2246   nsresult rv =
2247       NS_NewStorageStream(segsize, maxsize, getter_AddRefs(storStream));
2248   NS_ENSURE_SUCCESS(rv, rv);
2249
2250   NS_ENSURE_SUCCESS(CallQueryInterface(storStream, aOutputStream),
2251                     NS_ERROR_FAILURE);
2252   return NS_OK;
2253 }
2254
2255 void nsWebBrowserPersist::FinishDownload() {
2256   // We call FinishDownload when we run out of things to download for this
2257   // persist operation, by dispatching this method to the main thread. By now,
2258   // it's possible that we have been canceled or encountered an error earlier
2259   // in the download, or something else called EndDownload. In that case, don't
2260   // re-run EndDownload.
2261   if (mEndCalled) {
2262     return;
2263   }
2264   EndDownload(NS_OK);
2265 }
2266
2267 void nsWebBrowserPersist::EndDownload(nsresult aResult) {
2268   MOZ_ASSERT(NS_IsMainThread(), "Should end download on the main thread.");
2269
2270   // Really this should just never happen, but if it does, at least avoid
2271   // no-op notifications or pretending we succeeded if we already failed.
2272   if (mEndCalled && (NS_SUCCEEDED(aResult) || mPersistResult == aResult)) {
2273     return;
2274   }
2275
2276   // Store the error code in the result if it is an error
2277   if (NS_SUCCEEDED(mPersistResult) && NS_FAILED(aResult)) {
2278     mPersistResult = aResult;
2279   }
2280
2281   if (mEndCalled) {
2282     MOZ_ASSERT(!mEndCalled, "Should only end the download once.");
2283     return;
2284   }
2285   mEndCalled = true;
2286
2287   ClosePromise::All(GetCurrentSerialEventTarget(), mFileClosePromises)
2288       ->Then(GetCurrentSerialEventTarget(), __func__,
2289              [self = RefPtr{this}, aResult]() {
2290                self->EndDownloadInternal(aResult);
2291              });
2292 }
2293
2294 void nsWebBrowserPersist::EndDownloadInternal(nsresult aResult) {
2295   // mCompleted needs to be set before issuing the stop notification.
2296   // (Bug 1224437)
2297   mCompleted = true;
2298   // State stop notification
2299   if (mProgressListener) {
2300     mProgressListener->OnStateChange(
2301         nullptr, nullptr,
2302         nsIWebProgressListener::STATE_STOP |
2303             nsIWebProgressListener::STATE_IS_NETWORK,
2304         mPersistResult);
2305   }
2306
2307   // Do file cleanup if required
2308   if (NS_FAILED(aResult) &&
2309       (mPersistFlags & PERSIST_FLAGS_CLEANUP_ON_FAILURE)) {
2310     CleanupLocalFiles();
2311   }
2312
2313   // Cleanup the channels
2314   Cleanup();
2315
2316   mProgressListener = nullptr;
2317   mProgressListener2 = nullptr;
2318   mEventSink = nullptr;
2319 }
2320
2321 nsresult nsWebBrowserPersist::FixRedirectedChannelEntry(
2322     nsIChannel* aNewChannel) {
2323   NS_ENSURE_ARG_POINTER(aNewChannel);
2324
2325   // Iterate through existing open channels looking for one with a URI
2326   // matching the one specified.
2327   nsCOMPtr<nsIURI> originalURI;
2328   aNewChannel->GetOriginalURI(getter_AddRefs(originalURI));
2329   nsISupports* matchingKey = nullptr;
2330   for (nsISupports* key : mOutputMap.Keys()) {
2331     nsCOMPtr<nsIChannel> thisChannel = do_QueryInterface(key);
2332     nsCOMPtr<nsIURI> thisURI;
2333
2334     thisChannel->GetOriginalURI(getter_AddRefs(thisURI));
2335
2336     // Compare this channel's URI to the one passed in.
2337     bool matchingURI = false;
2338     thisURI->Equals(originalURI, &matchingURI);
2339     if (matchingURI) {
2340       matchingKey = key;
2341       break;
2342     }
2343   }
2344
2345   if (matchingKey) {
2346     // We only get called from OnStartRequest, so this is always on the
2347     // main thread. Make sure we don't pull the rug from under anything else.
2348     MutexAutoLock lock(mOutputMapMutex);
2349     // If a match was found, remove the data entry with the old channel
2350     // key and re-add it with the new channel key.
2351     mozilla::UniquePtr<OutputData> outputData;
2352     mOutputMap.Remove(matchingKey, &outputData);
2353     NS_ENSURE_TRUE(outputData, NS_ERROR_FAILURE);
2354
2355     // Store data again with new channel unless told to ignore redirects.
2356     if (!(mPersistFlags & PERSIST_FLAGS_IGNORE_REDIRECTED_DATA)) {
2357       nsCOMPtr<nsISupports> keyPtr = do_QueryInterface(aNewChannel);
2358       mOutputMap.InsertOrUpdate(keyPtr, std::move(outputData));
2359     }
2360   }
2361
2362   return NS_OK;
2363 }
2364
2365 void nsWebBrowserPersist::CalcTotalProgress() {
2366   mTotalCurrentProgress = 0;
2367   mTotalMaxProgress = 0;
2368
2369   if (mOutputMap.Count() > 0) {
2370     // Total up the progress of each output stream
2371     for (const auto& data : mOutputMap.Values()) {
2372       // Only count toward total progress if destination file is local.
2373       nsCOMPtr<nsIFileURL> fileURL = do_QueryInterface(data->mFile);
2374       if (fileURL) {
2375         mTotalCurrentProgress += data->mSelfProgress;
2376         mTotalMaxProgress += data->mSelfProgressMax;
2377       }
2378     }
2379   }
2380
2381   if (mUploadList.Count() > 0) {
2382     // Total up the progress of each upload
2383     for (const auto& data : mUploadList.Values()) {
2384       if (data) {
2385         mTotalCurrentProgress += data->mSelfProgress;
2386         mTotalMaxProgress += data->mSelfProgressMax;
2387       }
2388     }
2389   }
2390
2391   // XXX this code seems pretty bogus and pointless
2392   if (mTotalCurrentProgress == 0 && mTotalMaxProgress == 0) {
2393     // No output streams so we must be complete
2394     mTotalCurrentProgress = 10000;
2395     mTotalMaxProgress = 10000;
2396   }
2397 }
2398
2399 nsresult nsWebBrowserPersist::StoreURI(const nsACString& aURI,
2400                                        nsIWebBrowserPersistDocument* aDoc,
2401                                        nsContentPolicyType aContentPolicyType,
2402                                        bool aNeedsPersisting, URIData** aData) {
2403   nsCOMPtr<nsIURI> uri;
2404   nsresult rv = NS_NewURI(getter_AddRefs(uri), aURI, mCurrentCharset.get(),
2405                           mCurrentBaseURI);
2406   NS_ENSURE_SUCCESS(rv, rv);
2407
2408   return StoreURI(uri, aDoc, aContentPolicyType, aNeedsPersisting, aData);
2409 }
2410
2411 nsresult nsWebBrowserPersist::StoreURI(nsIURI* aURI,
2412                                        nsIWebBrowserPersistDocument* aDoc,
2413                                        nsContentPolicyType aContentPolicyType,
2414                                        bool aNeedsPersisting, URIData** aData) {
2415   NS_ENSURE_ARG_POINTER(aURI);
2416   if (aData) {
2417     *aData = nullptr;
2418   }
2419
2420   // Test if this URI should be persisted. By default
2421   // we should assume the URI  is persistable.
2422   bool doNotPersistURI;
2423   nsresult rv = NS_URIChainHasFlags(
2424       aURI, nsIProtocolHandler::URI_NON_PERSISTABLE, &doNotPersistURI);
2425   if (NS_FAILED(rv)) {
2426     doNotPersistURI = false;
2427   }
2428
2429   if (doNotPersistURI) {
2430     return NS_OK;
2431   }
2432
2433   URIData* data = nullptr;
2434   MakeAndStoreLocalFilenameInURIMap(aURI, aDoc, aContentPolicyType,
2435                                     aNeedsPersisting, &data);
2436   if (aData) {
2437     *aData = data;
2438   }
2439
2440   return NS_OK;
2441 }
2442
2443 nsresult nsWebBrowserPersist::URIData::GetLocalURI(nsIURI* targetBaseURI,
2444                                                    nsCString& aSpecOut) {
2445   aSpecOut.SetIsVoid(true);
2446   if (!mNeedsFixup) {
2447     return NS_OK;
2448   }
2449   nsresult rv;
2450   nsCOMPtr<nsIURI> fileAsURI;
2451   if (mFile) {
2452     fileAsURI = mFile;
2453   } else {
2454     fileAsURI = mDataPath;
2455     rv = AppendPathToURI(fileAsURI, mFilename, fileAsURI);
2456     NS_ENSURE_SUCCESS(rv, rv);
2457   }
2458
2459   // remove username/password if present
2460   Unused << NS_MutateURI(fileAsURI).SetUserPass(""_ns).Finalize(fileAsURI);
2461
2462   // reset node attribute
2463   // Use relative or absolute links
2464   if (mDataPathIsRelative) {
2465     bool isEqual = false;
2466     if (NS_SUCCEEDED(mRelativeDocumentURI->Equals(targetBaseURI, &isEqual)) &&
2467         isEqual) {
2468       nsCOMPtr<nsIURL> url(do_QueryInterface(fileAsURI));
2469       if (!url) {
2470         return NS_ERROR_FAILURE;
2471       }
2472
2473       nsAutoCString filename;
2474       url->GetFileName(filename);
2475
2476       nsAutoCString rawPathURL(mRelativePathToData);
2477       rawPathURL.Append(filename);
2478
2479       rv = NS_EscapeURL(rawPathURL, esc_FilePath, aSpecOut, fallible);
2480       NS_ENSURE_SUCCESS(rv, rv);
2481     } else {
2482       nsAutoCString rawPathURL;
2483
2484       nsCOMPtr<nsIFile> dataFile;
2485       rv = GetLocalFileFromURI(mFile, getter_AddRefs(dataFile));
2486       NS_ENSURE_SUCCESS(rv, rv);
2487
2488       nsCOMPtr<nsIFile> docFile;
2489       rv = GetLocalFileFromURI(targetBaseURI, getter_AddRefs(docFile));
2490       NS_ENSURE_SUCCESS(rv, rv);
2491
2492       nsCOMPtr<nsIFile> parentDir;
2493       rv = docFile->GetParent(getter_AddRefs(parentDir));
2494       NS_ENSURE_SUCCESS(rv, rv);
2495
2496       rv = dataFile->GetRelativePath(parentDir, rawPathURL);
2497       NS_ENSURE_SUCCESS(rv, rv);
2498
2499       rv = NS_EscapeURL(rawPathURL, esc_FilePath, aSpecOut, fallible);
2500       NS_ENSURE_SUCCESS(rv, rv);
2501     }
2502   } else {
2503     fileAsURI->GetSpec(aSpecOut);
2504   }
2505   if (mIsSubFrame) {
2506     AppendUTF16toUTF8(mSubFrameExt, aSpecOut);
2507   }
2508
2509   return NS_OK;
2510 }
2511
2512 bool nsWebBrowserPersist::DocumentEncoderExists(const char* aContentType) {
2513   return do_getDocumentTypeSupportedForEncoding(aContentType);
2514 }
2515
2516 nsresult nsWebBrowserPersist::SaveSubframeContent(
2517     nsIWebBrowserPersistDocument* aFrameContent,
2518     nsIWebBrowserPersistDocument* aParentDocument, const nsCString& aURISpec,
2519     URIData* aData) {
2520   NS_ENSURE_ARG_POINTER(aData);
2521
2522   // Extract the content type for the frame's contents.
2523   nsAutoCString contentType;
2524   nsresult rv = aFrameContent->GetContentType(contentType);
2525   NS_ENSURE_SUCCESS(rv, rv);
2526
2527   nsString ext;
2528   GetExtensionForContentType(NS_ConvertASCIItoUTF16(contentType).get(),
2529                              getter_Copies(ext));
2530
2531   // We must always have an extension so we will try to re-assign
2532   // the original extension if GetExtensionForContentType fails.
2533   if (ext.IsEmpty()) {
2534     nsCOMPtr<nsIURI> docURI;
2535     rv = NS_NewURI(getter_AddRefs(docURI), aURISpec, mCurrentCharset.get());
2536     NS_ENSURE_SUCCESS(rv, rv);
2537
2538     nsCOMPtr<nsIURL> url(do_QueryInterface(docURI, &rv));
2539     nsAutoCString extension;
2540     if (NS_SUCCEEDED(rv)) {
2541       url->GetFileExtension(extension);
2542     } else {
2543       extension.AssignLiteral("htm");
2544     }
2545     aData->mSubFrameExt.Assign(char16_t('.'));
2546     AppendUTF8toUTF16(extension, aData->mSubFrameExt);
2547   } else {
2548     aData->mSubFrameExt.Assign(char16_t('.'));
2549     aData->mSubFrameExt.Append(ext);
2550   }
2551
2552   nsString filenameWithExt = aData->mFilename;
2553   filenameWithExt.Append(aData->mSubFrameExt);
2554
2555   // Work out the path for the subframe
2556   nsCOMPtr<nsIURI> frameURI = mCurrentDataPath;
2557   rv = AppendPathToURI(frameURI, filenameWithExt, frameURI);
2558   NS_ENSURE_SUCCESS(rv, rv);
2559
2560   // Work out the path for the subframe data
2561   nsCOMPtr<nsIURI> frameDataURI = mCurrentDataPath;
2562   nsAutoString newFrameDataPath(aData->mFilename);
2563
2564   // Append _data
2565   newFrameDataPath.AppendLiteral("_data");
2566   rv = AppendPathToURI(frameDataURI, newFrameDataPath, frameDataURI);
2567   NS_ENSURE_SUCCESS(rv, rv);
2568
2569   // Make frame document & data path conformant and unique
2570   nsCOMPtr<nsIURI> out;
2571   rv = CalculateUniqueFilename(frameURI, out);
2572   NS_ENSURE_SUCCESS(rv, rv);
2573   frameURI = out;
2574
2575   rv = CalculateUniqueFilename(frameDataURI, out);
2576   NS_ENSURE_SUCCESS(rv, rv);
2577   frameDataURI = out;
2578
2579   mCurrentThingsToPersist++;
2580
2581   // We shouldn't use SaveDocumentInternal for the contents
2582   // of frames that are not documents, e.g. images.
2583   if (DocumentEncoderExists(contentType.get())) {
2584     auto toWalk = mozilla::MakeUnique<WalkData>();
2585     toWalk->mDocument = aFrameContent;
2586     toWalk->mFile = frameURI;
2587     toWalk->mDataPath = frameDataURI;
2588     mWalkStack.AppendElement(std::move(toWalk));
2589   } else {
2590     nsContentPolicyType policyType = nsIContentPolicy::TYPE_OTHER;
2591     if (StringBeginsWith(contentType, "image/"_ns)) {
2592       policyType = nsIContentPolicy::TYPE_IMAGE;
2593     } else if (StringBeginsWith(contentType, "audio/"_ns) ||
2594                StringBeginsWith(contentType, "video/"_ns)) {
2595       policyType = nsIContentPolicy::TYPE_MEDIA;
2596     }
2597     rv = StoreURI(aURISpec, aParentDocument, policyType);
2598   }
2599   NS_ENSURE_SUCCESS(rv, rv);
2600
2601   // Store the updated uri to the frame
2602   aData->mFile = frameURI;
2603   aData->mSubFrameExt.Truncate();  // we already put this in frameURI
2604
2605   return NS_OK;
2606 }
2607
2608 nsresult nsWebBrowserPersist::CreateChannelFromURI(nsIURI* aURI,
2609                                                    nsIChannel** aChannel) {
2610   nsresult rv = NS_OK;
2611   *aChannel = nullptr;
2612
2613   rv = NS_NewChannel(aChannel, aURI, nsContentUtils::GetSystemPrincipal(),
2614                      nsILoadInfo::SEC_ALLOW_CROSS_ORIGIN_SEC_CONTEXT_IS_NULL,
2615                      nsIContentPolicy::TYPE_OTHER);
2616   NS_ENSURE_SUCCESS(rv, rv);
2617   NS_ENSURE_ARG_POINTER(*aChannel);
2618
2619   rv = (*aChannel)->SetNotificationCallbacks(
2620       static_cast<nsIInterfaceRequestor*>(this));
2621   NS_ENSURE_SUCCESS(rv, rv);
2622   return NS_OK;
2623 }
2624
2625 // we store the current location as the key (absolutized version of domnode's
2626 // attribute's value)
2627 nsresult nsWebBrowserPersist::MakeAndStoreLocalFilenameInURIMap(
2628     nsIURI* aURI, nsIWebBrowserPersistDocument* aDoc,
2629     nsContentPolicyType aContentPolicyType, bool aNeedsPersisting,
2630     URIData** aData) {
2631   NS_ENSURE_ARG_POINTER(aURI);
2632
2633   nsAutoCString spec;
2634   nsresult rv = aURI->GetSpec(spec);
2635   NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
2636
2637   // Create a sensibly named filename for the URI and store in the URI map
2638   URIData* data;
2639   if (mURIMap.Get(spec, &data)) {
2640     if (aNeedsPersisting) {
2641       data->mNeedsPersisting = true;
2642     }
2643     if (aData) {
2644       *aData = data;
2645     }
2646     return NS_OK;
2647   }
2648
2649   // Create a unique file name for the uri
2650   nsString filename;
2651   rv = MakeFilenameFromURI(aURI, filename);
2652   NS_ENSURE_SUCCESS(rv, NS_ERROR_FAILURE);
2653
2654   // Store the file name
2655   data = new URIData;
2656
2657   data->mContentPolicyType = aContentPolicyType;
2658   data->mNeedsPersisting = aNeedsPersisting;
2659   data->mNeedsFixup = true;
2660   data->mFilename = filename;
2661   data->mSaved = false;
2662   data->mIsSubFrame = false;
2663   data->mDataPath = mCurrentDataPath;
2664   data->mDataPathIsRelative = mCurrentDataPathIsRelative;
2665   data->mRelativePathToData = mCurrentRelativePathToData;
2666   data->mRelativeDocumentURI = mTargetBaseURI;
2667   data->mCharset = mCurrentCharset;
2668
2669   aDoc->GetPrincipal(getter_AddRefs(data->mTriggeringPrincipal));
2670   aDoc->GetCookieJarSettings(getter_AddRefs(data->mCookieJarSettings));
2671
2672   if (aNeedsPersisting) mCurrentThingsToPersist++;
2673
2674   mURIMap.InsertOrUpdate(spec, UniquePtr<URIData>(data));
2675   if (aData) {
2676     *aData = data;
2677   }
2678
2679   return NS_OK;
2680 }
2681
2682 // Decide if we need to apply conversion to the passed channel.
2683 void nsWebBrowserPersist::SetApplyConversionIfNeeded(nsIChannel* aChannel) {
2684   nsresult rv = NS_OK;
2685   nsCOMPtr<nsIEncodedChannel> encChannel = do_QueryInterface(aChannel, &rv);
2686   if (NS_FAILED(rv)) return;
2687
2688   // Set the default conversion preference:
2689   encChannel->SetApplyConversion(false);
2690
2691   nsCOMPtr<nsIURI> thisURI;
2692   aChannel->GetURI(getter_AddRefs(thisURI));
2693   nsCOMPtr<nsIURL> sourceURL(do_QueryInterface(thisURI));
2694   if (!sourceURL) return;
2695   nsAutoCString extension;
2696   sourceURL->GetFileExtension(extension);
2697
2698   nsCOMPtr<nsIUTF8StringEnumerator> encEnum;
2699   encChannel->GetContentEncodings(getter_AddRefs(encEnum));
2700   if (!encEnum) return;
2701   nsCOMPtr<nsIExternalHelperAppService> helperAppService =
2702       do_GetService(NS_EXTERNALHELPERAPPSERVICE_CONTRACTID, &rv);
2703   if (NS_FAILED(rv)) return;
2704   bool hasMore;
2705   rv = encEnum->HasMore(&hasMore);
2706   if (NS_SUCCEEDED(rv) && hasMore) {
2707     nsAutoCString encType;
2708     rv = encEnum->GetNext(encType);
2709     if (NS_SUCCEEDED(rv)) {
2710       bool applyConversion = false;
2711       rv = helperAppService->ApplyDecodingForExtension(extension, encType,
2712                                                        &applyConversion);
2713       if (NS_SUCCEEDED(rv)) encChannel->SetApplyConversion(applyConversion);
2714     }
2715   }
2716 }