fix to pass vulnerabilities check

oap-project · Apr 9, 2021 · bec76cf · bec76cf
1 parent 31b19de
commit bec76cf
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 18 deletions.
diff --git a/src/main/java/org/apache/spark/util/collection/unsafe/sort/PMemWriter.java b/src/main/java/org/apache/spark/util/collection/unsafe/sort/PMemWriter.java
@@ -169,12 +169,13 @@ private void dumpPageToPMem(MemoryBlock page) {
         MemoryBlock pMemBlock = pageMap.get(page);
         if (pMemBlock == null) {
             logger.error("Fail to copy data to PMem as no corresponding PMem page found.");
+        } else {
+            PersistentMemoryPlatform.copyMemory(
+                    page.getBaseObject(), page.getBaseOffset(),
+                    null, pMemBlock.getBaseOffset(), page.size(),
+                    pMemClflushEnabled);
+            writeMetrics.incBytesWritten(page.size());
         }
-        PersistentMemoryPlatform.copyMemory(
-            page.getBaseObject(), page.getBaseOffset(),
-            null, pMemBlock.getBaseOffset(), page.size(),
-            pMemClflushEnabled);
-        writeMetrics.incBytesWritten(page.size());
     }
 
     public void updateLongArray(LongArray sortedArray, int numRecords, int position) {
@@ -186,22 +187,27 @@ public void updateLongArray(LongArray sortedArray, int numRecords, int position)
             long baseOffset = page == null? 0: page.getBaseOffset();
             long offset = taskMemoryManager.getOffsetInPage(originalRecordPointer) - baseOffset;
             MemoryBlock pMemBlock = pageMap.get(page);
-            long pMemOffset = pMemBlock.getBaseOffset() + offset;
-            sortedArray.set(position, pMemOffset);
-            position += 2;
+            if (pMemBlock == null) {
+                logger.error("No PMem page found corresponding with LongArrayPage");
+            } else {
+                long pMemOffset = pMemBlock.getBaseOffset() + offset;
+                sortedArray.set(position, pMemOffset);
+                position += 2;
+            }
         }
         // copy the LongArray to PMem
         MemoryBlock arrayBlock = sortedArray.memoryBlock();
         MemoryBlock pMemBlock = pageMap.get(arrayBlock);
         if (pMemBlock == null) {
             logger.error("failed to update LongArray as no pair PMem page found.");
+        } else {
+            PersistentMemoryPlatform.copyMemory(
+                    arrayBlock.getBaseObject(), arrayBlock.getBaseOffset(),
+                    null, pMemBlock.getBaseOffset(), arrayBlock.size(),
+                    pMemClflushEnabled);
+            writeMetrics.incBytesWritten(pMemBlock.size());
+            this.sortedArray = new LongArray(pMemBlock);
         }
-        PersistentMemoryPlatform.copyMemory(
-            arrayBlock.getBaseObject(), arrayBlock.getBaseOffset(),
-            null, pMemBlock.getBaseOffset(), arrayBlock.size(),
-            pMemClflushEnabled);
-        writeMetrics.incBytesWritten(pMemBlock.size());
-        this.sortedArray = new LongArray(pMemBlock);
     }
 
     @Override

diff --git a/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -264,10 +264,12 @@ public SpillWriterForUnsafeSorter spillWithWriter(
         isSorted);
     if (spillWriter == null) {
       logger.error("failed to get SpillWriter, please check related configuration.");
+      return null;
+    } else {
+      spillWriter.write();
+      spillWriters.add(spillWriter);
+      return spillWriter;
     }
-    spillWriter.write();
-    spillWriters.add(spillWriter);
-    return spillWriter;
   }
 
   /**
@@ -549,7 +551,7 @@ public long spill() throws IOException {
         long released = 0L;
         SpillWriterForUnsafeSorter spillWriter = spillWithWriter(inMemIterator, numRecords, writeMetrics, true);
         nextUpstream = spillWriter.getSpillReader();
-
+        assert(nextUpstream != null);
         synchronized (UnsafeExternalSorter.this) {
           // release the pages except the one that is used. There can still be a caller that
           // is accessing the current record. We free this page in that caller's next loadNext()