From 3743ec86d2f8833ad28d3be598467a4f6b516c96 Mon Sep 17 00:00:00 2001
From: Cory Barr <cory.barr@gmail.com>
Date: Tue, 13 Feb 2024 18:56:46 -0800
Subject: [PATCH] add friend methods to OmniRasterOverlay to lock down access
 to Native objects

---
 .../include/cesium/omniverse/OmniRasterOverlay.h   |  9 ++++++++-
 src/core/include/cesium/omniverse/OmniTileset.h    |  6 +++---
 src/core/src/FabricPrepareRenderResources.cpp      |  4 ++--
 src/core/src/OmniTileset.cpp                       | 14 +++++++++-----
 4 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/core/include/cesium/omniverse/OmniRasterOverlay.h b/src/core/include/cesium/omniverse/OmniRasterOverlay.h
index 0a187d9e..183c793b 100644
--- a/src/core/include/cesium/omniverse/OmniRasterOverlay.h
+++ b/src/core/include/cesium/omniverse/OmniRasterOverlay.h
@@ -1,5 +1,7 @@
 #pragma once
 
+#include "cesium/omniverse/OmniTileset.h"
+
 #include <CesiumRasterOverlays/RasterOverlay.h>
 #include <pxr/usd/sdf/path.h>
 
@@ -10,9 +12,14 @@ class RasterOverlay;
 namespace cesium::omniverse {
 
 class Context;
+class OmniTileset;
 enum class FabricOverlayRenderMethod;
 
 class OmniRasterOverlay {
+    friend void OmniTileset::addRasterOverlayIfExists(const OmniRasterOverlay* pOverlay);
+    friend pxr::SdfPath
+    OmniTileset::getRasterOverlayPathIfExists(const CesiumRasterOverlays::RasterOverlay& rasterOverlay);
+
   public:
     OmniRasterOverlay(Context* pContext, const pxr::SdfPath& path);
     virtual ~OmniRasterOverlay() = default;
@@ -32,11 +39,11 @@ class OmniRasterOverlay {
 
     [[nodiscard]] CesiumRasterOverlays::RasterOverlayOptions createRasterOverlayOptions() const;
 
-    [[nodiscard]] virtual CesiumRasterOverlays::RasterOverlay* getRasterOverlay() const = 0;
     void updateRasterOverlayOptions() const;
     virtual void reload() = 0;
 
   protected:
+    [[nodiscard]] virtual CesiumRasterOverlays::RasterOverlay* getRasterOverlay() const = 0;
     Context* _pContext;
     pxr::SdfPath _path;
 
diff --git a/src/core/include/cesium/omniverse/OmniTileset.h b/src/core/include/cesium/omniverse/OmniTileset.h
index 80bf64f1..299329e7 100644
--- a/src/core/include/cesium/omniverse/OmniTileset.h
+++ b/src/core/include/cesium/omniverse/OmniTileset.h
@@ -1,7 +1,5 @@
 #pragma once
 
-#include "cesium/omniverse/OmniTileset.h"
-
 #include <glm/glm.hpp>
 #include <pxr/usd/sdf/path.h>
 
@@ -31,6 +29,7 @@ namespace cesium::omniverse {
 
 class Context;
 class FabricPrepareRenderResources;
+class OmniRasterOverlay;
 struct TilesetStatistics;
 struct Viewport;
 
@@ -82,10 +81,11 @@ class OmniTileset {
     void updateTilesetOptions();
 
     void reload();
-    [[nodiscard]] pxr::SdfPath getRasterOverlayPath(const CesiumRasterOverlays::RasterOverlay& rasterOverlay) const;
+    [[nodiscard]] pxr::SdfPath getRasterOverlayPathIfExists(const CesiumRasterOverlays::RasterOverlay& rasterOverlay);
     void updateRasterOverlayAlpha(const pxr::SdfPath& rasterOverlayPath);
     void updateShaderInput(const pxr::SdfPath& shaderPath, const pxr::TfToken& attributeName);
     void updateDisplayColorAndOpacity();
+    void addRasterOverlayIfExists(const OmniRasterOverlay* overlay);
 
     void onUpdateFrame(const gsl::span<const Viewport>& viewports);
 
diff --git a/src/core/src/FabricPrepareRenderResources.cpp b/src/core/src/FabricPrepareRenderResources.cpp
index 60ed4ec0..dbc7e913 100644
--- a/src/core/src/FabricPrepareRenderResources.cpp
+++ b/src/core/src/FabricPrepareRenderResources.cpp
@@ -538,7 +538,7 @@ void FabricPrepareRenderResources::attachRasterInMainThread(
         return;
     }
 
-    const auto rasterOverlayPath = _pTileset->getRasterOverlayPath(rasterTile.getOverlay());
+    const auto rasterOverlayPath = _pTileset->getRasterOverlayPathIfExists(rasterTile.getOverlay());
 
     if (rasterOverlayPath.IsEmpty()) {
         return;
@@ -600,7 +600,7 @@ void FabricPrepareRenderResources::detachRasterInMainThread(
         return;
     }
 
-    const auto rasterOverlayPath = _pTileset->getRasterOverlayPath(rasterTile.getOverlay());
+    const auto rasterOverlayPath = _pTileset->getRasterOverlayPathIfExists(rasterTile.getOverlay());
 
     if (rasterOverlayPath.IsEmpty()) {
         return;
diff --git a/src/core/src/OmniTileset.cpp b/src/core/src/OmniTileset.cpp
index 4fefe9a3..e35e18fa 100644
--- a/src/core/src/OmniTileset.cpp
+++ b/src/core/src/OmniTileset.cpp
@@ -577,15 +577,12 @@ void OmniTileset::reload() {
         const auto pOmniRasterOverlay = _pContext->getAssetRegistry().getRasterOverlay(boundRasterOverlayPath);
         if (pOmniRasterOverlay) {
             pOmniRasterOverlay->reload();
-            const auto pNativeRasterOverlay = pOmniRasterOverlay->getRasterOverlay();
-            if (pNativeRasterOverlay) {
-                _pTileset->getOverlays().add(pNativeRasterOverlay);
-            }
+            addRasterOverlayIfExists(pOmniRasterOverlay);
         }
     }
 }
 
-pxr::SdfPath OmniTileset::getRasterOverlayPath(const CesiumRasterOverlays::RasterOverlay& rasterOverlay) const {
+pxr::SdfPath OmniTileset::getRasterOverlayPathIfExists(const CesiumRasterOverlays::RasterOverlay& rasterOverlay) {
     const auto rasterOverlayPaths = getRasterOverlayPaths();
 
     for (const auto& rasterOverlayPath : rasterOverlayPaths) {
@@ -787,4 +784,11 @@ void OmniTileset::destroyNativeTileset() {
     _pRenderResourcesPreparer = nullptr;
 }
 
+void OmniTileset::addRasterOverlayIfExists(const OmniRasterOverlay* pOmniRasterOverlay) {
+    const auto pNativeRasterOverlay = pOmniRasterOverlay->getRasterOverlay();
+    if (pNativeRasterOverlay) {
+        _pTileset->getOverlays().add(pNativeRasterOverlay);
+    }
+}
+
 } // namespace cesium::omniverse