Update IMG4 transform to use DecodeWithContext and synthesize a filename for the payload.

4 files changed, 40 insertions, 8 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index bd4f6de5..99ba7044 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -8929,6 +8929,7 @@ namespace BinaryNinja {
 	};
 
 	class TransformContext;
+	typedef BNTransformCapabilities TransformCapabilities;
 
 	/*! Allows users to implement custom transformations.
 
diff --git a/transform.cpp b/transform.cpp
index 735fb1b9..c2e51e80 100644
--- a/transform.cpp
+++ b/transform.cpp
@@ -114,8 +114,7 @@ bool Transform::DecodeWithContextCallback(void* ctxt, BNTransformContext* contex
 		paramMap[params[i].name] = DataBuffer(BNDuplicateDataBuffer(params[i].value));
 
 	CallbackRef<Transform> xform(ctxt);
-	Ref<TransformContext> contextRef = new TransformContext(BNNewTransformContextReference(context));
-	return xform->DecodeWithContext(contextRef, paramMap);
+	return xform->DecodeWithContext(new TransformContext(context), paramMap);
 }
 
 
diff --git a/transformcontext.cpp b/transformcontext.cpp
index 186b24b9..5bddda80 100644
--- a/transformcontext.cpp
+++ b/transformcontext.cpp
@@ -50,7 +50,7 @@ Ref<TransformContext> TransformContext::GetParent() const
 	BNTransformContext* parent = BNTransformContextGetParent(m_object);
 	if (!parent)
 		return nullptr;
-	return new TransformContext(BNNewTransformContextReference(parent));
+	return new TransformContext(parent);
 }
 
 
@@ -81,7 +81,7 @@ Ref<TransformContext> TransformContext::GetChild(const string& filename) const
 	BNTransformContext* child = BNTransformContextGetChild(m_object, filename.c_str());
 	if (!child)
 		return nullptr;
-	return new TransformContext(BNNewTransformContextReference(child));
+	return new TransformContext(child);
 }
 
 
@@ -90,7 +90,7 @@ Ref<TransformContext> TransformContext::CreateChild(const DataBuffer& data, cons
 	BNTransformContext* child = BNTransformContextCreateChild(m_object, data.GetBufferObject(), filename.c_str());
 	if (!child)
 		return nullptr;
-	return new TransformContext(BNNewTransformContextReference(child));
+	return new TransformContext(child);
 }
 
 
diff --git a/view/kernelcache/core/transformers/KernelCacheTransforms.cpp b/view/kernelcache/core/transformers/KernelCacheTransforms.cpp
index 89beda8d..6158e7ad 100644
--- a/view/kernelcache/core/transformers/KernelCacheTransforms.cpp
+++ b/view/kernelcache/core/transformers/KernelCacheTransforms.cpp
@@ -10,12 +10,13 @@ using namespace BinaryNinja;
 #include "libDER/libDER.h"
 #include "libimg4/img4.h"
 #include "liblzfse/lzfse.h"
+#include <algorithm>
 
 class IMG4PayloadTransform : public Transform
 {
 
 public:
-    IMG4PayloadTransform(): Transform(DecodeTransform, TransformSupportsDetection, "IMG4", "IMG4", "Container")
+    IMG4PayloadTransform(): Transform(DecodeTransform, TransformCapabilities(TransformSupportsDetection | TransformSupportsContext), "IMG4", "IMG4", "Container")
     {
     }
 
@@ -23,8 +24,7 @@ public:
     {
         DERItem item = {};
         item.data = (DERByte *)input.GetData();
-        item.length = input.GetLength();
-
+        item.length = static_cast<DERSize>(std::min(input.GetLength(), (size_t)std::numeric_limits<DERSize>::max()));
         Img4Payload payload = {};
         if (auto result = DERImg4DecodePayload(&item, &payload); (result != DR_Success) && (result != DR_DecodeError))
             return false;
@@ -36,6 +36,38 @@ public:
         return true;
     }
 
+    virtual bool DecodeWithContext(Ref<TransformContext> context, const std::map<std::string, DataBuffer>& params) override
+    {
+        if (!context || !context->GetInput())
+            return false;
+
+        // TODO: DataBuffers are not zero-copy, so this is inefficient. Investigate a better way.
+        auto input = context->GetInput()->ReadBuffer(0, context->GetInput()->GetLength());
+        DERItem item = {};
+        item.data = (DERByte *)input.GetData();
+        item.length = static_cast<DERSize>(std::min(input.GetLength(), (size_t)std::numeric_limits<DERSize>::max()));
+        Img4Payload payload = {};
+        if (auto result = DERImg4DecodePayload(&item, &payload); (result != DR_Success) && (result != DR_DecodeError))
+            return false;
+        if (!payload.payload.data || !payload.payload.length)
+            return false;
+
+        // Synthesize name: <type>[.<version>]
+        std::string filename = "";
+        if (payload.type.data && payload.type.length)
+            filename = std::string((const char*)payload.type.data, payload.type.length);
+        if (payload.version.data && payload.version.length)
+        {
+            if (!filename.empty())
+                filename += ".";
+            filename += std::string((const char*)payload.version.data, payload.version.length);
+        }
+
+        context->CreateChild(DataBuffer(payload.payload.data, payload.payload.length), filename);
+
+        return true;
+    }
+
     static void der_put_len(std::vector<uint8_t>& v, size_t len) {
         if (len < 0x80) { v.push_back(static_cast<uint8_t>(len)); return; }
         uint8_t tmp[9]; size_t n = 0;