Make creation of CPDFSDK_Document clearer

Move the creation of the CPDFSDK_Document into FPDFDOC_InitFormFillEnvironment
instead of hidden inside a Get method in CDPFXFA_Document.

Review-Url: https://codereview.chromium.org/2353303004

5 files changed, 30 insertions, 17 deletions

diff --git a/fpdfsdk/cpdfsdk_environment.cpp b/fpdfsdk/cpdfsdk_environment.cpp
index 2d80240366..b17b782367 100644
--- a/fpdfsdk/cpdfsdk_environment.cpp
+++ b/fpdfsdk/cpdfsdk_environment.cpp
@@ -8,6 +8,7 @@
 
 #include "fpdfsdk/formfiller/cffl_interactiveformfiller.h"
 #include "fpdfsdk/include/cpdfsdk_annothandlermgr.h"
+#include "fpdfsdk/include/cpdfsdk_document.h"
 #include "fpdfsdk/include/fsdk_actionhandler.h"
 #include "fpdfsdk/javascript/ijs_runtime.h"
 
@@ -28,9 +29,10 @@ FPDF_WIDESTRING AsFPDFWideString(CFX_ByteString* bsUTF16LE) {
 
 CPDFSDK_Environment::CPDFSDK_Environment(UnderlyingDocumentType* pDoc,
                                          FPDF_FORMFILLINFO* pFFinfo)
-    : m_pInfo(pFFinfo), m_pSDKDoc(nullptr), m_pUnderlyingDoc(pDoc) {
-  m_pSysHandler.reset(new CFX_SystemHandler(this));
-}
+    : m_pInfo(pFFinfo),
+      m_pSDKDoc(new CPDFSDK_Document(pDoc, this)),
+      m_pUnderlyingDoc(pDoc),
+      m_pSysHandler(new CFX_SystemHandler(this)) {}
 
 CPDFSDK_Environment::~CPDFSDK_Environment() {
 #ifdef PDF_ENABLE_XFA
diff --git a/fpdfsdk/fpdfformfill.cpp b/fpdfsdk/fpdfformfill.cpp
index 7e4ce36f5f..9baabc4de4 100644
--- a/fpdfsdk/fpdfformfill.cpp
+++ b/fpdfsdk/fpdfformfill.cpp
@@ -242,14 +242,22 @@ FPDFDOC_InitFormFillEnvironment(FPDF_DOCUMENT document,
   if (!pDocument)
     return nullptr;
 
+#ifdef PDF_ENABLE_XFA
+  // If the CPDFXFA_Document has a SDKDocument already then we've done this
+  // and can just return the old Env. Otherwise, we'll end up setting a new
+  // SDKDocument into the XFADocument and, that could get weird.
+  if (pDocument->GetSDKDoc())
+    return pDocument->GetSDKDoc()->GetEnv();
+#endif
+
   CPDFSDK_Environment* pEnv = new CPDFSDK_Environment(pDocument, formInfo);
+
 #ifdef PDF_ENABLE_XFA
-  pEnv->SetSDKDocument(pDocument->GetSDKDocument(pEnv));
-  CPDFXFA_App* pApp = CPDFXFA_App::GetInstance();
-  pApp->AddFormFillEnv(pEnv);
-#else  // PDF_ENABLE_XFA
-  pEnv->SetSDKDocument(new CPDFSDK_Document(pDocument, pEnv));
+  // Ownership of the SDKDocument is passed to the CPDFXFA_Document.
+  pDocument->SetSDKDoc(WrapUnique(pEnv->GetSDKDocument()));
+  CPDFXFA_App::GetInstance()->AddFormFillEnv(pEnv);
 #endif  // PDF_ENABLE_XFA
+
   return pEnv;
 }
 
@@ -257,16 +265,18 @@ DLLEXPORT void STDCALL
 FPDFDOC_ExitFormFillEnvironment(FPDF_FORMHANDLE hHandle) {
   if (!hHandle)
     return;
+
   CPDFSDK_Environment* pEnv = HandleToCPDFSDKEnvironment(hHandle);
+
 #ifdef PDF_ENABLE_XFA
-  CPDFXFA_App* pApp = CPDFXFA_App::GetInstance();
-  pApp->RemoveFormFillEnv(pEnv);
+  CPDFXFA_App::GetInstance()->RemoveFormFillEnv(pEnv);
 #else   // PDF_ENABLE_XFA
   if (CPDFSDK_Document* pSDKDoc = pEnv->GetSDKDocument()) {
     pEnv->SetSDKDocument(nullptr);
     delete pSDKDoc;
   }
 #endif  // PDF_ENABLE_XFA
+
   delete pEnv;
 }
 
diff --git a/fpdfsdk/fpdfxfa/fpdfxfa_doc.cpp b/fpdfsdk/fpdfxfa/fpdfxfa_doc.cpp
index 8db16bf4e9..2d74c02d23 100644
--- a/fpdfsdk/fpdfxfa/fpdfxfa_doc.cpp
+++ b/fpdfsdk/fpdfxfa/fpdfxfa_doc.cpp
@@ -191,11 +191,8 @@ void CPDFXFA_Document::RemovePage(CPDFXFA_Page* page) {
   m_XFAPageList.SetAt(page->GetPageIndex(), nullptr);
 }
 
-CPDFSDK_Document* CPDFXFA_Document::GetSDKDocument(
-    CPDFSDK_Environment* pFormFillEnv) {
-  if (!m_pSDKDoc && pFormFillEnv)
-    m_pSDKDoc.reset(new CPDFSDK_Document(this, pFormFillEnv));
-  return m_pSDKDoc.get();
+void CPDFXFA_Document::SetSDKDoc(std::unique_ptr<CPDFSDK_Document> pSDKDoc) {
+  m_pSDKDoc.reset(pSDKDoc.release());
 }
 
 void CPDFXFA_Document::ClearChangeMark() {
diff --git a/fpdfsdk/fpdfxfa/include/fpdfxfa_doc.h b/fpdfsdk/fpdfxfa/include/fpdfxfa_doc.h
index ed788eac55..5398c57519 100644
--- a/fpdfsdk/fpdfxfa/include/fpdfxfa_doc.h
+++ b/fpdfsdk/fpdfxfa/include/fpdfxfa_doc.h
@@ -38,9 +38,11 @@ class CPDFXFA_Document {
   CPDF_Document* GetPDFDoc() { return m_pPDFDoc.get(); }
   CXFA_FFDoc* GetXFADoc() { return m_pXFADoc.get(); }
   CXFA_FFDocView* GetXFADocView() { return m_pXFADocView; }
-  CPDFSDK_Document* GetSDKDocument(CPDFSDK_Environment* pFormFillEnv);
   int GetDocType() const { return m_iDocType; }
 
+  CPDFSDK_Document* GetSDKDoc() const { return m_pSDKDoc.get(); }
+  void SetSDKDoc(std::unique_ptr<CPDFSDK_Document> pSDKDoc);
+
   void DeletePage(int page_index);
   int GetPageCount() const;
 
@@ -54,7 +56,6 @@ class CPDFXFA_Document {
  protected:
   friend class CPDFXFA_DocEnvironment;
 
-  CPDFSDK_Document* GetSDKDoc() { return m_pSDKDoc.get(); }
   int GetOriginalPageCount() const { return m_nPageCount; }
   void SetOriginalPageCount(int count) {
     m_nPageCount = count;
diff --git a/fpdfsdk/include/cpdfsdk_environment.h b/fpdfsdk/include/cpdfsdk_environment.h
index 61916e46ad..b694359320 100644
--- a/fpdfsdk/include/cpdfsdk_environment.h
+++ b/fpdfsdk/include/cpdfsdk_environment.h
@@ -166,6 +166,9 @@ class CPDFSDK_Environment final {
   std::unique_ptr<CPDFSDK_ActionHandler> m_pActionHandler;
   std::unique_ptr<IJS_Runtime> m_pJSRuntime;
   FPDF_FORMFILLINFO* const m_pInfo;
+  // Ownership of |m_pSDKDoc| depends on if this is XFA. If we're in XFA then
+  // the object is owned by the CPDFXFA_Document. In non-xfa then we own
+  // the pointer.
   CPDFSDK_Document* m_pSDKDoc;
   UnderlyingDocumentType* const m_pUnderlyingDoc;
   std::unique_ptr<CFFL_InteractiveFormFiller> m_pFormFiller;