From 6b4d713cb075b4e01550b5842ef36355ca9d05d3 Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Mon, 3 Aug 2026 05:03:58 -0700 Subject: [PATCH 1/7] [API] Add FFI wrapper modifier to strong typedefs --- base/strong_typedef.h | 56 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/base/strong_typedef.h b/base/strong_typedef.h index 1a2db87d0..136b266dc 100644 --- a/base/strong_typedef.h +++ b/base/strong_typedef.h @@ -22,6 +22,7 @@ #include // IWYU pragma: keep #include +#include #include #include #include @@ -89,6 +90,8 @@ // Hashable usable as a key in std and absl hash containers // Formattable formattable with fmt, forwarding format specs such as {:#x} to the // underlying type's formatter +// FfiWrapper implicit conversion to and from a C-style wrapper structure through +// the specified member, without making the underlying type implicit // NonExtractable removes the explicit operator T() and the Value() accessor, so the // underlying value can be constructed but never read back out. Every // other modifier continues to function as normal. @@ -115,6 +118,25 @@ namespace detail { template inline constexpr bool HasModifier = (std::is_same_v || ...); +template +concept FfiWrapperFor = requires(const Ffi& value) { + typename Mod::FfiType; + requires std::same_as; + { Mod::template FromFfi(value) } -> std::same_as; +}; + +template +inline constexpr bool HasFfiWrapper = (FfiWrapperFor || ...); + +template +constexpr T FromFfi(const Ffi& value) +{ + if constexpr (FfiWrapperFor) + return First::template FromFfi(value); + else + return FromFfi(value); +} + // Internal access to a StrongTypedef's underlying value so that modifiers have // access to it even when NonExtractable is in use. struct Access @@ -409,6 +431,33 @@ struct Formattable }; }; +// Enables implicit conversion to and from a C-style wrapper structure whose selected +// member stores the StrongTypedef's underlying value. Conversion to the underlying type +// itself remains explicit. +template +struct FfiWrapper +{ + using FfiType = Ffi; + + template + static constexpr T FromFfi(const Ffi& value) + { + return T(value.*ValueMember); + } + + template + requires requires(Ffi ffi, const T& value) { ffi.*ValueMember = value; } + struct Apply + { + constexpr operator Ffi() const + { + Ffi result {}; + result.*ValueMember = detail::Access::Get(static_cast(*this)); + return result; + } + }; +}; + // Disable both the explicit operator T() and the Value() accessor. // All other modifiers continue to function as normal. struct NonExtractable @@ -437,6 +486,13 @@ class BN_EMPTY_BASES StrongTypedef : public Mods::template Apply + requires detail::HasFfiWrapper, T, Mods...> + constexpr StrongTypedef(Ffi&& value) + : m_value(detail::FromFfi, Mods...>(value)) + { + } + explicit constexpr operator T() const noexcept(std::is_nothrow_copy_constructible_v) requires (!detail::HasModifier) { From e9bf26207a216cc17bcc2767236a8daea4a11326 Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Mon, 3 Aug 2026 05:04:20 -0700 Subject: [PATCH 2/7] Add Binary Similarity APIs --- binaryninjacore.h | 435 +++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 433 insertions(+), 2 deletions(-) diff --git a/binaryninjacore.h b/binaryninjacore.h index b1e83870a..925c49968 100644 --- a/binaryninjacore.h +++ b/binaryninjacore.h @@ -37,14 +37,14 @@ // Current ABI version for linking to the core. This is incremented any time // there are changes to the API that affect linking, including new functions, // new types, or modifications to existing functions or types. -#define BN_CURRENT_CORE_ABI_VERSION 179 +#define BN_CURRENT_CORE_ABI_VERSION 180 // Minimum ABI version that is supported for loading of plugins. Plugins that // are linked to an ABI version less than this will not be able to load and // will require rebuilding. The minimum version is increased when there are // incompatible changes that break binary compatibility, such as changes to // existing types or functions. -#define BN_MINIMUM_CORE_ABI_VERSION 179 +#define BN_MINIMUM_CORE_ABI_VERSION 180 #define BN_DEMANGLER_MSVC "MS" #define BN_DEMANGLER_GNU3 "GNU3" @@ -359,6 +359,20 @@ extern "C" typedef struct BNConstantRenderer BNConstantRenderer; typedef struct BNStringRecognizer BNStringRecognizer; typedef struct BNCustomStringType BNCustomStringType; + typedef struct BNSimilarityProviderType BNSimilarityProviderType; + typedef struct BNSimilarityProvider BNSimilarityProvider; + typedef struct BNSimilarityProviderResults BNSimilarityProviderResults; + typedef struct BNSimilarityRenderContext BNSimilarityRenderContext; + typedef struct BNSimilarityView BNSimilarityView; + typedef struct BNDiffRenderer BNDiffRenderer; + typedef struct BNSimilaritySessionResolverType BNSimilaritySessionResolverType; + typedef struct BNSimilaritySessionResolver BNSimilaritySessionResolver; + typedef struct BNSimilaritySessionNode BNSimilaritySessionNode; + typedef struct BNSimilaritySessionGraph BNSimilaritySessionGraph; + typedef struct BNSimilaritySessionGraphReceiver BNSimilaritySessionGraphReceiver; + typedef struct BNSimilaritySessionCompletion BNSimilaritySessionCompletion; + typedef struct BNSimilaritySessionReceiver BNSimilaritySessionReceiver; + typedef struct BNSimilaritySession BNSimilaritySession; typedef struct BNVersionInfo { uint32_t major; @@ -4300,6 +4314,423 @@ extern "C" char* stringPostfix; } BNCustomStringTypeInfo; + typedef struct BNSimilarityEntityId + { + uint32_t value; + } BNSimilarityEntityId; + + typedef struct BNSimilarityResultId + { + uint64_t value; + } BNSimilarityResultId; + + typedef struct BNSimilaritySessionNodeId + { + uint32_t value; + } BNSimilaritySessionNodeId; + + typedef struct BNSimilaritySessionId + { + uint32_t value; + } BNSimilaritySessionId; + + typedef struct BNSimilarityProviderId + { + uint32_t value; + } BNSimilarityProviderId; + + typedef struct BNSimilaritySessionResolverId + { + uint32_t value; + } BNSimilaritySessionResolverId; + + typedef struct BNSimilarityEntityRef + { + BNSimilaritySessionNodeId nodeId; + BNSimilarityEntityId entityId; + } BNSimilarityEntityRef; + + BN_ENUM(uint8_t, BNSimilarityEntityType) + { + SimilarityEntityFunction = 0, + }; + + BN_ENUM(uint8_t, BNSimilarityApplyStatus) + { + SimilarityApplySuccess = 0, + SimilarityApplyNodeInactive = 1, + SimilarityApplyEntityNotFound = 2, + SimilarityApplyUnsupported = 3, + SimilarityApplyFailed = 4, + }; + + typedef struct BNSimilarityResult + { + BNSimilarityProviderId providerId; + uint8_t similarity; + uint8_t confidence; + BNSimilarityEntityRef target; + } BNSimilarityResult; + + typedef struct BNSimilarityEntityInfo + { + BNSimilarityEntityType type; + uint64_t address; + const char* name; + } BNSimilarityEntityInfo; + + BN_ENUM(uint8_t, BNSimilarityViewType) { + SimilarityViewFlowGraph = 0, + SimilarityViewLinear = 1, + }; + + BN_ENUM(uint8_t, BNSimilarityAnnotationType) { + SimilarityAnnotationAdded = 0, + SimilarityAnnotationRemoved = 1, + SimilarityAnnotationChanged = 2, + }; + + typedef struct BNSimilarityRangeAnnotation + { + uint64_t start; + uint64_t end; + BNSimilarityAnnotationType type; + } BNSimilarityRangeAnnotation; + + typedef struct BNSimilaritySessionCompletionQuery + { + bool hasNodeId; + BNSimilaritySessionNodeId nodeId; + bool hasProviderId; + BNSimilarityProviderId providerId; + bool hasResolverId; + BNSimilaritySessionResolverId resolverId; + } BNSimilaritySessionCompletionQuery; + + typedef struct BNCustomSimilarityProvider + { + void* context; + void (*externalRefTaken)(void* ctxt); + void (*externalRefReleased)(void* ctxt); + bool (*updateSettings)(void* ctxt, BNSettings* settings); + bool (*visitNode)(void* ctxt, BNSimilaritySessionNode* node, BNSimilarityProviderResults* results, + BNSimilaritySessionCompletion* completion); + bool (*visitNodeEdge)(void* ctxt, BNSimilaritySessionNode* from, BNSimilaritySessionNode* to, + BNSimilarityProviderResults* results, BNSimilaritySessionCompletion* completion); + char* (*getName)(void* ctxt, BNSimilaritySessionNode* node, BNSimilarityEntityId entity, + BNSimilarityResultId result); + BNSimilarityApplyStatus (*apply)(void* ctxt, BNSimilaritySessionNode* node, BNSimilarityEntityId entity, + BNSimilarityResultId result); + void (*render)(void* ctxt, BNSimilaritySessionNode* node, BNSimilarityEntityId entity, + BNSimilarityRenderContext* context, BNSimilarityResultId result); + void (*free)(void* ctxt); + } BNCustomSimilarityProvider; + + typedef struct BNCustomSimilarityProviderType + { + void* context; + BNSimilarityProvider* (*create)(void* ctxt, BNSettings* settings); + BNSettings* (*getDefaultSettings)(void* ctxt); + } BNCustomSimilarityProviderType; + + typedef struct BNCustomSimilaritySessionResolver + { + void* context; + void (*externalRefTaken)(void* ctxt); + void (*externalRefReleased)(void* ctxt); + bool (*updateSettings)(void* ctxt, BNSettings* settings); + void (*prepareForNode)(void* ctxt, BNSimilaritySession* session, BNSimilaritySessionNode* node, + BNSimilaritySessionCompletion* completion, BNSimilaritySessionResolverId resolverId); + void (*resolveForNode)(void* ctxt, BNSimilaritySession* session, BNSimilaritySessionNode* node, + const BNSimilarityEntityId* entities, size_t entityCount, BNSimilaritySessionCompletion* completion, + BNSimilaritySessionResolverId resolverId); + void (*free)(void* ctxt); + } BNCustomSimilaritySessionResolver; + + typedef struct BNCustomSimilaritySessionResolverType + { + void* context; + BNSimilaritySessionResolver* (*create)(void* ctxt, BNSimilaritySession* session, BNSettings* settings); + BNSettings* (*getDefaultSettings)(void* ctxt); + } BNCustomSimilaritySessionResolverType; + + typedef struct BNCustomSimilaritySessionReceiver + { + void* context; + void (*externalRefTaken)(void* ctxt); + void (*externalRefReleased)(void* ctxt); + void (*onStarted)(void* ctxt, BNSimilaritySessionCompletion* completion); + void (*onUpdated)(void* ctxt, BNSimilaritySessionNode* node, BNSimilarityProvider* provider, + const BNSimilarityEntityId* entities, size_t count); + void (*free)(void* ctxt); + } BNCustomSimilaritySessionReceiver; + + typedef struct BNCustomSimilaritySessionGraphReceiver + { + void* context; + void (*externalRefTaken)(void* ctxt); + void (*externalRefReleased)(void* ctxt); + void (*onGraphChanged)(void* ctxt); + void (*free)(void* ctxt); + } BNCustomSimilaritySessionGraphReceiver; + + BINARYNINJACOREAPI BNSimilarityProviderType* BNRegisterSimilarityProviderType( + const char* name, const char* description, BNCustomSimilarityProviderType* type); + BINARYNINJACOREAPI BNSimilarityProviderType* BNGetSimilarityProviderTypeByName(const char* name); + BINARYNINJACOREAPI BNSimilarityProviderType** BNGetSimilarityProviderTypeList(size_t* count); + BINARYNINJACOREAPI void BNFreeSimilarityProviderTypeList(BNSimilarityProviderType** types); + BINARYNINJACOREAPI char* BNSimilarityProviderTypeGetName(BNSimilarityProviderType* type); + BINARYNINJACOREAPI char* BNSimilarityProviderTypeGetDescription(BNSimilarityProviderType* type); + /*! Returns `nullptr` outside the Ultimate edition. */ + BINARYNINJACOREAPI BNSimilarityProvider* BNSimilarityProviderTypeCreateProvider( + BNSimilarityProviderType* type, BNSettings* settings); + BINARYNINJACOREAPI BNSettings* BNSimilarityProviderTypeGetDefaultSettings(BNSimilarityProviderType* type); + + BINARYNINJACOREAPI BNSimilarityProvider* BNCreateCustomSimilarityProvider( + BNSimilarityProviderType* type, BNCustomSimilarityProvider* callbacks); + BINARYNINJACOREAPI void BNSimilarityProviderVisitNode( + BNSimilarityProvider* provider, BNSimilaritySessionNode* node, BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI void BNSimilarityProviderVisitNodeEdge(BNSimilarityProvider* provider, + BNSimilaritySessionNode* from, BNSimilaritySessionNode* to, BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI bool BNSimilarityProviderPerformVisitNode(BNSimilarityProvider* provider, + BNSimilaritySessionNode* node, BNSimilarityProviderResults* results, BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI bool BNSimilarityProviderPerformVisitNodeEdge(BNSimilarityProvider* provider, + BNSimilaritySessionNode* from, BNSimilaritySessionNode* to, BNSimilarityProviderResults* results, + BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI BNSimilarityResultId BNSimilarityProviderResultsAddResult(BNSimilarityProviderResults* results, + const BNSimilarityEntityRef* source, const BNSimilarityEntityRef* target, + uint8_t similarity, uint8_t confidence); + BINARYNINJACOREAPI BNSimilarityProviderType* BNSimilarityProviderGetType(BNSimilarityProvider* provider); + BINARYNINJACOREAPI BNSimilarityProviderId BNSimilarityProviderGetId(BNSimilarityProvider* provider); + BINARYNINJACOREAPI char* BNSimilarityProviderGetName(BNSimilarityProvider* provider, + BNSimilaritySessionNode* node, BNSimilarityEntityId entity, BNSimilarityResultId result); + BINARYNINJACOREAPI BNSimilarityApplyStatus BNSimilarityProviderApply(BNSimilarityProvider* provider, + BNSimilaritySessionNode* node, BNSimilarityEntityId entity, BNSimilarityResultId result); + BINARYNINJACOREAPI BNSimilarityApplyStatus BNSimilaritySessionNodeApplyTarget(BNSimilaritySessionNode* node, + BNSimilarityEntityId entity, const BNSimilarityEntityRef* target); + BINARYNINJACOREAPI void BNSimilarityProviderRender( + BNSimilarityProvider* provider, BNSimilaritySessionNode* node, BNSimilarityEntityId entity, + BNSimilarityRenderContext* context, BNSimilarityResultId result); + + BINARYNINJACOREAPI BNSimilarityRenderContext* BNCreateSimilarityRenderContext(void); + BINARYNINJACOREAPI void BNSimilarityRenderContextSetPreferredViewType( + BNSimilarityRenderContext* context, BNFunctionViewType type); + BINARYNINJACOREAPI BNFunctionGraphType BNSimilarityRenderContextGetPreferredViewType( + BNSimilarityRenderContext* context); + BINARYNINJACOREAPI char* BNSimilarityRenderContextGetPreferredViewTypeName(BNSimilarityRenderContext* context); + BINARYNINJACOREAPI void BNSimilarityRenderContextAddFlowGraph( + BNSimilarityRenderContext* context, const char* group, BNFlowGraph* graph); + BINARYNINJACOREAPI void BNSimilarityRenderContextAddFlowGraphForEntity( + BNSimilarityRenderContext* context, const char* group, BNFlowGraph* graph, const BNSimilarityEntityRef* entity); + BINARYNINJACOREAPI void BNSimilarityRenderContextAddLinearView( + BNSimilarityRenderContext* context, const char* group, BNBinaryView* data, BNLinearViewObject* linearView); + BINARYNINJACOREAPI void BNSimilarityRenderContextAddLinearViewForEntity(BNSimilarityRenderContext* context, + const char* group, BNBinaryView* data, BNLinearViewObject* linearView, const BNSimilarityEntityRef* entity); + BINARYNINJACOREAPI BNSimilarityView** BNGetSimilarityRenderContextViews( + BNSimilarityRenderContext* context, size_t* count); + BINARYNINJACOREAPI void BNFreeSimilarityViewList(BNSimilarityView** views, size_t count); + BINARYNINJACOREAPI char* BNSimilarityViewGetGroup(BNSimilarityView* view); + BINARYNINJACOREAPI BNSimilarityViewType BNSimilarityViewGetType(BNSimilarityView* view); + BINARYNINJACOREAPI BNFlowGraph* BNSimilarityViewGetFlowGraph(BNSimilarityView* view); + BINARYNINJACOREAPI BNBinaryView* BNSimilarityViewGetLinearViewData(BNSimilarityView* view); + BINARYNINJACOREAPI BNLinearViewObject* BNSimilarityViewGetLinearView(BNSimilarityView* view); + BINARYNINJACOREAPI bool BNSimilarityViewGetEntity(BNSimilarityView* view, BNSimilarityEntityRef* entity); + BINARYNINJACOREAPI BNSimilarityView* BNNewSimilarityViewReference(BNSimilarityView* view); + BINARYNINJACOREAPI void BNFreeSimilarityView(BNSimilarityView* view); + BINARYNINJACOREAPI BNSimilarityRenderContext* BNNewSimilarityRenderContextReference( + BNSimilarityRenderContext* context); + BINARYNINJACOREAPI void BNFreeSimilarityRenderContext(BNSimilarityRenderContext* context); + BINARYNINJACOREAPI BNDiffRenderer* BNCreateDiffRenderer(void); + BINARYNINJACOREAPI void BNDiffRendererAddRangeAnnotation( + BNDiffRenderer* renderer, uint64_t start, uint64_t end, BNSimilarityAnnotationType type); + BINARYNINJACOREAPI void BNDiffRendererRenderFunction( + BNDiffRenderer* renderer, BNSimilarityRenderContext* context, BNFunction* function); + BINARYNINJACOREAPI void BNDiffRendererRenderFunctionForEntity(BNDiffRenderer* renderer, + BNSimilarityRenderContext* context, BNFunction* function, const BNSimilarityEntityRef* entity); + BINARYNINJACOREAPI void BNDiffRendererRenderFlowGraph( + BNDiffRenderer* renderer, BNSimilarityRenderContext* context, const char* group, BNFlowGraph* graph); + BINARYNINJACOREAPI void BNDiffRendererRenderFlowGraphForEntity(BNDiffRenderer* renderer, + BNSimilarityRenderContext* context, const char* group, BNFlowGraph* graph, const BNSimilarityEntityRef* entity); + BINARYNINJACOREAPI void BNDiffRendererRenderLinearView(BNDiffRenderer* renderer, BNSimilarityRenderContext* context, + const char* group, BNBinaryView* data, BNLinearViewObject* linearView); + BINARYNINJACOREAPI void BNDiffRendererRenderLinearViewForEntity(BNDiffRenderer* renderer, + BNSimilarityRenderContext* context, const char* group, BNBinaryView* data, BNLinearViewObject* linearView, + const BNSimilarityEntityRef* entity); + BINARYNINJACOREAPI BNDiffRenderer* BNNewDiffRendererReference(BNDiffRenderer* renderer); + BINARYNINJACOREAPI void BNFreeDiffRenderer(BNDiffRenderer* renderer); + BINARYNINJACOREAPI BNSimilarityProvider* BNNewSimilarityProviderReference(BNSimilarityProvider* provider); + BINARYNINJACOREAPI void BNFreeSimilarityProvider(BNSimilarityProvider* provider); + BINARYNINJACOREAPI void BNFreeSimilarityResultIdList(BNSimilarityResultId* results); + + BINARYNINJACOREAPI BNSimilaritySessionResolverType* BNRegisterSimilaritySessionResolverType( + const char* name, const char* description, BNCustomSimilaritySessionResolverType* type); + BINARYNINJACOREAPI BNSimilaritySessionResolverType* BNGetSimilaritySessionResolverTypeByName(const char* name); + BINARYNINJACOREAPI BNSimilaritySessionResolverType** BNGetSimilaritySessionResolverTypeList(size_t* count); + BINARYNINJACOREAPI void BNFreeSimilaritySessionResolverTypeList(BNSimilaritySessionResolverType** types); + BINARYNINJACOREAPI char* BNSimilaritySessionResolverTypeGetName(BNSimilaritySessionResolverType* type); + BINARYNINJACOREAPI char* BNSimilaritySessionResolverTypeGetDescription(BNSimilaritySessionResolverType* type); + BINARYNINJACOREAPI BNSimilaritySessionResolver* BNSimilaritySessionResolverTypeCreateResolver( + BNSimilaritySessionResolverType* type, BNSimilaritySession* session, BNSettings* settings); + BINARYNINJACOREAPI BNSettings* BNSimilaritySessionResolverTypeGetDefaultSettings( + BNSimilaritySessionResolverType* type); + + BINARYNINJACOREAPI BNSimilaritySessionResolver* BNCreateCustomSimilaritySessionResolver( + BNSimilaritySessionResolverType* type, BNSimilaritySession* session, + BNCustomSimilaritySessionResolver* callbacks); + BINARYNINJACOREAPI BNSimilaritySessionResolverType* BNSimilaritySessionResolverGetType( + BNSimilaritySessionResolver* resolver); + BINARYNINJACOREAPI BNSimilaritySessionResolverId BNSimilaritySessionResolverGetId( + BNSimilaritySessionResolver* resolver); + BINARYNINJACOREAPI void BNSimilaritySessionResolverPrepareForNode( + BNSimilaritySessionResolver* resolver, BNSimilaritySession* session, BNSimilaritySessionNode* node, + BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI void BNSimilaritySessionResolverResolveForNode(BNSimilaritySessionResolver* resolver, + BNSimilaritySession* session, BNSimilaritySessionNode* node, const BNSimilarityEntityId* entities, + size_t entityCount, BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI BNSimilaritySessionResolver* BNNewSimilaritySessionResolverReference(BNSimilaritySessionResolver* resolver); + BINARYNINJACOREAPI void BNFreeSimilaritySessionResolver(BNSimilaritySessionResolver* resolver); + BINARYNINJACOREAPI void BNFreeSimilaritySessionResolverList(BNSimilaritySessionResolver** resolvers, size_t count); + + BINARYNINJACOREAPI BNSimilaritySessionNode* BNCreateSimilaritySessionNode(BNBinaryView* view); + BINARYNINJACOREAPI BNSimilaritySessionNode* BNCreateSimilaritySessionNodeFromFile(BNFileMetadata* file); + BINARYNINJACOREAPI BNBinaryView* BNSimilaritySessionNodeGetView(BNSimilaritySessionNode* node); + BINARYNINJACOREAPI void BNSimilaritySessionNodeSetView(BNSimilaritySessionNode* node, BNBinaryView* view); + BINARYNINJACOREAPI BNFileMetadata* BNSimilaritySessionNodeGetFile(BNSimilaritySessionNode* node); + BINARYNINJACOREAPI BNSettings* BNSimilaritySessionNodeGetLoadOptions(BNSimilaritySessionNode* node); + BINARYNINJACOREAPI BNSimilaritySessionNodeId BNSimilaritySessionNodeGetId(BNSimilaritySessionNode* node); + BINARYNINJACOREAPI BNSimilarityEntityId BNSimilaritySessionNodeCreateEntity( + BNSimilaritySessionNode* node, const BNSimilarityEntityInfo* info); + BINARYNINJACOREAPI bool BNSimilaritySessionNodeRemoveEntity( + BNSimilaritySessionNode* node, BNSimilarityEntityId id); + BINARYNINJACOREAPI bool BNSimilaritySessionNodeGetEntity( + BNSimilaritySessionNode* node, BNSimilarityEntityId id, BNSimilarityEntityInfo* result); + BINARYNINJACOREAPI void BNFreeSimilarityEntityInfo(BNSimilarityEntityInfo* info); + BINARYNINJACOREAPI BNSimilarityEntityId* BNSimilaritySessionNodeGetEntities( + BNSimilaritySessionNode* node, size_t* count); + BINARYNINJACOREAPI bool BNSimilaritySessionNodeAddScheduledEntity( + BNSimilaritySessionNode* node, BNSimilarityEntityId id); + BINARYNINJACOREAPI bool BNSimilaritySessionNodeRemoveScheduledEntity( + BNSimilaritySessionNode* node, BNSimilarityEntityId id); + BINARYNINJACOREAPI BNSimilarityEntityId* BNSimilaritySessionNodeGetScheduledEntities( + BNSimilaritySessionNode* node, size_t* count); + BINARYNINJACOREAPI BNFunction* BNSimilaritySessionNodeGetEntityFunction( + BNSimilaritySessionNode* node, BNSimilarityEntityId id); + BINARYNINJACOREAPI BNSimilarityResultId* BNSimilaritySessionNodeGetResults( + BNSimilaritySessionNode* node, BNSimilarityEntityId entity, size_t* count); + BINARYNINJACOREAPI bool BNSimilaritySessionNodeGetResult( + BNSimilaritySessionNode* node, BNSimilarityResultId resultId, BNSimilarityResult* result); + BINARYNINJACOREAPI bool BNSimilaritySessionNodeSetResolvedResult( + BNSimilaritySessionNode* node, BNSimilarityEntityId entity, BNSimilarityResultId result); + BINARYNINJACOREAPI bool BNSimilaritySessionNodeGetResolvedResult( + BNSimilaritySessionNode* node, BNSimilarityEntityId entity, BNSimilarityResultId* result); + BINARYNINJACOREAPI bool BNSimilaritySessionNodeClearResolvedResult( + BNSimilaritySessionNode* node, BNSimilarityEntityId entity); + BINARYNINJACOREAPI void BNFreeSimilarityEntityList(BNSimilarityEntityId* entities); + BINARYNINJACOREAPI BNSimilaritySessionNodeId* BNSimilaritySessionNodeGetIncomingEdges(BNSimilaritySessionNode* node, size_t* count); + BINARYNINJACOREAPI BNSimilaritySessionNodeId* BNSimilaritySessionNodeGetOutgoingEdges(BNSimilaritySessionNode* node, size_t* count); + BINARYNINJACOREAPI void BNFreeSimilaritySessionNodeEdgeList(BNSimilaritySessionNodeId* edges); + BINARYNINJACOREAPI BNSimilaritySessionNode** BNSimilaritySessionNodeGetIncomingNodes( + BNSimilaritySessionNode* node, size_t* count); + BINARYNINJACOREAPI BNSimilaritySessionNode** BNSimilaritySessionNodeGetOutgoingNodes( + BNSimilaritySessionNode* node, size_t* count); + BINARYNINJACOREAPI BNSimilaritySessionNode* BNNewSimilaritySessionNodeReference(BNSimilaritySessionNode* node); + BINARYNINJACOREAPI void BNFreeSimilaritySessionNode(BNSimilaritySessionNode* node); + + BINARYNINJACOREAPI void BNSimilaritySessionGraphAddNode( + BNSimilaritySessionGraph* graph, BNSimilaritySessionNode* node); + BINARYNINJACOREAPI void BNSimilaritySessionGraphRemoveNode( + BNSimilaritySessionGraph* graph, BNSimilaritySessionNode* node); + BINARYNINJACOREAPI bool BNSimilaritySessionGraphIsValidEdge( + BNSimilaritySessionGraph* graph, BNSimilaritySessionNode* from, BNSimilaritySessionNode* to); + BINARYNINJACOREAPI bool BNSimilaritySessionGraphAddEdge( + BNSimilaritySessionGraph* graph, BNSimilaritySessionNode* from, BNSimilaritySessionNode* to); + BINARYNINJACOREAPI bool BNSimilaritySessionGraphRemoveEdge( + BNSimilaritySessionGraph* graph, BNSimilaritySessionNode* from, BNSimilaritySessionNode* to); + BINARYNINJACOREAPI BNSimilaritySessionNode* BNSimilaritySessionGraphGetNode( + BNSimilaritySessionGraph* graph, BNSimilaritySessionNodeId id); + BINARYNINJACOREAPI BNSimilaritySessionNode** BNSimilaritySessionGraphGetNodes( + BNSimilaritySessionGraph* graph, size_t* count); + BINARYNINJACOREAPI void BNFreeSimilaritySessionNodeList(BNSimilaritySessionNode** nodes, size_t count); + BINARYNINJACOREAPI BNSimilaritySessionNode*** BNSimilaritySessionGraphGetSchedule( + BNSimilaritySessionGraph* graph, size_t** nodeCounts, size_t* levelCount); + BINARYNINJACOREAPI void BNFreeSimilaritySessionNodeSchedule( + BNSimilaritySessionNode*** schedule, size_t* nodeCounts, size_t levelCount); + BINARYNINJACOREAPI BNSimilaritySessionGraph* BNNewSimilaritySessionGraphReference(BNSimilaritySessionGraph* node); + BINARYNINJACOREAPI void BNFreeSimilaritySessionGraph(BNSimilaritySessionGraph* graph); + BINARYNINJACOREAPI BNSimilaritySessionGraphReceiver* BNCreateCustomSimilaritySessionGraphReceiver( + BNCustomSimilaritySessionGraphReceiver* callbacks); + BINARYNINJACOREAPI BNSimilaritySessionGraphReceiver* BNNewSimilaritySessionGraphReceiverReference( + BNSimilaritySessionGraphReceiver* receiver); + BINARYNINJACOREAPI void BNFreeSimilaritySessionGraphReceiver(BNSimilaritySessionGraphReceiver* receiver); + BINARYNINJACOREAPI void BNFreeSimilaritySessionGraphReceiverList( + BNSimilaritySessionGraphReceiver** receivers, size_t count); + BINARYNINJACOREAPI void BNSimilaritySessionGraphReceiverNotifyGraphChanged( + BNSimilaritySessionGraphReceiver* receiver); + BINARYNINJACOREAPI void BNSimilaritySessionGraphAddReceiver( + BNSimilaritySessionGraph* graph, BNSimilaritySessionGraphReceiver* receiver); + BINARYNINJACOREAPI void BNSimilaritySessionGraphRemoveReceiver( + BNSimilaritySessionGraph* graph, BNSimilaritySessionGraphReceiver* receiver); + BINARYNINJACOREAPI BNSimilaritySessionGraphReceiver** BNSimilaritySessionGraphGetReceivers( + BNSimilaritySessionGraph* graph, size_t* count); + + BINARYNINJACOREAPI BNSimilaritySession* BNCreateSimilaritySession(); + BINARYNINJACOREAPI BNSimilaritySessionId BNSimilaritySessionGetId(BNSimilaritySession* session); + BINARYNINJACOREAPI BNSimilaritySessionReceiver* BNCreateCustomSimilaritySessionReceiver( + BNCustomSimilaritySessionReceiver* callbacks); + BINARYNINJACOREAPI BNSimilaritySessionReceiver* BNNewSimilaritySessionReceiverReference( + BNSimilaritySessionReceiver* receiver); + BINARYNINJACOREAPI void BNFreeSimilaritySessionReceiver(BNSimilaritySessionReceiver* receiver); + BINARYNINJACOREAPI void BNFreeSimilaritySessionReceiverList(BNSimilaritySessionReceiver** receivers, size_t count); + BINARYNINJACOREAPI void BNSimilaritySessionReceiverNotifyStart( + BNSimilaritySessionReceiver* receiver, BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI void BNSimilaritySessionReceiverNotifyBatch(BNSimilaritySessionReceiver* receiver, + BNSimilaritySessionNode* node, BNSimilarityProvider* provider, + const BNSimilarityEntityId* entities, size_t count); + BINARYNINJACOREAPI void BNSimilaritySessionAddProvider( + BNSimilaritySession* session, BNSimilarityProvider* provider); + BINARYNINJACOREAPI void BNSimilaritySessionRemoveProvider( + BNSimilaritySession* session, BNSimilarityProvider* provider); + BINARYNINJACOREAPI bool BNSimilaritySessionUpdateProviderSettings( + BNSimilaritySession* session, BNSimilarityProvider* provider, BNSettings* settings); + BINARYNINJACOREAPI BNSimilarityProvider* BNSimilaritySessionGetProvider( + BNSimilaritySession* session, BNSimilarityProviderId id); + BINARYNINJACOREAPI BNSimilarityProvider** BNSimilaritySessionGetProviders( + BNSimilaritySession* session, size_t* count); + BINARYNINJACOREAPI void BNFreeSimilarityProviderList(BNSimilarityProvider** providers, size_t count); + BINARYNINJACOREAPI bool BNSimilaritySessionAddResolver( + BNSimilaritySession* session, BNSimilaritySessionResolver* resolver); + BINARYNINJACOREAPI bool BNSimilaritySessionRemoveResolver( + BNSimilaritySession* session, BNSimilaritySessionResolver* resolver); + BINARYNINJACOREAPI bool BNSimilaritySessionUpdateResolverSettings( + BNSimilaritySession* session, BNSimilaritySessionResolver* resolver, BNSettings* settings); + BINARYNINJACOREAPI BNSimilaritySessionResolver* BNSimilaritySessionGetResolver( + BNSimilaritySession* session, BNSimilaritySessionResolverId id); + BINARYNINJACOREAPI BNSimilaritySessionResolver** BNSimilaritySessionGetResolvers( + BNSimilaritySession* session, size_t* count); + BINARYNINJACOREAPI void BNSimilaritySessionAddReceiver( + BNSimilaritySession* session, BNSimilaritySessionReceiver* receiver); + BINARYNINJACOREAPI void BNSimilaritySessionRemoveReceiver( + BNSimilaritySession* session, BNSimilaritySessionReceiver* receiver); + BINARYNINJACOREAPI BNSimilaritySessionReceiver** BNSimilaritySessionGetReceivers( + BNSimilaritySession* session, size_t* count); + BINARYNINJACOREAPI BNSimilaritySessionGraph* BNSimilaritySessionGetGraph(BNSimilaritySession* session); + BINARYNINJACOREAPI BNSimilaritySessionCompletion* BNSimilaritySessionRun(BNSimilaritySession* session); + BINARYNINJACOREAPI BNSimilaritySession* BNNewSimilaritySessionReference(BNSimilaritySession* session); + BINARYNINJACOREAPI void BNFreeSimilaritySession(BNSimilaritySession* session); + + BINARYNINJACOREAPI BNSimilaritySessionCompletion* BNCreateSimilaritySessionCompletion(); + BINARYNINJACOREAPI bool BNSimilaritySessionCompletionIsFinished(BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI void BNSimilaritySessionCompletionRequestStop(BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI bool BNSimilaritySessionCompletionIsStopRequested(BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI double BNSimilaritySessionCompletionGetProgress( + BNSimilaritySessionCompletion* completion, const BNSimilaritySessionCompletionQuery* query); + BINARYNINJACOREAPI void BNSimilaritySessionCompletionSetProgress( + BNSimilaritySessionCompletion* completion, const BNSimilaritySessionCompletionQuery* query, double progress); + BINARYNINJACOREAPI uint64_t BNSimilaritySessionCompletionGetTiming( + BNSimilaritySessionCompletion* completion, const BNSimilaritySessionCompletionQuery* query); + BINARYNINJACOREAPI BNSimilaritySessionCompletion* BNNewSimilaritySessionCompletionReference(BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI void BNFreeSimilaritySessionCompletion(BNSimilaritySessionCompletion* completion); + BINARYNINJACOREAPI char* BNAllocString(const char* contents); BINARYNINJACOREAPI char* BNAllocStringWithLength(const char* contents, size_t len); BINARYNINJACOREAPI void BNFreeString(char* str); From dc8e3512c897d525b1285f8b73a34a7d911f1f1c Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Mon, 3 Aug 2026 05:04:20 -0700 Subject: [PATCH 3/7] [Rust] Add FFI conversions to ID types --- rust/src/ffi.rs | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/rust/src/ffi.rs b/rust/src/ffi.rs index 66e5f70f1..4ffc89836 100644 --- a/rust/src/ffi.rs +++ b/rust/src/ffi.rs @@ -59,7 +59,8 @@ macro_rules! ffi_span { } macro_rules! new_id_type { - ($name:ident, $inner_type:ty) => { + ($(#[$meta:meta])* $name:ident, $inner_type:ty $(, $ffi_type:ty, $ffi_field:ident)?) => { + $(#[$meta])* #[derive(std::fmt::Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] pub struct $name(pub $inner_type); @@ -80,5 +81,19 @@ macro_rules! new_id_type { write!(f, "{}", self.0) } } + + $( + impl From<$ffi_type> for $name { + fn from(value: $ffi_type) -> Self { + Self(value.$ffi_field) + } + } + + impl From<$name> for $ffi_type { + fn from(value: $name) -> Self { + Self { $ffi_field: value.0 } + } + } + )? }; } From 73787d7bad34d01344bbe9390f3412ecff7d31d8 Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Mon, 3 Aug 2026 05:04:20 -0700 Subject: [PATCH 4/7] [Rust] Support owned settings handles --- rust/src/settings.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 4e35f20f6..1693b5a35 100644 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -35,6 +35,10 @@ pub struct Settings { } impl Settings { + pub(crate) unsafe fn from_raw(handle: *mut BNSettings) -> Self { + Self { handle } + } + pub(crate) unsafe fn ref_from_raw(handle: *mut BNSettings) -> Ref { debug_assert!(!handle.is_null()); Ref::new(Self { handle }) @@ -517,11 +521,11 @@ impl Settings { unsafe { BNSettingsRegisterGroup(self.handle, group.as_ptr(), title.as_ptr()) } } - pub fn register_setting_json(&self, group: &str, properties: &str) -> bool { - let group = group.to_cstr(); + pub fn register_setting_json(&self, key: &str, properties: &str) -> bool { + let key = key.to_cstr(); let properties = properties.to_cstr(); - unsafe { BNSettingsRegisterSetting(self.handle, group.as_ptr(), properties.as_ptr()) } + unsafe { BNSettingsRegisterSetting(self.handle, key.as_ptr(), properties.as_ptr()) } } // TODO: register_setting but type-safely turn it into json From 7d5a58559ab3648a4a7adb982b50d70890cff61d Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Mon, 3 Aug 2026 05:04:20 -0700 Subject: [PATCH 5/7] Add Binary Similarity APIs --- rust/src/lib.rs | 1 + rust/src/similarity.rs | 196 ++++++++ rust/src/similarity/graph.rs | 218 +++++++++ rust/src/similarity/node.rs | 324 +++++++++++++ rust/src/similarity/provider.rs | 589 ++++++++++++++++++++++++ rust/src/similarity/render.rs | 329 +++++++++++++ rust/src/similarity/session.rs | 267 +++++++++++ rust/src/similarity/session/receiver.rs | 152 ++++++ rust/src/similarity/session/resolver.rs | 393 ++++++++++++++++ rust/tests/similarity.rs | 453 ++++++++++++++++++ 10 files changed, 2922 insertions(+) create mode 100644 rust/src/similarity.rs create mode 100644 rust/src/similarity/graph.rs create mode 100644 rust/src/similarity/node.rs create mode 100644 rust/src/similarity/provider.rs create mode 100644 rust/src/similarity/render.rs create mode 100644 rust/src/similarity/session.rs create mode 100644 rust/src/similarity/session/receiver.rs create mode 100644 rust/src/similarity/session/resolver.rs create mode 100644 rust/tests/similarity.rs diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 295f7474c..3e3d1f88c 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -81,6 +81,7 @@ pub mod secrets_provider; pub mod section; pub mod segment; pub mod settings; +pub mod similarity; pub mod string; pub mod string_detection; pub mod symbol; diff --git a/rust/src/similarity.rs b/rust/src/similarity.rs new file mode 100644 index 000000000..ceb638821 --- /dev/null +++ b/rust/src/similarity.rs @@ -0,0 +1,196 @@ +//! Function similarity providers, sessions, and result rendering. + +use binaryninjacore_sys::{ + BNSimilarityAnnotationType, BNSimilarityApplyStatus, BNSimilarityEntityId, + BNSimilarityEntityRef, BNSimilarityEntityType, BNSimilarityProviderId, BNSimilarityResultId, + BNSimilaritySessionCompletionQuery, BNSimilaritySessionId, BNSimilaritySessionNodeId, + BNSimilaritySessionResolverId, BNSimilarityViewType, +}; + +pub mod graph; +pub mod node; +pub mod provider; +pub mod render; +pub mod session; + +pub use graph::*; +pub use node::*; +pub use provider::*; +pub use render::*; +pub use session::*; + +/// The kind of object represented by a similarity entity. +pub type SimilarityEntityType = BNSimilarityEntityType; + +/// The result of applying a similarity match. +pub type SimilarityApplyStatus = BNSimilarityApplyStatus; + +/// The kind of view produced when rendering a result. +pub type SimilarityViewType = BNSimilarityViewType; + +/// The change represented by a rendered address range. +pub type SimilarityAnnotationType = BNSimilarityAnnotationType; + +new_id_type!( + /// Identifies an entity within a similarity session node. + SimilarityEntityId, + u32, + BNSimilarityEntityId, + value +); + +new_id_type!( + /// Identifies a result within a similarity session node. + SimilarityResultId, + u64, + BNSimilarityResultId, + value +); + +new_id_type!( + /// Identifies a similarity session node. + SimilaritySessionNodeId, + u32, + BNSimilaritySessionNodeId, + value +); + +new_id_type!( + /// Identifies a similarity session. + SimilaritySessionId, + u32, + BNSimilaritySessionId, + value +); + +new_id_type!( + /// Identifies a similarity provider instance. + SimilarityProviderId, + u32, + BNSimilarityProviderId, + value +); + +new_id_type!( + /// Identifies a similarity resolver instance. + SimilaritySessionResolverId, + u32, + BNSimilaritySessionResolverId, + value +); + +/// Chooses which similarity session completion data to read or update. +/// +/// A query cannot select both a provider and a resolver. An empty query selects the whole session. +#[derive(Debug, Copy, Clone, Default, PartialEq, Eq)] +pub struct SimilaritySessionCompletionQuery { + node_id: Option, + provider_id: Option, + resolver_id: Option, +} + +impl SimilaritySessionCompletionQuery { + /// Selects the whole session. + pub fn for_session() -> Self { + Self::default() + } + /// Selects a node. + pub fn for_node(node_id: SimilaritySessionNodeId) -> Self { + Self { + node_id: Some(node_id), + ..Self::default() + } + } + /// Selects a provider across the session. + pub fn for_provider(provider_id: SimilarityProviderId) -> Self { + Self { + provider_id: Some(provider_id), + ..Self::default() + } + } + /// Selects a resolver across the session. + pub fn for_resolver(resolver_id: SimilaritySessionResolverId) -> Self { + Self { + resolver_id: Some(resolver_id), + ..Self::default() + } + } + /// Selects a provider within the current selection. + pub fn with_provider(mut self, provider_id: SimilarityProviderId) -> Self { + self.provider_id = Some(provider_id); + self.resolver_id = None; + self + } + /// Selects a resolver within the current selection. + pub fn with_resolver(mut self, resolver_id: SimilaritySessionResolverId) -> Self { + self.provider_id = None; + self.resolver_id = Some(resolver_id); + self + } + + /// Returns the selected node, if any. + pub fn node_id(&self) -> Option { + self.node_id + } + + /// Returns the selected provider, if any. + pub fn provider_id(&self) -> Option { + self.provider_id + } + + /// Returns the selected resolver, if any. + pub fn resolver_id(&self) -> Option { + self.resolver_id + } +} + +impl From for BNSimilaritySessionCompletionQuery { + fn from(value: SimilaritySessionCompletionQuery) -> Self { + Self { + hasNodeId: value.node_id.is_some(), + nodeId: value.node_id.unwrap_or(SimilaritySessionNodeId(0)).into(), + hasProviderId: value.provider_id.is_some(), + providerId: value.provider_id.unwrap_or(SimilarityProviderId(0)).into(), + hasResolverId: value.resolver_id.is_some(), + resolverId: value + .resolver_id + .unwrap_or(SimilaritySessionResolverId(0)) + .into(), + } + } +} + +/// Identifies an entity within a session node. +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +pub struct SimilarityEntityRef { + pub node_id: SimilaritySessionNodeId, + pub entity_id: SimilarityEntityId, +} + +/// Information about an entity. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct SimilarityEntityInfo { + pub entity_type: SimilarityEntityType, + /// The address of the entity within its [`crate::binary_view::BinaryView`]. + pub address: u64, + /// The display name of the entity. + pub name: String, +} + +impl From for SimilarityEntityRef { + fn from(value: BNSimilarityEntityRef) -> Self { + Self { + node_id: value.nodeId.into(), + entity_id: value.entityId.into(), + } + } +} + +impl From for BNSimilarityEntityRef { + fn from(value: SimilarityEntityRef) -> Self { + Self { + nodeId: value.node_id.into(), + entityId: value.entity_id.into(), + } + } +} diff --git a/rust/src/similarity/graph.rs b/rust/src/similarity/graph.rs new file mode 100644 index 000000000..90b74180a --- /dev/null +++ b/rust/src/similarity/graph.rs @@ -0,0 +1,218 @@ +use super::node::SimilaritySessionNode; +use super::SimilaritySessionNodeId; +use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable}; +use binaryninjacore_sys::*; +use std::ffi::c_void; + +/// Receives notifications after nodes or edges are added to or removed from a session graph. +pub trait SimilaritySessionGraphReceiver: Send + Sync + 'static { + fn on_graph_changed(&self); +} + +/// A core-backed session graph receiver. +pub struct CoreSimilaritySessionGraphReceiver { + pub(crate) handle: *mut BNSimilaritySessionGraphReceiver, +} + +impl CoreSimilaritySessionGraphReceiver { + pub unsafe fn from_raw(handle: *mut BNSimilaritySessionGraphReceiver) -> Self { + Self { handle } + } + + pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySessionGraphReceiver) -> Ref { + Ref::new(Self { handle }) + } + + pub fn create( + receiver: C, + ) -> Ref { + let receiver = Box::into_raw(Box::new(receiver)); + let mut callbacks = BNCustomSimilaritySessionGraphReceiver { + context: receiver.cast(), + externalRefTaken: None, + externalRefReleased: None, + onGraphChanged: Some(cb_on_graph_changed::), + free: Some(cb_receiver_free::), + }; + let raw_receiver = unsafe { BNCreateCustomSimilaritySessionGraphReceiver(&mut callbacks) }; + unsafe { CoreSimilaritySessionGraphReceiver::ref_from_raw(raw_receiver) } + } +} + +impl SimilaritySessionGraphReceiver for CoreSimilaritySessionGraphReceiver { + fn on_graph_changed(&self) { + unsafe { BNSimilaritySessionGraphReceiverNotifyGraphChanged(self.handle) } + } +} + +unsafe impl Send for CoreSimilaritySessionGraphReceiver {} +unsafe impl Sync for CoreSimilaritySessionGraphReceiver {} + +impl ToOwned for CoreSimilaritySessionGraphReceiver { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for CoreSimilaritySessionGraphReceiver { + unsafe fn inc_ref(handle: &Self) -> Ref { + Ref::new(Self { + handle: BNNewSimilaritySessionGraphReceiverReference(handle.handle), + }) + } + + unsafe fn dec_ref(handle: &Self) { + BNFreeSimilaritySessionGraphReceiver(handle.handle); + } +} + +impl CoreArrayProvider for CoreSimilaritySessionGraphReceiver { + type Raw = *mut BNSimilaritySessionGraphReceiver; + type Context = (); + type Wrapped<'a> = Guard<'a, Self>; +} + +unsafe impl CoreArrayProviderInner for CoreSimilaritySessionGraphReceiver { + unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) { + BNFreeSimilaritySessionGraphReceiverList(raw, count) + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { + Guard::new(Self::from_raw(*raw), context) + } +} + +unsafe extern "C" fn cb_on_graph_changed(ctxt: *mut c_void) { + ffi_wrap!("SimilaritySessionGraphReceiver::on_graph_changed", unsafe { + let ctxt: &C = &*(ctxt as *const C); + ctxt.on_graph_changed(); + }) +} + +unsafe extern "C" fn cb_receiver_free(ctxt: *mut c_void) { + ffi_wrap!("SimilaritySessionGraphReceiver::free", unsafe { + let _ = Box::from_raw(ctxt as *mut C); + }) +} + +/// A graph that controls node processing order and cannot contain cycles. +/// +/// Nodes and edges cannot be changed during a run. +pub struct SimilaritySessionGraph { + pub(crate) handle: *mut BNSimilaritySessionGraph, +} + +impl SimilaritySessionGraph { + pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySessionGraph) -> Ref { + Ref::new(Self { handle }) + } + + /// Adds a node, moving it from its current graph if necessary. + /// + /// If either graph is running, the node is unchanged. + pub fn add_node(&self, node: &SimilaritySessionNode) { + unsafe { BNSimilaritySessionGraphAddNode(self.handle, node.handle) } + } + + /// Removes a node and its edges from the graph. + pub fn remove_node(&self, node: &SimilaritySessionNode) { + unsafe { BNSimilaritySessionGraphRemoveNode(self.handle, node.handle) } + } + + /// Returns a node by ID. + pub fn node(&self, id: SimilaritySessionNodeId) -> Option> { + let handle = unsafe { BNSimilaritySessionGraphGetNode(self.handle, id.into()) }; + match handle.is_null() { + true => None, + false => Some(unsafe { SimilaritySessionNode::ref_from_raw(handle) }), + } + } + + /// Returns all nodes in the graph. + pub fn nodes(&self) -> Array { + let mut count = 0; + let result = unsafe { BNSimilaritySessionGraphGetNodes(self.handle, &mut count) }; + unsafe { Array::new(result, count, ()) } + } + + /// Returns whether an edge can be added without creating a cycle. + pub fn is_valid_edge(&self, from: &SimilaritySessionNode, to: &SimilaritySessionNode) -> bool { + unsafe { BNSimilaritySessionGraphIsValidEdge(self.handle, from.handle, to.handle) } + } + + /// Adds an edge if both nodes are present and it would not create a cycle. + pub fn add_edge(&self, from: &SimilaritySessionNode, to: &SimilaritySessionNode) -> bool { + unsafe { BNSimilaritySessionGraphAddEdge(self.handle, from.handle, to.handle) } + } + + /// Removes an edge from the graph. + pub fn remove_edge(&self, from: &SimilaritySessionNode, to: &SimilaritySessionNode) -> bool { + unsafe { BNSimilaritySessionGraphRemoveEdge(self.handle, from.handle, to.handle) } + } + + /// Adds a graph-change receiver. + pub fn add_receiver(&self, receiver: &CoreSimilaritySessionGraphReceiver) { + unsafe { BNSimilaritySessionGraphAddReceiver(self.handle, receiver.handle) } + } + + /// Removes a graph-change receiver. + pub fn remove_receiver(&self, receiver: &CoreSimilaritySessionGraphReceiver) { + unsafe { BNSimilaritySessionGraphRemoveReceiver(self.handle, receiver.handle) } + } + + /// Returns the graph-change receivers registered with this graph. + pub fn receivers(&self) -> Array { + let mut count = 0; + let result = unsafe { BNSimilaritySessionGraphGetReceivers(self.handle, &mut count) }; + unsafe { Array::new(result, count, ()) } + } + + /// Returns groups of nodes in processing order. Nodes in the same group may run in parallel. + pub fn schedule(&self) -> Vec>> { + let mut level_count = 0; + let mut node_counts_ptr: *mut usize = std::ptr::null_mut(); + let raw_schedule = unsafe { + BNSimilaritySessionGraphGetSchedule(self.handle, &mut node_counts_ptr, &mut level_count) + }; + + let mut result = Vec::with_capacity(level_count); + unsafe { + let raw_levels = std::slice::from_raw_parts(raw_schedule, level_count); + let node_counts = std::slice::from_raw_parts(node_counts_ptr, level_count); + for (&raw_level, &count) in raw_levels.iter().zip(node_counts) { + let level = std::slice::from_raw_parts(raw_level, count) + .iter() + .map(|&node| { + SimilaritySessionNode::ref_from_raw(BNNewSimilaritySessionNodeReference( + node, + )) + }) + .collect(); + result.push(level); + } + BNFreeSimilaritySessionNodeSchedule(raw_schedule, node_counts_ptr, level_count); + } + result + } +} + +impl ToOwned for SimilaritySessionGraph { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for SimilaritySessionGraph { + unsafe fn inc_ref(handle: &Self) -> Ref { + Ref::new(Self { + handle: BNNewSimilaritySessionGraphReference(handle.handle), + }) + } + unsafe fn dec_ref(handle: &Self) { + BNFreeSimilaritySessionGraph(handle.handle); + } +} diff --git a/rust/src/similarity/node.rs b/rust/src/similarity/node.rs new file mode 100644 index 000000000..cda7badb1 --- /dev/null +++ b/rust/src/similarity/node.rs @@ -0,0 +1,324 @@ +use super::{ + SimilarityApplyStatus, SimilarityEntityId, SimilarityEntityInfo, SimilarityEntityRef, + SimilarityResult, SimilarityResultId, SimilaritySessionNodeId, +}; +use crate::binary_view::BinaryView; +use crate::file_metadata::FileMetadata; +use crate::function::Function; +use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable}; +use crate::settings::Settings; +use binaryninjacore_sys::*; +use std::ffi::{CStr, CString}; + +/// The main unit of similarity processing. +pub struct SimilaritySessionNode { + pub(crate) handle: *mut BNSimilaritySessionNode, +} + +impl SimilaritySessionNode { + /// Creates a node for an open view and schedules its functions. + pub fn new(view: &BinaryView) -> Ref { + let handle = unsafe { BNCreateSimilaritySessionNode(view.handle) }; + unsafe { Ref::new(Self { handle }) } + } + + /// Creates a node that opens its view and schedules its functions when run. + /// + /// The session closes the view when the run no longer needs it, so the view may be unavailable at + /// other times. + pub fn new_from_file(file: &FileMetadata) -> Ref { + let handle = unsafe { BNCreateSimilaritySessionNodeFromFile(file.handle) }; + unsafe { Ref::new(Self { handle }) } + } + + pub unsafe fn from_raw(handle: *mut BNSimilaritySessionNode) -> Self { + Self { handle } + } + + pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySessionNode) -> Ref { + Ref::new(Self { handle }) + } + + /// Returns the node's open view if one is available. + /// + /// If the node was created with [`SimilaritySessionNode::new_from_file`], its view may be + /// unavailable outside a session run. The session closes the view when the run no longer needs + /// it. + pub fn view(&self) -> Option> { + let view = unsafe { BNSimilaritySessionNodeGetView(self.handle) }; + if view.is_null() { + None + } else { + Some(unsafe { BinaryView::ref_from_raw(view) }) + } + } + + /// Replaces the node's view and creates entities for its functions. + /// + /// A view backed by a different [`FileMetadata`] is ignored. Do not mix files. + pub fn set_view(&self, view: Option<&BinaryView>) { + unsafe { + BNSimilaritySessionNodeSetView( + self.handle, + view.map(|view| view.handle).unwrap_or(std::ptr::null_mut()), + ) + } + } + + /// Returns the file used by the node. + pub fn file(&self) -> Ref { + let file = unsafe { BNSimilaritySessionNodeGetFile(self.handle) }; + FileMetadata::ref_from_raw(file) + } + + /// Returns the settings used when this node opens its view. + /// + /// Modify the returned settings before running the session. + pub fn load_options(&self) -> Ref { + let settings = unsafe { BNSimilaritySessionNodeGetLoadOptions(self.handle) }; + unsafe { Settings::ref_from_raw(settings) } + } + + /// Returns the node's ID. + pub fn id(&self) -> SimilaritySessionNodeId { + unsafe { BNSimilaritySessionNodeGetId(self.handle) }.into() + } + + /// Adds an entity without scheduling it. + /// + /// An existing entity with the same type and address is reused. A non-empty name refreshes its + /// display name. + pub fn create_entity(&self, info: SimilarityEntityInfo) -> SimilarityEntityId { + let name = + CString::new(info.name).expect("similarity entity names cannot contain null bytes"); + let raw_info = BNSimilarityEntityInfo { + type_: info.entity_type, + address: info.address, + name: name.as_ptr(), + }; + unsafe { BNSimilaritySessionNodeCreateEntity(self.handle, &raw_info) }.into() + } + + /// Removes an entity, its schedule, its provider results, and its selected result. + /// + /// If you are looking to unschedule an entity, use [`SimilaritySessionNode::remove_scheduled_entity`]. + pub fn remove_entity(&self, id: SimilarityEntityId) -> bool { + unsafe { BNSimilaritySessionNodeRemoveEntity(self.handle, id.into()) } + } + + /// Returns information about an entity. + pub fn entity(&self, id: SimilarityEntityId) -> Option { + let mut info = BNSimilarityEntityInfo::default(); + let success = + unsafe { BNSimilaritySessionNodeGetEntity(self.handle, id.into(), &mut info) }; + if success { + let result = SimilarityEntityInfo { + entity_type: info.type_, + address: info.address, + name: if info.name.is_null() { + String::new() + } else { + unsafe { CStr::from_ptr(info.name) } + .to_string_lossy() + .into_owned() + }, + }; + unsafe { BNFreeSimilarityEntityInfo(&mut info) }; + Some(result) + } else { + None + } + } + + /// Returns all entities in the node, including entities used only as match targets. + pub fn entities(&self) -> Array { + let mut count = 0; + let entities = unsafe { BNSimilaritySessionNodeGetEntities(self.handle, &mut count) }; + unsafe { Array::new(entities, count, ()) } + } + + /// Schedules an entity for the next provider round. + /// + /// The session consumes each scheduled batch before resolution. Resolvers can schedule an entity again to + /// request another round. + /// + /// New nodes schedule all available entities, so this is mainly needed for entities added later. + pub fn add_scheduled_entity(&self, id: SimilarityEntityId) -> bool { + unsafe { BNSimilaritySessionNodeAddScheduledEntity(self.handle, id.into()) } + } + + /// Unschedules an entity without removing it from the node. + /// + /// If you are looking to remove an entity, use [`SimilaritySessionNode::remove_entity`]. + pub fn remove_scheduled_entity(&self, id: SimilarityEntityId) -> bool { + unsafe { BNSimilaritySessionNodeRemoveScheduledEntity(self.handle, id.into()) } + } + + /// Returns the entities waiting for provider processing. + pub fn scheduled_entities(&self) -> Array { + let mut count = 0; + let entities = + unsafe { BNSimilaritySessionNodeGetScheduledEntities(self.handle, &mut count) }; + unsafe { Array::new(entities, count, ()) } + } + + /// Returns the function represented by an entity, if it is available. + pub fn entity_function(&self, id: SimilarityEntityId) -> Option> { + let function = unsafe { BNSimilaritySessionNodeGetEntityFunction(self.handle, id.into()) }; + if function.is_null() { + None + } else { + Some(unsafe { Function::ref_from_raw(function) }) + } + } + + /// Returns the result IDs for an entity. + pub fn results(&self, entity: SimilarityEntityId) -> Vec { + let mut count = 0; + let results = + unsafe { BNSimilaritySessionNodeGetResults(self.handle, entity.into(), &mut count) }; + let output = unsafe { std::slice::from_raw_parts(results, count) } + .iter() + .copied() + .map(SimilarityResultId::from) + .collect(); + unsafe { BNFreeSimilarityResultIdList(results) }; + output + } + + /// Returns a stored result by its ID, which is unique within the node. + pub fn result(&self, result: SimilarityResultId) -> Option { + let mut output = BNSimilarityResult::default(); + unsafe { + BNSimilaritySessionNodeGetResult(self.handle, result.into(), &mut output) + .then(|| output.into()) + } + } + + /// Applies the standard metadata transfer from a target entity. + pub fn apply_target( + &self, + entity: SimilarityEntityId, + target: SimilarityEntityRef, + ) -> SimilarityApplyStatus { + let target = target.into(); + unsafe { BNSimilaritySessionNodeApplyTarget(self.handle, entity.into(), &target) } + } + + /// Selects a provider result for an entity. + /// + pub fn set_resolved_result( + &self, + entity: SimilarityEntityId, + result: SimilarityResultId, + ) -> bool { + unsafe { + BNSimilaritySessionNodeSetResolvedResult(self.handle, entity.into(), result.into()) + } + } + + /// Returns the selected result for an entity. + pub fn resolved_result(&self, entity: SimilarityEntityId) -> Option { + let mut result = BNSimilarityResultId::default(); + let success = unsafe { + BNSimilaritySessionNodeGetResolvedResult(self.handle, entity.into(), &mut result) + }; + success.then(|| result.into()) + } + + /// Clears the selected result for an entity. + pub fn clear_resolved_result(&self, entity: SimilarityEntityId) -> bool { + unsafe { BNSimilaritySessionNodeClearResolvedResult(self.handle, entity.into()) } + } + + /// Returns the IDs of nodes with edges into this node, in ascending order. + pub fn incoming_edges(&self) -> Vec { + let mut count = 0; + let edges = unsafe { BNSimilaritySessionNodeGetIncomingEdges(self.handle, &mut count) }; + let result = unsafe { std::slice::from_raw_parts(edges, count) } + .iter() + .copied() + .map(Into::into) + .collect(); + unsafe { BNFreeSimilaritySessionNodeEdgeList(edges) }; + result + } + + /// Returns the IDs of nodes with edges out of this node, in ascending order. + pub fn outgoing_edges(&self) -> Vec { + let mut count = 0; + let edges = unsafe { BNSimilaritySessionNodeGetOutgoingEdges(self.handle, &mut count) }; + let result = unsafe { std::slice::from_raw_parts(edges, count) } + .iter() + .copied() + .map(Into::into) + .collect(); + unsafe { BNFreeSimilaritySessionNodeEdgeList(edges) }; + result + } + + /// Returns nodes with edges into this node, ordered by ID. + pub fn incoming_nodes(&self) -> Array { + let mut count = 0; + let result = unsafe { BNSimilaritySessionNodeGetIncomingNodes(self.handle, &mut count) }; + unsafe { Array::new(result, count, ()) } + } + + /// Returns nodes with edges out of this node, ordered by ID. + pub fn outgoing_nodes(&self) -> Array { + let mut count = 0; + let result = unsafe { BNSimilaritySessionNodeGetOutgoingNodes(self.handle, &mut count) }; + unsafe { Array::new(result, count, ()) } + } +} + +unsafe impl RefCountable for SimilaritySessionNode { + unsafe fn inc_ref(handle: &Self) -> Ref { + Ref::new(Self { + handle: BNNewSimilaritySessionNodeReference(handle.handle), + }) + } + unsafe fn dec_ref(handle: &Self) { + BNFreeSimilaritySessionNode(handle.handle); + } +} + +impl ToOwned for SimilaritySessionNode { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +impl CoreArrayProvider for SimilaritySessionNode { + type Raw = *mut BNSimilaritySessionNode; + type Context = (); + type Wrapped<'a> = Guard<'a, Self>; +} + +impl CoreArrayProvider for SimilarityEntityId { + type Raw = BNSimilarityEntityId; + type Context = (); + type Wrapped<'a> = SimilarityEntityId; +} + +unsafe impl CoreArrayProviderInner for SimilarityEntityId { + unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) { + BNFreeSimilarityEntityList(raw) + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> { + (*raw).into() + } +} + +unsafe impl CoreArrayProviderInner for SimilaritySessionNode { + unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) { + BNFreeSimilaritySessionNodeList(raw, count) + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { + Guard::new(Self::from_raw(*raw), context) + } +} diff --git a/rust/src/similarity/provider.rs b/rust/src/similarity/provider.rs new file mode 100644 index 000000000..06191d637 --- /dev/null +++ b/rust/src/similarity/provider.rs @@ -0,0 +1,589 @@ +use super::node::SimilaritySessionNode; +use super::render::SimilarityRenderContext; +use super::{ + SimilarityApplyStatus, SimilarityEntityId, SimilarityEntityRef, SimilarityProviderId, + SimilarityResultId, SimilaritySessionCompletion, +}; +use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable}; +use crate::settings::Settings; +use crate::string::{BnString, IntoCStr}; +use binaryninjacore_sys::*; +use std::ffi::{c_char, c_void}; +use std::marker::PhantomData; +use std::rc::Rc; + +/// Registers a similarity provider type. +pub fn register_similarity_provider(provider_ty: C) -> (&'static C, CoreSimilarityProviderType) +where + C: SimilarityProviderType, +{ + let name = C::NAME.to_cstr(); + let description = C::DESCRIPTION.to_cstr(); + // Provider types remain registered for the lifetime of the process. + let leaked_provider: &'static C = Box::leak(Box::new(provider_ty)); + let result = unsafe { + BNRegisterSimilarityProviderType( + name.as_ptr(), + description.as_ptr(), + &mut BNCustomSimilarityProviderType { + context: leaked_provider as *const C as *mut c_void, + create: Some(cb_create_provider::), + getDefaultSettings: Some(cb_default_settings::), + }, + ) + }; + let core_provider_ty = unsafe { CoreSimilarityProviderType::from_raw(result) }; + (leaked_provider, core_provider_ty) +} + +/// Creates similarity providers with the given settings. +pub trait SimilarityProviderType: Sync + 'static { + type SimilarityProvider: SimilarityProvider; + + /// The provider name shown to users. + const NAME: &'static str; + + /// A short description of the provider. + const DESCRIPTION: &'static str; + + /// Creates a provider with the given settings. + fn create_provider(&self, settings: &Settings) -> Self::SimilarityProvider; + + /// Returns the settings used to configure new providers, if any. + fn default_settings(&self) -> Option>; +} + +/// Produces and applies similarity results for session entities. +/// +/// Visits made by a session keep the visited node and both edge endpoints active. Direct calls must provide active +/// views. +pub trait SimilarityProvider: Send + Sync + 'static { + /// Replaces this provider's settings only if they are valid. + /// + /// Return `false` without changing the current settings when the new settings are invalid or updates are not + /// supported. Use + /// [`SimilaritySession::update_provider_settings`](super::SimilaritySession::update_provider_settings) so affected + /// entities are scheduled again. + fn update_settings(&self, _settings: &Settings) -> bool { + false + } + + /// Visits the node's scheduled entities and writes results for the node. + /// + /// A successful visit replaces earlier results from this provider for scheduled entities. Results for unscheduled + /// entities remain unchanged. Return `false` to discard the visit. + fn visit_node( + &self, + _node: &SimilaritySessionNode, + _results: &mut SimilarityProviderResults<'_>, + _completion: &SimilaritySessionCompletion, + ) -> bool { + true + } + + /// Visits an edge after both endpoint nodes have been visited and writes results for the edge. + /// + /// Return `false` to discard the visit. + fn visit_node_edge( + &self, + _from: &SimilaritySessionNode, + _to: &SimilaritySessionNode, + _results: &mut SimilarityProviderResults<'_>, + _completion: &SimilaritySessionCompletion, + ) -> bool { + true + } + + /// Returns the display name for a result. + fn result_name( + &self, + node: &SimilaritySessionNode, + entity: SimilarityEntityId, + result: SimilarityResultId, + ) -> Option; + + /// Applies a result. + /// + /// The default implementation performs the standard metadata transfer from the result target. + fn apply_result( + &self, + node: &SimilaritySessionNode, + entity: SimilarityEntityId, + result: SimilarityResultId, + ) -> SimilarityApplyStatus { + let Some(result) = node.result(result) else { + return SimilarityApplyStatus::SimilarityApplyFailed; + }; + node.apply_target(entity, result.target) + } + + /// Adds views for a result to a render context. + fn render_result( + &self, + _node: &SimilaritySessionNode, + _entity: SimilarityEntityId, + _context: &SimilarityRenderContext, + _result: SimilarityResultId, + ) { + } +} + +/// Writes results for one provider node or edge visit. +/// +/// Only use this writer during the provider callback that received it. +pub struct SimilarityProviderResults<'a> { + handle: *mut BNSimilarityProviderResults, + _lifetime: PhantomData<&'a mut BNSimilarityProviderResults>, + _not_send_or_sync: PhantomData>, +} + +impl<'a> SimilarityProviderResults<'a> { + unsafe fn from_raw(handle: *mut BNSimilarityProviderResults) -> Self { + Self { + handle, + _lifetime: PhantomData, + _not_send_or_sync: PhantomData, + } + } + + /// Adds a result for a scheduled entity and returns its ID, or zero on failure. The ID is unique + /// within the node. A later visit replaces the result and gives it a new ID. + pub fn add_result( + &mut self, + source: SimilarityEntityRef, + target: SimilarityEntityRef, + similarity: u8, + confidence: u8, + ) -> SimilarityResultId { + let source = BNSimilarityEntityRef::from(source); + let target = BNSimilarityEntityRef::from(target); + unsafe { + BNSimilarityProviderResultsAddResult( + self.handle, + &source, + &target, + similarity, + confidence, + ) + .into() + } + } +} + +/// A registered similarity provider type. +pub struct CoreSimilarityProviderType { + pub(crate) handle: *mut BNSimilarityProviderType, +} + +impl CoreSimilarityProviderType { + pub unsafe fn from_raw(handle: *mut BNSimilarityProviderType) -> Self { + Self { handle } + } + + /// Returns a registered provider type by name. + pub fn by_name(name: &str) -> Option { + let name = name.to_cstr(); + let raw_type = unsafe { BNGetSimilarityProviderTypeByName(name.as_ptr()) }; + match raw_type.is_null() { + true => None, + false => Some(unsafe { Self::from_raw(raw_type) }), + } + } + + /// Returns all registered provider types. + pub fn all() -> Array { + let mut count = 0; + let result = unsafe { BNGetSimilarityProviderTypeList(&mut count) }; + unsafe { Array::new(result, count, ()) } + } + + /// Returns the registered name. + pub fn name(&self) -> String { + unsafe { BnString::into_string(BNSimilarityProviderTypeGetName(self.handle)) } + } + + /// Returns the provider description. + pub fn description(&self) -> String { + unsafe { BnString::into_string(BNSimilarityProviderTypeGetDescription(self.handle)) } + } + + /// Creates a provider with the given settings, if supported. Returns `None` outside Ultimate. + pub fn create_provider(&self, settings: &Settings) -> Option> { + let provider_raw = + unsafe { BNSimilarityProviderTypeCreateProvider(self.handle, settings.handle) }; + (!provider_raw.is_null()) + .then(|| unsafe { CoreSimilarityProvider::ref_from_raw(provider_raw) }) + } + + /// Returns the default provider settings, if available. + pub fn default_settings(&self) -> Option> { + let settings_raw = unsafe { BNSimilarityProviderTypeGetDefaultSettings(self.handle) }; + (!settings_raw.is_null()).then(|| unsafe { Settings::ref_from_raw(settings_raw) }) + } +} + +impl CoreArrayProvider for CoreSimilarityProviderType { + type Raw = *mut BNSimilarityProviderType; + type Context = (); + type Wrapped<'a> = CoreSimilarityProviderType; +} + +unsafe impl CoreArrayProviderInner for CoreSimilarityProviderType { + unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) { + BNFreeSimilarityProviderTypeList(raw) + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> { + CoreSimilarityProviderType::from_raw(*raw) + } +} + +/// A core-backed similarity provider. +pub struct CoreSimilarityProvider { + pub(crate) handle: *mut BNSimilarityProvider, +} + +impl CoreSimilarityProvider { + pub unsafe fn from_raw(handle: *mut BNSimilarityProvider) -> Self { + Self { handle } + } + + pub unsafe fn ref_from_raw(handle: *mut BNSimilarityProvider) -> Ref { + Ref::new(Self { handle }) + } + + /// Returns the provider's ID. + pub fn id(&self) -> SimilarityProviderId { + unsafe { BNSimilarityProviderGetId(self.handle) }.into() + } + + /// Wraps a custom provider in a core provider. + pub fn create( + ty: &CoreSimilarityProviderType, + provider: C, + ) -> Ref { + let provider = Box::into_raw(Box::new(provider)); + let mut callbacks = BNCustomSimilarityProvider { + context: provider.cast(), + externalRefTaken: None, + externalRefReleased: None, + updateSettings: Some(cb_update_provider_settings::), + visitNode: Some(cb_visit_node::), + visitNodeEdge: Some(cb_visit_node_edge::), + getName: Some(cb_provider_get_name::), + apply: Some(cb_provider_apply::), + render: Some(cb_provider_render::), + free: Some(cb_provider_free::), + }; + + let raw_provider = unsafe { BNCreateCustomSimilarityProvider(ty.handle, &mut callbacks) }; + unsafe { CoreSimilarityProvider::ref_from_raw(raw_provider) } + } + + /// Returns the registered type that created this provider. + pub fn provider_type(&self) -> CoreSimilarityProviderType { + let handle = unsafe { BNSimilarityProviderGetType(self.handle) }; + unsafe { CoreSimilarityProviderType::from_raw(handle) } + } + + /// Performs a complete node visit with the core managing result updates. + pub fn visit_node( + &self, + node: &SimilaritySessionNode, + completion: &SimilaritySessionCompletion, + ) { + unsafe { BNSimilarityProviderVisitNode(self.handle, node.handle, completion.handle) } + } + + /// Performs a complete edge visit with the core managing result updates. + pub fn visit_node_edge( + &self, + from: &SimilaritySessionNode, + to: &SimilaritySessionNode, + completion: &SimilaritySessionCompletion, + ) { + unsafe { + BNSimilarityProviderVisitNodeEdge( + self.handle, + from.handle, + to.handle, + completion.handle, + ) + } + } + + /// Calls this provider's node visit with an existing result writer. + pub fn perform_visit_node( + &self, + node: &SimilaritySessionNode, + results: &mut SimilarityProviderResults<'_>, + completion: &SimilaritySessionCompletion, + ) -> bool { + unsafe { + BNSimilarityProviderPerformVisitNode( + self.handle, + node.handle, + results.handle, + completion.handle, + ) + } + } + + /// Calls this provider's edge visit with an existing result writer. + pub fn perform_visit_node_edge( + &self, + from: &SimilaritySessionNode, + to: &SimilaritySessionNode, + results: &mut SimilarityProviderResults<'_>, + completion: &SimilaritySessionCompletion, + ) -> bool { + unsafe { + BNSimilarityProviderPerformVisitNodeEdge( + self.handle, + from.handle, + to.handle, + results.handle, + completion.handle, + ) + } + } + + /// Returns the display name for a result. + pub fn result_name( + &self, + node: &SimilaritySessionNode, + entity: SimilarityEntityId, + result: SimilarityResultId, + ) -> Option { + let name_raw = unsafe { + BNSimilarityProviderGetName(self.handle, node.handle, entity.into(), result.into()) + }; + if name_raw.is_null() { + return None; + } + Some(unsafe { BnString::into_string(name_raw) }) + } + + /// Applies a result. + pub fn apply_result( + &self, + node: &SimilaritySessionNode, + entity: SimilarityEntityId, + result: SimilarityResultId, + ) -> SimilarityApplyStatus { + unsafe { BNSimilarityProviderApply(self.handle, node.handle, entity.into(), result.into()) } + } + + /// Adds views for a result to a render context. + pub fn render_result( + &self, + node: &SimilaritySessionNode, + entity: SimilarityEntityId, + context: &SimilarityRenderContext, + result: SimilarityResultId, + ) { + unsafe { + BNSimilarityProviderRender( + self.handle, + node.handle, + entity.into(), + context.handle, + result.into(), + ) + } + } +} + +unsafe impl Send for CoreSimilarityProvider {} +unsafe impl Sync for CoreSimilarityProvider {} + +impl ToOwned for CoreSimilarityProvider { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for CoreSimilarityProvider { + unsafe fn inc_ref(handle: &Self) -> Ref { + Ref::new(Self { + handle: BNNewSimilarityProviderReference(handle.handle), + }) + } + + unsafe fn dec_ref(handle: &Self) { + BNFreeSimilarityProvider(handle.handle); + } +} + +impl CoreArrayProvider for CoreSimilarityProvider { + type Raw = *mut BNSimilarityProvider; + type Context = (); + type Wrapped<'a> = Guard<'a, Self>; +} + +unsafe impl CoreArrayProviderInner for CoreSimilarityProvider { + unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) { + BNFreeSimilarityProviderList(raw, count) + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { + Guard::new(Self::from_raw(*raw), context) + } +} + +/// A match produced by a similarity provider. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct SimilarityResult { + /// The provider which produced the match. + pub provider_id: SimilarityProviderId, + /// The similarity of the two entities. + pub similarity: u8, + /// The provider's confidence in the match. + pub confidence: u8, + /// The entity matched by this result. + /// + /// This can be used to transfer metadata from the target to the source entity. + pub target: SimilarityEntityRef, +} + +impl From for SimilarityResult { + fn from(value: BNSimilarityResult) -> Self { + Self { + provider_id: value.providerId.into(), + similarity: value.similarity, + confidence: value.confidence, + target: value.target.into(), + } + } +} + +impl From for BNSimilarityResult { + fn from(value: SimilarityResult) -> Self { + Self { + providerId: value.provider_id.into(), + similarity: value.similarity, + confidence: value.confidence, + target: value.target.into(), + } + } +} + +unsafe extern "C" fn cb_create_provider( + ctxt: *mut c_void, + settings: *mut BNSettings, +) -> *mut BNSimilarityProvider { + ffi_wrap!("SimilarityProviderType::create_provider", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let settings = Settings::from_raw(settings); + let provider = ctxt.create_provider(&settings); + let core_type = CoreSimilarityProviderType::by_name(C::NAME).unwrap(); + let core_provider = CoreSimilarityProvider::create(&core_type, provider); + Ref::into_raw(core_provider).handle + }) +} + +unsafe extern "C" fn cb_default_settings( + ctxt: *mut c_void, +) -> *mut BNSettings { + ffi_wrap!("SimilarityProviderType::default_settings", unsafe { + let ctxt: &C = &*(ctxt as *const C); + ctxt.default_settings() + .map(|settings| Ref::into_raw(settings).handle) + .unwrap_or(std::ptr::null_mut()) + }) +} + +unsafe extern "C" fn cb_update_provider_settings( + ctxt: *mut c_void, + settings: *mut BNSettings, +) -> bool { + ffi_wrap!("SimilarityProvider::update_settings", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let settings = Settings::from_raw(settings); + ctxt.update_settings(&settings) + }) +} + +unsafe extern "C" fn cb_visit_node( + ctxt: *mut c_void, + node: *mut BNSimilaritySessionNode, + results: *mut BNSimilarityProviderResults, + completion: *mut BNSimilaritySessionCompletion, +) -> bool { + ffi_wrap!("SimilarityProvider::visit_node", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let node = SimilaritySessionNode::from_raw(node); + let mut results = SimilarityProviderResults::from_raw(results); + let completion = SimilaritySessionCompletion::from_raw(completion); + ctxt.visit_node(&node, &mut results, &completion) + }) +} + +unsafe extern "C" fn cb_visit_node_edge( + ctxt: *mut c_void, + from: *mut BNSimilaritySessionNode, + to: *mut BNSimilaritySessionNode, + results: *mut BNSimilarityProviderResults, + completion: *mut BNSimilaritySessionCompletion, +) -> bool { + ffi_wrap!("SimilarityProvider::visit_node_edge", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let from = SimilaritySessionNode::from_raw(from); + let to = SimilaritySessionNode::from_raw(to); + let mut results = SimilarityProviderResults::from_raw(results); + let completion = SimilaritySessionCompletion::from_raw(completion); + ctxt.visit_node_edge(&from, &to, &mut results, &completion) + }) +} + +unsafe extern "C" fn cb_provider_get_name( + ctxt: *mut c_void, + node: *mut BNSimilaritySessionNode, + entity: BNSimilarityEntityId, + result: BNSimilarityResultId, +) -> *mut c_char { + ffi_wrap!("SimilarityProvider::result_name", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let node = SimilaritySessionNode::from_raw(node); + let Some(name) = ctxt.result_name(&node, entity.into(), result.into()) else { + return std::ptr::null_mut(); + }; + BnString::into_raw(BnString::new(name)) + }) +} + +unsafe extern "C" fn cb_provider_apply( + ctxt: *mut c_void, + node: *mut BNSimilaritySessionNode, + entity: BNSimilarityEntityId, + result: BNSimilarityResultId, +) -> BNSimilarityApplyStatus { + ffi_wrap!("SimilarityProvider::apply_result", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let node = SimilaritySessionNode::from_raw(node); + ctxt.apply_result(&node, entity.into(), result.into()) + }) +} + +unsafe extern "C" fn cb_provider_render( + ctxt: *mut c_void, + node: *mut BNSimilaritySessionNode, + entity: BNSimilarityEntityId, + context: *mut BNSimilarityRenderContext, + result: BNSimilarityResultId, +) { + ffi_wrap!("SimilarityProvider::render_result", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let node = SimilaritySessionNode::from_raw(node); + let context = SimilarityRenderContext::from_raw(context); + ctxt.render_result(&node, entity.into(), &context, result.into()); + }) +} + +unsafe extern "C" fn cb_provider_free(ctxt: *mut c_void) { + ffi_wrap!("SimilarityProvider::free", unsafe { + let _ = Box::from_raw(ctxt as *mut C); + }) +} diff --git a/rust/src/similarity/render.rs b/rust/src/similarity/render.rs new file mode 100644 index 000000000..8b29eedf7 --- /dev/null +++ b/rust/src/similarity/render.rs @@ -0,0 +1,329 @@ +use super::{SimilarityAnnotationType, SimilarityEntityRef, SimilarityViewType}; +use crate::binary_view::BinaryView; +use crate::flowgraph::FlowGraph; +use crate::function::{Function, FunctionViewType}; +use crate::linear_view::LinearViewObject; +use crate::rc::{Ref, RefCountable}; +use crate::string::{BnString, IntoCStr}; +use binaryninjacore_sys::*; + +/// A graph or linear view used to display a similarity result. +pub struct SimilarityView { + /// The group used to arrange related views. + pub group: String, + /// The kind of view stored in this entry. + pub view_type: SimilarityViewType, + /// The flow graph, when `view_type` is a graph. + pub graph: Option>, + /// The binary view backing a linear view. + pub data: Option>, + /// The linear view object, when `view_type` is linear. + pub linear_view: Option>, + /// The session entity for this view, if its renderer provided one. + pub entity: Option, +} + +/// An added, removed, or changed address range `[start, end)`. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct SimilarityRangeAnnotation { + pub start: u64, + pub end: u64, + pub annotation_type: SimilarityAnnotationType, +} + +/// Holds the views used to display a similarity result. +pub struct SimilarityRenderContext { + pub(crate) handle: *mut BNSimilarityRenderContext, +} + +impl SimilarityRenderContext { + /// Creates an empty render context. + pub fn new() -> Ref { + unsafe { Self::ref_from_raw(BNCreateSimilarityRenderContext()) } + } + + pub unsafe fn from_raw(handle: *mut BNSimilarityRenderContext) -> Self { + Self { handle } + } + + pub unsafe fn ref_from_raw(handle: *mut BNSimilarityRenderContext) -> Ref { + Ref::new(Self { handle }) + } + + /// Sets the function representation preferred by renderers writing to this context. + pub fn set_preferred_view_type(&self, view_type: FunctionViewType) { + let raw = FunctionViewType::into_raw(view_type); + unsafe { BNSimilarityRenderContextSetPreferredViewType(self.handle, raw) }; + FunctionViewType::free_raw(raw); + } + + /// Returns the function representation preferred by renderers writing to this context. + pub fn preferred_view_type(&self) -> FunctionViewType { + let type_ = unsafe { BNSimilarityRenderContextGetPreferredViewType(self.handle) }; + let name = if type_ == BNFunctionGraphType::HighLevelLanguageRepresentationFunctionGraph { + unsafe { BNSimilarityRenderContextGetPreferredViewTypeName(self.handle) } + } else { + std::ptr::null_mut() + }; + FunctionViewType::from_owned_raw(BNFunctionViewType { type_, name }).unwrap() + } + + /// Adds a flow graph to a view group. + pub fn add_flow_graph(&self, group: &str, graph: &FlowGraph) { + let group = group.to_cstr(); + unsafe { BNSimilarityRenderContextAddFlowGraph(self.handle, group.as_ptr(), graph.handle) } + } + + /// Adds a flow graph for a session entity to a view group. + pub fn add_flow_graph_for_entity( + &self, + group: &str, + graph: &FlowGraph, + entity: SimilarityEntityRef, + ) { + let group = group.to_cstr(); + let entity = BNSimilarityEntityRef::from(entity); + unsafe { + BNSimilarityRenderContextAddFlowGraphForEntity( + self.handle, + group.as_ptr(), + graph.handle, + &entity, + ) + } + } + + /// Adds a linear view to a view group. + pub fn add_linear_view(&self, group: &str, data: &BinaryView, linear_view: &LinearViewObject) { + let group = group.to_cstr(); + unsafe { + BNSimilarityRenderContextAddLinearView( + self.handle, + group.as_ptr(), + data.handle, + linear_view.handle, + ) + } + } + + /// Adds a linear view for a session entity to a view group. + pub fn add_linear_view_for_entity( + &self, + group: &str, + data: &BinaryView, + linear_view: &LinearViewObject, + entity: SimilarityEntityRef, + ) { + let group = group.to_cstr(); + let entity = BNSimilarityEntityRef::from(entity); + unsafe { + BNSimilarityRenderContextAddLinearViewForEntity( + self.handle, + group.as_ptr(), + data.handle, + linear_view.handle, + &entity, + ) + } + } + + /// Returns the views in insertion order. + pub fn views(&self) -> Vec { + let mut count = 0; + let raw = unsafe { BNGetSimilarityRenderContextViews(self.handle, &mut count) }; + let views = unsafe { std::slice::from_raw_parts(raw, count) }; + let result = views + .iter() + .map(|view| { + let view = *view; + let graph = unsafe { BNSimilarityViewGetFlowGraph(view) }; + let data = unsafe { BNSimilarityViewGetLinearViewData(view) }; + let linear_view = unsafe { BNSimilarityViewGetLinearView(view) }; + let mut entity = std::mem::MaybeUninit::uninit(); + let entity = unsafe { BNSimilarityViewGetEntity(view, entity.as_mut_ptr()) } + .then(|| unsafe { entity.assume_init().into() }); + SimilarityView { + group: unsafe { BnString::into_string(BNSimilarityViewGetGroup(view)) }, + view_type: unsafe { BNSimilarityViewGetType(view) }, + graph: (!graph.is_null()).then(|| unsafe { FlowGraph::ref_from_raw(graph) }), + data: (!data.is_null()).then(|| unsafe { BinaryView::ref_from_raw(data) }), + linear_view: (!linear_view.is_null()) + .then(|| unsafe { LinearViewObject::ref_from_raw(linear_view) }), + entity, + } + }) + .collect(); + unsafe { BNFreeSimilarityViewList(raw, count) }; + result + } +} + +impl ToOwned for SimilarityRenderContext { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for SimilarityRenderContext { + unsafe fn inc_ref(handle: &Self) -> Ref { + Ref::new(Self { + handle: BNNewSimilarityRenderContextReference(handle.handle), + }) + } + + unsafe fn dec_ref(handle: &Self) { + BNFreeSimilarityRenderContext(handle.handle) + } +} + +/// Renders graph and linear views with range annotations. +pub struct DiffRenderer { + handle: *mut BNDiffRenderer, +} + +impl DiffRenderer { + /// Creates a renderer without annotations. + pub fn new() -> Ref { + unsafe { + Ref::new(Self { + handle: BNCreateDiffRenderer(), + }) + } + } + + /// Adds a range annotation to later renders. + /// + /// NOTE: Empty ranges are ignored. + pub fn add_range_annotation(&self, annotation: SimilarityRangeAnnotation) { + unsafe { + BNDiffRendererAddRangeAnnotation( + self.handle, + annotation.start, + annotation.end, + annotation.annotation_type, + ) + } + } + + /// Renders the graph and linear views for a function. + pub fn render_function(&self, context: &SimilarityRenderContext, function: &Function) { + unsafe { BNDiffRendererRenderFunction(self.handle, context.handle, function.handle) } + } + + /// Renders graph and linear views for a function and session entity. + pub fn render_function_for_entity( + &self, + context: &SimilarityRenderContext, + function: &Function, + entity: SimilarityEntityRef, + ) { + let entity = BNSimilarityEntityRef::from(entity); + unsafe { + BNDiffRendererRenderFunctionForEntity( + self.handle, + context.handle, + function.handle, + &entity, + ) + } + } + + /// Renders an annotated flow graph. + pub fn render_flow_graph( + &self, + context: &SimilarityRenderContext, + group: &str, + graph: &FlowGraph, + ) { + let group = group.to_cstr(); + unsafe { + BNDiffRendererRenderFlowGraph(self.handle, context.handle, group.as_ptr(), graph.handle) + } + } + + /// Renders an annotated flow graph for a session entity. + pub fn render_flow_graph_for_entity( + &self, + context: &SimilarityRenderContext, + group: &str, + graph: &FlowGraph, + entity: SimilarityEntityRef, + ) { + let group = group.to_cstr(); + let entity = BNSimilarityEntityRef::from(entity); + unsafe { + BNDiffRendererRenderFlowGraphForEntity( + self.handle, + context.handle, + group.as_ptr(), + graph.handle, + &entity, + ) + } + } + + /// Renders an annotated linear view. + pub fn render_linear_view( + &self, + context: &SimilarityRenderContext, + group: &str, + data: &BinaryView, + linear_view: &LinearViewObject, + ) { + let group = group.to_cstr(); + unsafe { + BNDiffRendererRenderLinearView( + self.handle, + context.handle, + group.as_ptr(), + data.handle, + linear_view.handle, + ) + } + } + + /// Renders an annotated linear view for a session entity. + pub fn render_linear_view_for_entity( + &self, + context: &SimilarityRenderContext, + group: &str, + data: &BinaryView, + linear_view: &LinearViewObject, + entity: SimilarityEntityRef, + ) { + let group = group.to_cstr(); + let entity = BNSimilarityEntityRef::from(entity); + unsafe { + BNDiffRendererRenderLinearViewForEntity( + self.handle, + context.handle, + group.as_ptr(), + data.handle, + linear_view.handle, + &entity, + ) + } + } +} + +impl ToOwned for DiffRenderer { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for DiffRenderer { + unsafe fn inc_ref(handle: &Self) -> Ref { + Ref::new(Self { + handle: BNNewDiffRendererReference(handle.handle), + }) + } + + unsafe fn dec_ref(handle: &Self) { + BNFreeDiffRenderer(handle.handle) + } +} diff --git a/rust/src/similarity/session.rs b/rust/src/similarity/session.rs new file mode 100644 index 000000000..0ae67ac63 --- /dev/null +++ b/rust/src/similarity/session.rs @@ -0,0 +1,267 @@ +use super::graph::SimilaritySessionGraph; +use super::provider::CoreSimilarityProvider; +use super::{ + SimilarityProviderId, SimilaritySessionCompletionQuery, SimilaritySessionId, + SimilaritySessionResolverId, +}; +use crate::rc::{Array, Ref, RefCountable}; +use crate::settings::Settings; +use binaryninjacore_sys::*; +use std::time::Duration; + +pub mod receiver; +pub mod resolver; + +pub use receiver::*; +pub use resolver::*; + +/// Coordinates providers and resolvers across a graph of binaries. +pub struct SimilaritySession { + pub(crate) handle: *mut BNSimilaritySession, +} + +impl SimilaritySession { + /// Creates an empty session. + pub fn new() -> Ref { + let handle = unsafe { BNCreateSimilaritySession() }; + unsafe { Ref::new(Self { handle }) } + } + + pub unsafe fn from_raw(handle: *mut BNSimilaritySession) -> Self { + Self { handle } + } + + pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySession) -> Ref { + Ref::new(Self { handle }) + } + + /// Returns the session's unique ID. + pub fn id(&self) -> SimilaritySessionId { + unsafe { BNSimilaritySessionGetId(self.handle) }.into() + } + + /// Adds a provider and schedules entities processed by earlier runs for the next run. + /// + /// NOTE: Ignored while a run is active. + pub fn add_provider(&self, provider: &CoreSimilarityProvider) { + unsafe { BNSimilaritySessionAddProvider(self.handle, provider.handle) } + } + + /// Removes a provider, clears its results, and marks affected entities for resolution. + /// + /// NOTE: Ignored while a run is active. + pub fn remove_provider(&self, provider: &CoreSimilarityProvider) { + unsafe { BNSimilaritySessionRemoveProvider(self.handle, provider.handle) } + } + + /// Updates a provider already in the session and schedules previously processed entities again. + /// + /// Returns `false` during a run, when the provider is absent, or when it rejects the settings. + pub fn update_provider_settings( + &self, + provider: &CoreSimilarityProvider, + settings: &Settings, + ) -> bool { + unsafe { + BNSimilaritySessionUpdateProviderSettings(self.handle, provider.handle, settings.handle) + } + } + + /// Returns a provider by ID. + pub fn provider(&self, id: SimilarityProviderId) -> Option> { + let handle = unsafe { BNSimilaritySessionGetProvider(self.handle, id.into()) }; + match handle.is_null() { + true => None, + false => unsafe { Some(CoreSimilarityProvider::ref_from_raw(handle)) }, + } + } + + /// Returns the session's providers. + pub fn providers(&self) -> Array { + let mut count = 0; + let result = unsafe { BNSimilaritySessionGetProviders(self.handle, &mut count) }; + unsafe { Array::new(result, count, ()) } + } + + /// Adds a resolver created for this session and marks entities processed by earlier runs for resolution. + /// + /// NOTE: Returns `false` during a run, for a duplicate, or for a resolver created for another session. + pub fn add_resolver(&self, resolver: &CoreSimilaritySessionResolver) -> bool { + unsafe { BNSimilaritySessionAddResolver(self.handle, resolver.handle) } + } + + /// Removes a resolver from the session. + /// + /// NOTE: Returns `false` during a run, or if it is absent or belongs to another session. + pub fn remove_resolver(&self, resolver: &CoreSimilaritySessionResolver) -> bool { + unsafe { BNSimilaritySessionRemoveResolver(self.handle, resolver.handle) } + } + + /// Updates a resolver already in the session and marks previously processed entities for resolution. + /// + /// Returns `false` during a run, when the resolver is absent or belongs to another session, or when it rejects the + /// settings. + pub fn update_resolver_settings( + &self, + resolver: &CoreSimilaritySessionResolver, + settings: &Settings, + ) -> bool { + unsafe { + BNSimilaritySessionUpdateResolverSettings(self.handle, resolver.handle, settings.handle) + } + } + + /// Returns a resolver by ID. + pub fn resolver( + &self, + id: SimilaritySessionResolverId, + ) -> Option> { + let handle = unsafe { BNSimilaritySessionGetResolver(self.handle, id.into()) }; + match handle.is_null() { + true => None, + false => unsafe { Some(CoreSimilaritySessionResolver::ref_from_raw(handle)) }, + } + } + + /// Returns the session's resolvers. + pub fn resolvers(&self) -> Array { + let mut count = 0; + let result = unsafe { BNSimilaritySessionGetResolvers(self.handle, &mut count) }; + unsafe { Array::new(result, count, ()) } + } + + /// Adds an update receiver. + /// + /// A running session keeps using the receiver list it started with. + pub fn add_receiver(&self, receiver: &CoreSimilaritySessionReceiver) { + unsafe { BNSimilaritySessionAddReceiver(self.handle, receiver.handle) } + } + + /// Removes an update receiver. + /// + /// A running session keeps using the receiver list it started with. + pub fn remove_receiver(&self, receiver: &CoreSimilaritySessionReceiver) { + unsafe { BNSimilaritySessionRemoveReceiver(self.handle, receiver.handle) } + } + + /// Returns the session's update receivers. + pub fn receivers(&self) -> Array { + let mut count = 0; + let result = unsafe { BNSimilaritySessionGetReceivers(self.handle, &mut count) }; + unsafe { Array::new(result, count, ()) } + } + + /// Returns the session graph. + pub fn graph(&self) -> Ref { + let handle = unsafe { BNSimilaritySessionGetGraph(self.handle) }; + unsafe { SimilaritySessionGraph::ref_from_raw(handle) } + } + + /// Starts a background run with the current graph, providers, and resolvers. + /// + /// Changes to them are ignored until the run finishes. + /// + /// NOTE: Returns the completion state of the active run when already running. + pub fn run(&self) -> Ref { + unsafe { SimilaritySessionCompletion::ref_from_raw(BNSimilaritySessionRun(self.handle)) } + } +} + +impl ToOwned for SimilaritySession { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for SimilaritySession { + unsafe fn inc_ref(handle: &Self) -> Ref { + Ref::new(Self { + handle: BNNewSimilaritySessionReference(handle.handle), + }) + } + unsafe fn dec_ref(handle: &Self) { + BNFreeSimilaritySession(handle.handle); + } +} + +/// Stop requests, progress, and timing for a session run. +pub struct SimilaritySessionCompletion { + pub(crate) handle: *mut BNSimilaritySessionCompletion, +} + +impl SimilaritySessionCompletion { + /// Creates a new completion state, typically only done when calling into providers and resolvers + /// directly instead of from a session. + pub fn new() -> Ref { + unsafe { Self::ref_from_raw(BNCreateSimilaritySessionCompletion()) } + } + + pub unsafe fn from_raw(handle: *mut BNSimilaritySessionCompletion) -> Self { + Self { handle } + } + + pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySessionCompletion) -> Ref { + Ref::new(Self { handle }) + } + + /// Returns whether the session run has finished. + pub fn is_finished(&self) -> bool { + unsafe { BNSimilaritySessionCompletionIsFinished(self.handle) } + } + + /// Returns progress for the selected part of the run from `0.0` through `1.0`. + pub fn progress(&self, query: SimilaritySessionCompletionQuery) -> f64 { + let raw_query = query.into(); + unsafe { BNSimilaritySessionCompletionGetProgress(self.handle, &raw_query) } + } + + /// Asks the run to stop. + pub fn request_stop(&self) { + unsafe { BNSimilaritySessionCompletionRequestStop(self.handle) } + } + + /// Returns whether a stop has been requested. + pub fn is_stop_requested(&self) -> bool { + unsafe { BNSimilaritySessionCompletionIsStopRequested(self.handle) } + } + + /// Increases progress for a node and one provider or resolver. Progress cannot decrease. + /// + /// NOTE: Only call this from the provider or resolver selected by the query. + pub fn set_progress(&self, query: SimilaritySessionCompletionQuery, progress: f64) { + let raw_query = query.into(); + unsafe { BNSimilaritySessionCompletionSetProgress(self.handle, &raw_query, progress) } + } + + /// Returns timing for the selected part of the run. + pub fn timing(&self, query: SimilaritySessionCompletionQuery) -> Duration { + let raw_query = query.into(); + Duration::from_millis(unsafe { + BNSimilaritySessionCompletionGetTiming(self.handle, &raw_query) + }) + } +} + +unsafe impl Send for SimilaritySessionCompletion {} +unsafe impl Sync for SimilaritySessionCompletion {} + +impl ToOwned for SimilaritySessionCompletion { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for SimilaritySessionCompletion { + unsafe fn inc_ref(handle: &Self) -> Ref { + Ref::new(Self { + handle: BNNewSimilaritySessionCompletionReference(handle.handle), + }) + } + unsafe fn dec_ref(handle: &Self) { + BNFreeSimilaritySessionCompletion(handle.handle); + } +} diff --git a/rust/src/similarity/session/receiver.rs b/rust/src/similarity/session/receiver.rs new file mode 100644 index 000000000..38e4e5d57 --- /dev/null +++ b/rust/src/similarity/session/receiver.rs @@ -0,0 +1,152 @@ +use super::SimilaritySessionCompletion; +use crate::rc::{CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable}; +use crate::similarity::node::SimilaritySessionNode; +use crate::similarity::provider::CoreSimilarityProvider; +use crate::similarity::SimilarityEntityId; +use binaryninjacore_sys::*; +use std::ffi::c_void; + +/// Receives start and update notifications from a similarity session. +pub trait SimilaritySessionReceiver: Send + Sync + 'static { + /// Called when a session run starts. + /// + /// This is mainly used to get the completion state for the run. + fn on_started(&self, _completion: &SimilaritySessionCompletion) {} + + /// Called when a provider's results, resolution state, or applied metadata changes. + fn on_updated( + &self, + node: &SimilaritySessionNode, + provider: &CoreSimilarityProvider, + entities: &[SimilarityEntityId], + ); +} + +/// A core-backed similarity session receiver. +pub struct CoreSimilaritySessionReceiver { + pub(crate) handle: *mut BNSimilaritySessionReceiver, +} + +impl CoreSimilaritySessionReceiver { + pub unsafe fn from_raw(handle: *mut BNSimilaritySessionReceiver) -> Self { + Self { handle } + } + + pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySessionReceiver) -> Ref { + Ref::new(Self { handle }) + } + + /// Wraps a custom session receiver in a core receiver. + pub fn create(receiver: C) -> Ref { + let receiver = Box::into_raw(Box::new(receiver)); + let mut callbacks = BNCustomSimilaritySessionReceiver { + context: receiver.cast(), + externalRefTaken: None, + externalRefReleased: None, + onStarted: Some(cb_on_started::), + onUpdated: Some(cb_on_updated::), + free: Some(cb_receiver_free::), + }; + let raw_receiver = unsafe { BNCreateCustomSimilaritySessionReceiver(&mut callbacks) }; + unsafe { CoreSimilaritySessionReceiver::ref_from_raw(raw_receiver) } + } +} + +impl SimilaritySessionReceiver for CoreSimilaritySessionReceiver { + fn on_started(&self, completion: &SimilaritySessionCompletion) { + unsafe { BNSimilaritySessionReceiverNotifyStart(self.handle, completion.handle) } + } + + fn on_updated( + &self, + node: &SimilaritySessionNode, + provider: &CoreSimilarityProvider, + entities: &[SimilarityEntityId], + ) { + let raw_entities: Vec = + entities.iter().copied().map(Into::into).collect(); + unsafe { + BNSimilaritySessionReceiverNotifyBatch( + self.handle, + node.handle, + provider.handle, + raw_entities.as_ptr(), + raw_entities.len(), + ) + } + } +} + +unsafe impl Send for CoreSimilaritySessionReceiver {} +unsafe impl Sync for CoreSimilaritySessionReceiver {} + +impl ToOwned for CoreSimilaritySessionReceiver { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for CoreSimilaritySessionReceiver { + unsafe fn inc_ref(handle: &Self) -> Ref { + Ref::new(Self { + handle: BNNewSimilaritySessionReceiverReference(handle.handle), + }) + } + + unsafe fn dec_ref(handle: &Self) { + BNFreeSimilaritySessionReceiver(handle.handle); + } +} + +impl CoreArrayProvider for CoreSimilaritySessionReceiver { + type Raw = *mut BNSimilaritySessionReceiver; + type Context = (); + type Wrapped<'a> = Guard<'a, Self>; +} + +unsafe impl CoreArrayProviderInner for CoreSimilaritySessionReceiver { + unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) { + BNFreeSimilaritySessionReceiverList(raw, count) + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { + Guard::new(Self::from_raw(*raw), context) + } +} + +unsafe extern "C" fn cb_on_updated( + ctxt: *mut c_void, + node: *mut BNSimilaritySessionNode, + provider: *mut BNSimilarityProvider, + entities: *const BNSimilarityEntityId, + count: usize, +) { + ffi_wrap!("SimilaritySessionReceiver::on_updated", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let node = SimilaritySessionNode::from_raw(node); + let provider = CoreSimilarityProvider { handle: provider }; + let entity_slice = crate::ffi::slice_from_raw_parts(entities, count); + let mapped_entities: Vec = + entity_slice.iter().copied().map(Into::into).collect(); + ctxt.on_updated(&node, &provider, &mapped_entities); + }) +} + +unsafe extern "C" fn cb_on_started( + ctxt: *mut c_void, + completion: *mut BNSimilaritySessionCompletion, +) { + ffi_wrap!("SimilaritySessionReceiver::on_started", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let completion = SimilaritySessionCompletion::from_raw(completion); + ctxt.on_started(&completion); + }) +} + +unsafe extern "C" fn cb_receiver_free(ctxt: *mut c_void) { + ffi_wrap!("SimilaritySessionReceiver::free", unsafe { + let _ = Box::from_raw(ctxt as *mut C); + }) +} diff --git a/rust/src/similarity/session/resolver.rs b/rust/src/similarity/session/resolver.rs new file mode 100644 index 000000000..5a3e8581d --- /dev/null +++ b/rust/src/similarity/session/resolver.rs @@ -0,0 +1,393 @@ +use super::{SimilaritySession, SimilaritySessionCompletion}; +use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable}; +use crate::settings::Settings; +use crate::similarity::node::SimilaritySessionNode; +use crate::similarity::{SimilarityEntityId, SimilaritySessionResolverId}; +use crate::string::IntoCStr; +use binaryninjacore_sys::*; +use std::ffi::c_void; + +/// Creates similarity result resolvers with the given settings. +pub trait SimilaritySessionResolverType: Sync + 'static { + type SimilaritySessionResolver: SimilaritySessionResolver; + + /// The resolver name shown to users. + const NAME: &'static str; + /// A short description of the resolver. + const DESCRIPTION: &'static str; + + /// Creates a resolver for a session. + fn create_resolver( + &self, + session: &SimilaritySession, + settings: &Settings, + ) -> Self::SimilaritySessionResolver; + + /// Returns the settings used to configure new resolvers, if any. + fn default_settings(&self) -> Option>; +} + +/// Selects provider results for session entities. +/// +/// Calls for nodes in the same processing group may run at the same time. Do not keep the session after a callback +/// returns. +pub trait SimilaritySessionResolver: Send + Sync + 'static { + /// Replaces this resolver's settings only if they are valid. + /// + /// Return `false` without changing the current settings when the new settings are invalid or updates are not + /// supported. Use + /// [`SimilaritySession::update_resolver_settings`] so affected entities are resolved again. + fn update_settings(&self, _settings: &Settings) -> bool { + false + } + + /// Prepares a node after its entities are created and before providers run. + /// + /// This is useful for large graphs where not all views are available, but the resolver needs to + /// change the scheduled entities or add new entities. + fn prepare_for_node( + &self, + _session: &SimilaritySession, + _node: &SimilaritySessionNode, + _completion: &SimilaritySessionCompletion, + _resolver_id: SimilaritySessionResolverId, + ) { + } + + /// Selects results after a node's providers have run. + /// + /// To select a result, call [`SimilaritySessionNode::set_resolved_result`] for the given entity. + /// To request another provider and resolver round, call [`SimilaritySessionNode::add_scheduled_entity`]. + fn resolve_for_node( + &self, + session: &SimilaritySession, + node: &SimilaritySessionNode, + entities: &[SimilarityEntityId], + completion: &SimilaritySessionCompletion, + resolver_id: SimilaritySessionResolverId, + ); +} + +/// Registers a similarity session resolver type. +pub fn register_similarity_session_resolver( + resolver_ty: C, +) -> (&'static C, CoreSimilaritySessionResolverType) +where + C: SimilaritySessionResolverType, +{ + let name = C::NAME.to_cstr(); + let description = C::DESCRIPTION.to_cstr(); + // Resolver types remain registered for the lifetime of the process. + let leaked_resolver: &'static C = Box::leak(Box::new(resolver_ty)); + let result = unsafe { + BNRegisterSimilaritySessionResolverType( + name.as_ptr(), + description.as_ptr(), + &mut BNCustomSimilaritySessionResolverType { + context: leaked_resolver as *const C as *mut c_void, + create: Some(cb_create_resolver::), + getDefaultSettings: Some(cb_resolver_default_settings::), + }, + ) + }; + let core_resolver_ty = unsafe { CoreSimilaritySessionResolverType::from_raw(result) }; + (leaked_resolver, core_resolver_ty) +} + +/// A registered similarity session resolver type. +pub struct CoreSimilaritySessionResolverType { + pub(crate) handle: *mut BNSimilaritySessionResolverType, +} + +impl CoreSimilaritySessionResolverType { + pub unsafe fn from_raw(handle: *mut BNSimilaritySessionResolverType) -> Self { + Self { handle } + } + + /// Returns a registered resolver type by name. + pub fn by_name(name: &str) -> Option { + let name = name.to_cstr(); + let raw_type = unsafe { BNGetSimilaritySessionResolverTypeByName(name.as_ptr()) }; + match raw_type.is_null() { + true => None, + false => Some(unsafe { Self::from_raw(raw_type) }), + } + } + + /// Returns all registered resolver types. + pub fn all() -> Array { + let mut count = 0; + let result = unsafe { BNGetSimilaritySessionResolverTypeList(&mut count) }; + unsafe { Array::new(result, count, ()) } + } + + /// Returns the registered name. + pub fn name(&self) -> String { + unsafe { + crate::string::BnString::into_string(BNSimilaritySessionResolverTypeGetName( + self.handle, + )) + } + } + + /// Returns the resolver description. + pub fn description(&self) -> String { + unsafe { + crate::string::BnString::into_string(BNSimilaritySessionResolverTypeGetDescription( + self.handle, + )) + } + } + + /// Creates a resolver for a session, if supported. + pub fn create_resolver( + &self, + session: &SimilaritySession, + settings: &Settings, + ) -> Option> { + let handle = unsafe { + BNSimilaritySessionResolverTypeCreateResolver( + self.handle, + session.handle, + settings.handle, + ) + }; + (!handle.is_null()).then(|| unsafe { CoreSimilaritySessionResolver::ref_from_raw(handle) }) + } + + /// Returns the default resolver settings, if available. + pub fn default_settings(&self) -> Option> { + let handle = unsafe { BNSimilaritySessionResolverTypeGetDefaultSettings(self.handle) }; + (!handle.is_null()).then(|| unsafe { Settings::ref_from_raw(handle) }) + } +} + +impl CoreArrayProvider for CoreSimilaritySessionResolverType { + type Raw = *mut BNSimilaritySessionResolverType; + type Context = (); + type Wrapped<'a> = CoreSimilaritySessionResolverType; +} + +unsafe impl CoreArrayProviderInner for CoreSimilaritySessionResolverType { + unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) { + BNFreeSimilaritySessionResolverTypeList(raw) + } + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> { + CoreSimilaritySessionResolverType::from_raw(*raw) + } +} + +/// A core-backed similarity session resolver. +pub struct CoreSimilaritySessionResolver { + pub(crate) handle: *mut BNSimilaritySessionResolver, +} + +impl CoreSimilaritySessionResolver { + pub unsafe fn from_raw(handle: *mut BNSimilaritySessionResolver) -> Self { + Self { handle } + } + + pub unsafe fn ref_from_raw(handle: *mut BNSimilaritySessionResolver) -> Ref { + Ref::new(Self { handle }) + } + + /// Wraps a custom resolver for `session`. + pub fn create( + ty: &CoreSimilaritySessionResolverType, + session: &SimilaritySession, + resolver: C, + ) -> Ref { + let resolver = Box::into_raw(Box::new(resolver)); + let mut callbacks = BNCustomSimilaritySessionResolver { + context: resolver.cast(), + externalRefTaken: None, + externalRefReleased: None, + updateSettings: Some(cb_update_resolver_settings::), + prepareForNode: Some(cb_prepare_for_node::), + resolveForNode: Some(cb_resolve_for_node::), + free: Some(cb_resolver_free::), + }; + let raw_resolver = unsafe { + BNCreateCustomSimilaritySessionResolver(ty.handle, session.handle, &mut callbacks) + }; + unsafe { CoreSimilaritySessionResolver::ref_from_raw(raw_resolver) } + } + + /// Returns the resolver's ID. + pub fn id(&self) -> SimilaritySessionResolverId { + unsafe { BNSimilaritySessionResolverGetId(self.handle) }.into() + } + + /// Returns the registered type that created this resolver. + pub fn resolver_type(&self) -> CoreSimilaritySessionResolverType { + let handle = unsafe { BNSimilaritySessionResolverGetType(self.handle) }; + unsafe { CoreSimilaritySessionResolverType::from_raw(handle) } + } +} + +impl SimilaritySessionResolver for CoreSimilaritySessionResolver { + fn prepare_for_node( + &self, + session: &SimilaritySession, + node: &SimilaritySessionNode, + completion: &SimilaritySessionCompletion, + _resolver_id: SimilaritySessionResolverId, + ) { + unsafe { + BNSimilaritySessionResolverPrepareForNode( + self.handle, + session.handle, + node.handle, + completion.handle, + ) + } + } + + fn resolve_for_node( + &self, + session: &SimilaritySession, + node: &SimilaritySessionNode, + entities: &[SimilarityEntityId], + completion: &SimilaritySessionCompletion, + _resolver_id: SimilaritySessionResolverId, + ) { + let raw_entities: Vec = + entities.iter().copied().map(Into::into).collect(); + unsafe { + BNSimilaritySessionResolverResolveForNode( + self.handle, + session.handle, + node.handle, + raw_entities.as_ptr(), + raw_entities.len(), + completion.handle, + ) + } + } +} + +unsafe impl Send for CoreSimilaritySessionResolver {} +unsafe impl Sync for CoreSimilaritySessionResolver {} + +impl ToOwned for CoreSimilaritySessionResolver { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for CoreSimilaritySessionResolver { + unsafe fn inc_ref(handle: &Self) -> Ref { + Ref::new(Self { + handle: BNNewSimilaritySessionResolverReference(handle.handle), + }) + } + + unsafe fn dec_ref(handle: &Self) { + BNFreeSimilaritySessionResolver(handle.handle); + } +} + +impl CoreArrayProvider for CoreSimilaritySessionResolver { + type Raw = *mut BNSimilaritySessionResolver; + type Context = (); + type Wrapped<'a> = Guard<'a, Self>; +} + +unsafe impl CoreArrayProviderInner for CoreSimilaritySessionResolver { + unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) { + BNFreeSimilaritySessionResolverList(raw, count) + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { + Guard::new(Self::from_raw(*raw), context) + } +} + +unsafe extern "C" fn cb_create_resolver( + ctxt: *mut c_void, + session: *mut BNSimilaritySession, + settings: *mut BNSettings, +) -> *mut BNSimilaritySessionResolver { + ffi_wrap!("SimilaritySessionResolverType::create_resolver", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let session = SimilaritySession::from_raw(session); + let settings = Settings::from_raw(settings); + let resolver = ctxt.create_resolver(&session, &settings); + let core_type = CoreSimilaritySessionResolverType::by_name(C::NAME).unwrap(); + let core_resolver = CoreSimilaritySessionResolver::create(&core_type, &session, resolver); + Ref::into_raw(core_resolver).handle + }) +} + +unsafe extern "C" fn cb_resolver_default_settings( + ctxt: *mut c_void, +) -> *mut BNSettings { + ffi_wrap!("SimilaritySessionResolverType::default_settings", unsafe { + let ctxt: &C = &*(ctxt as *const C); + ctxt.default_settings() + .map(|settings| Ref::into_raw(settings).handle) + .unwrap_or(std::ptr::null_mut()) + }) +} + +unsafe extern "C" fn cb_update_resolver_settings( + ctxt: *mut c_void, + settings: *mut BNSettings, +) -> bool { + ffi_wrap!("SimilaritySessionResolver::update_settings", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let settings = Settings::from_raw(settings); + ctxt.update_settings(&settings) + }) +} + +unsafe extern "C" fn cb_resolve_for_node( + ctxt: *mut c_void, + session: *mut BNSimilaritySession, + node: *mut BNSimilaritySessionNode, + entities: *const BNSimilarityEntityId, + entity_count: usize, + completion: *mut BNSimilaritySessionCompletion, + resolver_id: BNSimilaritySessionResolverId, +) { + ffi_wrap!("SimilaritySessionResolver::resolve_for_node", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let session = SimilaritySession::from_raw(session); + let node = SimilaritySessionNode::from_raw(node); + let entity_slice = crate::ffi::slice_from_raw_parts(entities, entity_count); + let mapped_entities: Vec = + entity_slice.iter().copied().map(Into::into).collect(); + let completion = SimilaritySessionCompletion::from_raw(completion); + ctxt.resolve_for_node( + &session, + &node, + &mapped_entities, + &completion, + resolver_id.into(), + ); + }) +} + +unsafe extern "C" fn cb_prepare_for_node( + ctxt: *mut c_void, + session: *mut BNSimilaritySession, + node: *mut BNSimilaritySessionNode, + completion: *mut BNSimilaritySessionCompletion, + resolver_id: BNSimilaritySessionResolverId, +) { + ffi_wrap!("SimilaritySessionResolver::prepare_for_node", unsafe { + let ctxt: &C = &*(ctxt as *const C); + let session = SimilaritySession::from_raw(session); + let node = SimilaritySessionNode::from_raw(node); + let completion = SimilaritySessionCompletion::from_raw(completion); + ctxt.prepare_for_node(&session, &node, &completion, resolver_id.into()); + }) +} + +unsafe extern "C" fn cb_resolver_free(ctxt: *mut c_void) { + ffi_wrap!("SimilaritySessionResolver::free", unsafe { + let _ = Box::from_raw(ctxt as *mut C); + }) +} diff --git a/rust/tests/similarity.rs b/rust/tests/similarity.rs new file mode 100644 index 000000000..1792b1c32 --- /dev/null +++ b/rust/tests/similarity.rs @@ -0,0 +1,453 @@ +use binaryninja::function::FunctionViewType; +use binaryninja::headless::Session as HeadlessSession; +use binaryninja::rc::Ref; +use binaryninja::settings::Settings; +use binaryninja::similarity::{ + register_similarity_provider, register_similarity_session_resolver, CoreSimilarityProvider, + CoreSimilarityProviderType, CoreSimilaritySessionGraphReceiver, CoreSimilaritySessionReceiver, + SimilarityApplyStatus, SimilarityEntityId, SimilarityEntityInfo, SimilarityEntityRef, + SimilarityEntityType, SimilarityProvider, SimilarityProviderType, SimilarityRenderContext, + SimilarityResultId, SimilaritySession, SimilaritySessionCompletion, + SimilaritySessionCompletionQuery, SimilaritySessionGraphReceiver, SimilaritySessionNode, + SimilaritySessionNodeId, SimilaritySessionReceiver, SimilaritySessionResolver, + SimilaritySessionResolverId, SimilaritySessionResolverType, +}; +use std::collections::{HashMap, HashSet}; +use std::path::PathBuf; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +#[derive(Default)] +struct WorkflowState { + visited_nodes: Mutex>, + visited_edges: Mutex>, + visited_entities: Mutex>, + resolved_entities: Mutex>, + update_counts: Mutex>, + starts: AtomicUsize, + graph_changes: AtomicUsize, + provider_settings_updates: AtomicUsize, + resolver_settings_updates: AtomicUsize, +} + +struct TestProviderType { + state: Arc, +} + +impl SimilarityProviderType for TestProviderType { + type SimilarityProvider = TestProvider; + + const NAME: &'static str = "Rust Similarity Test"; + const DESCRIPTION: &'static str = "Matches each function to itself"; + + fn create_provider(&self, _settings: &Settings) -> Self::SimilarityProvider { + TestProvider { + state: self.state.clone(), + } + } + + fn default_settings(&self) -> Option> { + Some(Settings::new_with_id("rust.similarity.test.provider")) + } +} + +struct TestProvider { + state: Arc, +} + +impl SimilarityProvider for TestProvider { + fn update_settings(&self, _settings: &Settings) -> bool { + self.state + .provider_settings_updates + .fetch_add(1, Ordering::Relaxed); + true + } + + fn visit_node( + &self, + node: &SimilaritySessionNode, + results: &mut binaryninja::similarity::SimilarityProviderResults<'_>, + _completion: &SimilaritySessionCompletion, + ) -> bool { + let entities = node.scheduled_entities().to_vec(); + self.state.visited_nodes.lock().unwrap().insert(node.id()); + self.state + .visited_entities + .lock() + .unwrap() + .extend( + entities + .iter() + .copied() + .map(|entity_id| SimilarityEntityRef { + node_id: node.id(), + entity_id, + }), + ); + for entity in entities { + let entity_ref = SimilarityEntityRef { + node_id: node.id(), + entity_id: entity, + }; + let result_id = results.add_result(entity_ref, entity_ref, 255, 255); + assert_ne!(result_id, 0.into()); + } + true + } + + fn visit_node_edge( + &self, + from: &SimilaritySessionNode, + to: &SimilaritySessionNode, + _results: &mut binaryninja::similarity::SimilarityProviderResults<'_>, + _completion: &SimilaritySessionCompletion, + ) -> bool { + self.state + .visited_edges + .lock() + .unwrap() + .insert((from.id(), to.id())); + true + } + + fn result_name( + &self, + _node: &SimilaritySessionNode, + entity: SimilarityEntityId, + _result: SimilarityResultId, + ) -> Option { + Some(format!("match_{entity}")) + } + + fn render_result( + &self, + node: &SimilaritySessionNode, + entity: SimilarityEntityId, + context: &SimilarityRenderContext, + _result: SimilarityResultId, + ) { + if let Some(function) = node.entity_function(entity) { + binaryninja::similarity::DiffRenderer::new().render_function_for_entity( + context, + &function, + SimilarityEntityRef { + node_id: node.id(), + entity_id: entity, + }, + ); + } + } +} + +struct TestResolverType { + provider: Ref, + state: Arc, +} + +impl SimilaritySessionResolverType for TestResolverType { + type SimilaritySessionResolver = TestResolver; + + const NAME: &'static str = "Rust Similarity Test Resolver"; + const DESCRIPTION: &'static str = "Selects the first provider result"; + + fn create_resolver( + &self, + _session: &SimilaritySession, + _settings: &Settings, + ) -> Self::SimilaritySessionResolver { + TestResolver { + provider: self.provider.to_owned(), + state: self.state.clone(), + } + } + + fn default_settings(&self) -> Option> { + Some(Settings::new_with_id("rust.similarity.test.resolver")) + } +} + +struct TestResolver { + provider: Ref, + state: Arc, +} + +impl SimilaritySessionResolver for TestResolver { + fn update_settings(&self, _settings: &Settings) -> bool { + self.state + .resolver_settings_updates + .fetch_add(1, Ordering::Relaxed); + true + } + + fn resolve_for_node( + &self, + _session: &SimilaritySession, + node: &SimilaritySessionNode, + entities: &[SimilarityEntityId], + _completion: &SimilaritySessionCompletion, + _resolver_id: SimilaritySessionResolverId, + ) { + for entity in entities { + let Some(result) = node.results(*entity).into_iter().find(|result| { + node.result(*result) + .is_some_and(|result| result.provider_id == self.provider.id()) + }) else { + continue; + }; + if node.set_resolved_result(*entity, result) { + self.state + .resolved_entities + .lock() + .unwrap() + .insert(SimilarityEntityRef { + node_id: node.id(), + entity_id: *entity, + }); + } + } + } +} + +struct TestReceiver { + state: Arc, +} + +struct TestGraphReceiver { + state: Arc, +} + +impl SimilaritySessionGraphReceiver for TestGraphReceiver { + fn on_graph_changed(&self) { + self.state.graph_changes.fetch_add(1, Ordering::Relaxed); + } +} + +impl SimilaritySessionReceiver for TestReceiver { + fn on_started(&self, _completion: &SimilaritySessionCompletion) { + self.state.starts.fetch_add(1, Ordering::Relaxed); + } + + fn on_updated( + &self, + node: &SimilaritySessionNode, + _provider: &CoreSimilarityProvider, + entities: &[SimilarityEntityId], + ) { + let mut update_counts = self.state.update_counts.lock().unwrap(); + for &entity_id in entities { + let entity = SimilarityEntityRef { + node_id: node.id(), + entity_id, + }; + *update_counts.entry(entity).or_default() += 1; + } + } +} + +#[test] +fn similarity_session_workflow() { + let _headless_session = HeadlessSession::new().expect("Failed to initialize session"); + let out_dir = env!("OUT_DIR").parse::().unwrap(); + let fixture = out_dir.join("atox.obj"); + let root_view = binaryninja::load(&fixture).expect("Failed to create root view"); + let left_view = binaryninja::load(&fixture).expect("Failed to create left view"); + let right_view = binaryninja::load(&fixture).expect("Failed to create right view"); + let state = Arc::new(WorkflowState::default()); + + let (_, provider_type) = register_similarity_provider(TestProviderType { + state: state.clone(), + }); + assert!(CoreSimilarityProviderType::by_name(TestProviderType::NAME).is_some()); + assert_eq!(provider_type.name(), TestProviderType::NAME); + assert_eq!(provider_type.description(), TestProviderType::DESCRIPTION); + let provider_settings = provider_type + .default_settings() + .expect("test provider has default settings"); + let provider = provider_type + .create_provider(&provider_settings) + .expect("test provider can be created"); + assert_eq!(provider.provider_type().name(), TestProviderType::NAME); + + let session = SimilaritySession::new(); + session.add_provider(&provider); + assert_eq!(session.providers().len(), 1); + assert_eq!(session.provider(provider.id()).unwrap().id(), provider.id()); + + let (_, resolver_type) = register_similarity_session_resolver(TestResolverType { + provider: provider.to_owned(), + state: state.clone(), + }); + assert!( + binaryninja::similarity::CoreSimilaritySessionResolverType::by_name(TestResolverType::NAME) + .is_some() + ); + assert_eq!(resolver_type.name(), TestResolverType::NAME); + assert_eq!(resolver_type.description(), TestResolverType::DESCRIPTION); + let resolver_settings = resolver_type + .default_settings() + .expect("test resolver has default settings"); + let resolver = resolver_type + .create_resolver(&session, &resolver_settings) + .expect("test resolver can be created"); + assert_eq!(resolver.resolver_type().name(), TestResolverType::NAME); + assert!(session.add_resolver(&resolver)); + assert!(!session.add_resolver(&resolver)); + + let receiver = CoreSimilaritySessionReceiver::create(TestReceiver { + state: state.clone(), + }); + session.add_receiver(&receiver); + + let root_file = root_view.file(); + let root = SimilaritySessionNode::new_from_file(&root_file); + root.load_options().set_string("analysis.mode", "basic"); + assert_eq!(root.load_options().get_string("analysis.mode"), "basic"); + let left = SimilaritySessionNode::new(&left_view); + let right = SimilaritySessionNode::new(&right_view); + let catalog_only = root.create_entity(SimilarityEntityInfo { + entity_type: SimilarityEntityType::SimilarityEntityFunction, + address: u64::MAX, + name: "catalog-only".to_string(), + }); + assert_eq!(root.entity(catalog_only).unwrap().name, "catalog-only"); + assert!(!root + .scheduled_entities() + .iter() + .any(|entity| entity == catalog_only)); + + let graph = session.graph(); + let graph_receiver = CoreSimilaritySessionGraphReceiver::create(TestGraphReceiver { + state: state.clone(), + }); + graph.add_receiver(&graph_receiver); + for node in [&root, &left, &right] { + graph.add_node(node); + } + assert!(graph.add_edge(&root, &left)); + assert!(graph.add_edge(&root, &right)); + assert!(!graph.is_valid_edge(&left, &root)); + + let schedule = graph.schedule(); + assert_eq!(schedule.len(), 2); + assert_eq!(schedule[0][0].id(), root.id()); + assert_eq!( + schedule[1] + .iter() + .map(|node| node.id()) + .collect::>(), + HashSet::from([left.id(), right.id()]) + ); + + let left_entity = left.scheduled_entities().get(0); + let completion = session.run(); + let deadline = Instant::now() + Duration::from_secs(30); + while !completion.is_finished() { + assert!(Instant::now() < deadline, "similarity session timed out"); + std::thread::sleep(Duration::from_millis(10)); + } + + assert_eq!( + completion.progress(SimilaritySessionCompletionQuery::for_session()), + 1.0 + ); + for node in [&root, &left, &right] { + assert_eq!( + completion.progress( + SimilaritySessionCompletionQuery::for_node(node.id()).with_provider(provider.id()) + ), + 1.0 + ); + assert_eq!( + completion.progress( + SimilaritySessionCompletionQuery::for_node(node.id()).with_resolver(resolver.id()) + ), + 1.0 + ); + } + + let catalog_only_ref = SimilarityEntityRef { + node_id: root.id(), + entity_id: catalog_only, + }; + let expected_entities = [&root, &left, &right] + .into_iter() + .flat_map(|node| { + node.entities() + .into_iter() + .map(|entity_id| SimilarityEntityRef { + node_id: node.id(), + entity_id, + }) + .collect::>() + }) + .filter(|entity| *entity != catalog_only_ref) + .collect::>(); + assert_eq!(state.starts.load(Ordering::Relaxed), 1); + assert_eq!(state.graph_changes.load(Ordering::Relaxed), 5); + assert_eq!( + *state.visited_nodes.lock().unwrap(), + HashSet::from([root.id(), left.id(), right.id()]) + ); + assert_eq!( + *state.visited_edges.lock().unwrap(), + HashSet::from([(root.id(), left.id()), (root.id(), right.id())]) + ); + assert_eq!(*state.visited_entities.lock().unwrap(), expected_entities); + assert_eq!(*state.resolved_entities.lock().unwrap(), expected_entities); + let expected_update_counts = expected_entities + .iter() + .copied() + .map(|entity| (entity, 2)) + .collect::>(); + assert_eq!(*state.update_counts.lock().unwrap(), expected_update_counts); + assert!(root.resolved_result(catalog_only).is_none()); + assert!(root.view().is_none()); + for node in [&root, &left, &right] { + assert!(node.scheduled_entities().is_empty()); + } + + let result_id = left.results(left_entity)[0]; + let result = left.result(result_id).unwrap(); + assert_eq!(result.target.node_id, left.id()); + assert_eq!(left.result(result_id), Some(result)); + assert_eq!(left.result(999.into()), None); + assert_eq!( + provider.apply_result(&left, left_entity, result_id), + SimilarityApplyStatus::SimilarityApplySuccess + ); + + let render_context = SimilarityRenderContext::new(); + render_context.set_preferred_view_type(FunctionViewType::MediumLevelIL); + assert_eq!( + render_context.preferred_view_type(), + FunctionViewType::MediumLevelIL + ); + provider.render_result(&left, left_entity, &render_context, result_id); + let views = render_context.views(); + assert_eq!(views.len(), 2); + assert!(views.iter().all(|view| view.entity + == Some(SimilarityEntityRef { + node_id: left.id(), + entity_id: left_entity, + }))); + + assert!(left.clear_resolved_result(left_entity)); + assert!(left.resolved_result(left_entity).is_none()); + assert!(root.remove_entity(catalog_only)); + + assert_eq!(session.resolver(resolver.id()).unwrap().id(), resolver.id()); + assert!(session.update_resolver_settings(&resolver, &resolver_settings)); + assert_eq!(state.resolver_settings_updates.load(Ordering::Relaxed), 1); + assert!(session.update_provider_settings(&provider, &provider_settings)); + assert_eq!(state.provider_settings_updates.load(Ordering::Relaxed), 1); + assert!(session.remove_resolver(&resolver)); + assert!(!session.remove_resolver(&resolver)); + assert!(session.resolver(resolver.id()).is_none()); + session.remove_receiver(&receiver); + graph.remove_receiver(&graph_receiver); + assert!(graph.receivers().is_empty()); + session.remove_provider(&provider); + assert!(session.receivers().is_empty()); + assert!(session.providers().is_empty()); +} From 154b0c967a4d726bc18082902a6dc8e9c7be36d8 Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Mon, 3 Aug 2026 05:43:05 -0700 Subject: [PATCH 6/7] [Rust] Gate binary similarity tests to ultimate only --- rust/tests/similarity.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/rust/tests/similarity.rs b/rust/tests/similarity.rs index 1792b1c32..ce5865886 100644 --- a/rust/tests/similarity.rs +++ b/rust/tests/similarity.rs @@ -248,6 +248,10 @@ impl SimilaritySessionReceiver for TestReceiver { #[test] fn similarity_session_workflow() { let _headless_session = HeadlessSession::new().expect("Failed to initialize session"); + let is_ultimate = matches!( + binaryninja::product().as_str(), + "Binary Ninja Enterprise Client" | "Binary Ninja Ultimate" + ); let out_dir = env!("OUT_DIR").parse::().unwrap(); let fixture = out_dir.join("atox.obj"); let root_view = binaryninja::load(&fixture).expect("Failed to create root view"); @@ -264,9 +268,12 @@ fn similarity_session_workflow() { let provider_settings = provider_type .default_settings() .expect("test provider has default settings"); - let provider = provider_type - .create_provider(&provider_settings) - .expect("test provider can be created"); + let provider = provider_type.create_provider(&provider_settings); + if !is_ultimate { + assert!(provider.is_none()); + return; + } + let provider = provider.expect("test provider can be created in Ultimate"); assert_eq!(provider.provider_type().name(), TestProviderType::NAME); let session = SimilaritySession::new(); From d4023799c31b5961bd7bebd2711aa08a54d73cc0 Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Mon, 3 Aug 2026 04:16:03 -0700 Subject: [PATCH 7/7] Add WARP similarity provider --- plugins/warp/src/cache/container.rs | 4 +- plugins/warp/src/plugin.rs | 7 +- plugins/warp/src/plugin/load.rs | 9 +- plugins/warp/src/plugin/similarity.rs | 563 ++++++++++++++++++++++++++ plugins/warp/src/processor.rs | 140 +++++-- 5 files changed, 689 insertions(+), 34 deletions(-) create mode 100644 plugins/warp/src/plugin/similarity.rs diff --git a/plugins/warp/src/cache/container.rs b/plugins/warp/src/cache/container.rs index 99746d9ac..9ea0190f7 100644 --- a/plugins/warp/src/cache/container.rs +++ b/plugins/warp/src/cache/container.rs @@ -6,7 +6,7 @@ use std::sync::{Arc, OnceLock, RwLock}; pub static CONTAINER_CACHE: OnceLock>>>> = OnceLock::new(); -pub fn for_cached_containers(f: impl Fn(&dyn Container)) { +pub fn for_cached_containers(mut f: impl FnMut(&dyn Container)) { let containers_cache = CONTAINER_CACHE.get_or_init(Default::default); for container in containers_cache.iter() { if let Ok(guarded_container) = container.read() { @@ -15,7 +15,7 @@ pub fn for_cached_containers(f: impl Fn(&dyn Container)) { } } -pub fn for_cached_containers_mut(f: impl Fn(&mut dyn Container)) { +pub fn for_cached_containers_mut(mut f: impl FnMut(&mut dyn Container)) { let containers_cache = CONTAINER_CACHE.get_or_init(Default::default); for container in containers_cache.iter() { if let Ok(mut guarded_container) = container.write() { diff --git a/plugins/warp/src/plugin.rs b/plugins/warp/src/plugin.rs index bdb1f0a17..07ee118b6 100644 --- a/plugins/warp/src/plugin.rs +++ b/plugins/warp/src/plugin.rs @@ -7,17 +7,20 @@ use crate::container::network::{NetworkClient, NetworkContainer}; use crate::matcher::MatcherSettings; use crate::plugin::render_layer::HighlightRenderLayer; use crate::plugin::settings::PluginSettings; +use crate::plugin::similarity::WarpSimilarityProviderType; use crate::{core_signature_dir, user_signature_dir}; use binaryninja::background_task::BackgroundTask; use binaryninja::command::{register_command, register_command_for_function}; use binaryninja::is_ui_enabled; use binaryninja::settings::{QueryOptions, Settings}; +use binaryninja::similarity::register_similarity_provider; mod ffi; mod function; mod load; mod render_layer; -mod settings; +pub(crate) mod settings; +mod similarity; mod workflow; fn load_bundled_signatures() { @@ -180,6 +183,8 @@ fn plugin_init() -> bool { function::RemoveFunction {}, ); + register_similarity_provider(WarpSimilarityProviderType); + true } diff --git a/plugins/warp/src/plugin/load.rs b/plugins/warp/src/plugin/load.rs index 9ec6eaed1..92ca5ddde 100644 --- a/plugins/warp/src/plugin/load.rs +++ b/plugins/warp/src/plugin/load.rs @@ -12,7 +12,6 @@ use binaryninja::interaction::{ use binaryninja::rc::Ref; use std::collections::HashMap; use std::path::PathBuf; -use std::sync::atomic::AtomicBool; use std::thread; use warp::WarpFile; @@ -136,14 +135,14 @@ impl LoadSignatureFile { }; // Verify we have not already loaded the file. - let already_exists = AtomicBool::new(false); + let mut already_exists = false; for_cached_containers(|c| { - if let Ok(_) = c.source_path(&source_file_id) { + if c.source_path(&source_file_id).is_ok() { // TODO: What happens if path differs? Warn? - already_exists.store(true, std::sync::atomic::Ordering::SeqCst); + already_exists = true; } }); - if already_exists.load(std::sync::atomic::Ordering::SeqCst) { + if already_exists { let res = show_message_box( "Load again?", "File already loaded, would you like to load it again?", diff --git a/plugins/warp/src/plugin/similarity.rs b/plugins/warp/src/plugin/similarity.rs new file mode 100644 index 000000000..f64908238 --- /dev/null +++ b/plugins/warp/src/plugin/similarity.rs @@ -0,0 +1,563 @@ +use crate::cache::container::for_cached_containers; +use crate::cache::try_cached_function_guid; +use crate::container::disk::{DiskContainer, DiskContainerSource}; +use crate::container::{Container, SourceId, SourcePath}; +use crate::convert::platform_to_target; +use crate::convert::to_bn_symbol_at_address; +use crate::processor::WarpFileProcessor; +use crate::{basic_block_guid, function_guid, relocatable_regions, sorted_basic_blocks}; +use binaryninja::function::Function; +use binaryninja::rc::Ref; +use binaryninja::settings::Settings; +use binaryninja::similarity::{ + DiffRenderer, SimilarityAnnotationType, SimilarityApplyStatus, SimilarityEntityId, + SimilarityEntityInfo, SimilarityEntityRef, SimilarityEntityType, SimilarityProvider, + SimilarityProviderResults, SimilarityProviderType, SimilarityRangeAnnotation, + SimilarityRenderContext, SimilarityResultId, SimilaritySessionCompletion, + SimilaritySessionNode, SimilaritySessionNodeId, +}; +use std::collections::{HashMap, HashSet}; +use std::sync::RwLock; +use warp::signature::basic_block::BasicBlockGUID; +use warp::signature::function::{Function as WarpFunction, FunctionGUID}; + +type CatalogTargetKey = ( + SimilaritySessionNodeId, + String, + SourceId, + FunctionGUID, + String, +); +type AnalysisEntityIndex = HashMap>; + +pub struct WarpSimilarityProviderType; + +impl SimilarityProviderType for WarpSimilarityProviderType { + type SimilarityProvider = WarpSimilarityProvider; + const NAME: &'static str = "WARP"; + const DESCRIPTION: &'static str = "Uses WARP to find functions which are exact matches"; + + fn create_provider(&self, _settings: &Settings) -> Self::SimilarityProvider { + WarpSimilarityProvider::new() + } + + fn default_settings(&self) -> Option> { + None + } +} + +pub struct WarpSimilarityProvider { + container: RwLock, + mapped_targets: RwLock>, + target_data: RwLock>, +} + +enum TargetData { + Analysis(String), + Catalog(WarpFunction), +} + +impl TargetData { + fn name(&self) -> &str { + match self { + Self::Analysis(name) => name, + Self::Catalog(function) => &function.symbol.name, + } + } +} + +impl WarpSimilarityProvider { + pub fn new() -> Self { + Self { + container: RwLock::new(DiskContainer::new( + "Similarity Container".to_string(), + HashMap::new(), + )), + mapped_targets: RwLock::new(HashMap::new()), + target_data: RwLock::new(HashMap::new()), + } + } + + fn find_prepared_functions_with_matching_guid( + &self, + source: SourceId, + func: &Function, + ) -> Option> { + let func_lifted_il = func.lifted_il().ok()?; + let func_target = platform_to_target(&func.platform()); + let func_guid = function_guid(&func, &func_lifted_il); + let container = self.container.read().ok()?; + container + .functions_with_guid(&func_target, &source, &func_guid) + .ok() + } + + fn find_global_functions_with_matching_guid( + func: &Function, + ) -> Option> { + let func_lifted_il = func.lifted_il().ok()?; + let func_target = platform_to_target(&func.platform()); + let func_guid = function_guid(func, &func_lifted_il); + let mut matches = Vec::new(); + for_cached_containers(|container| { + let container_name = container.to_string(); + let Ok(sources) = container.sources_with_function_guid(&func_target, &func_guid) else { + return; + }; + for source in sources { + let Ok(functions) = + container.functions_with_guid(&func_target, &source, &func_guid) + else { + continue; + }; + matches.extend( + functions + .into_iter() + .map(|function| (container_name.clone(), source, function)), + ); + } + }); + + matches.sort_by(|left, right| { + ( + left.0.as_str(), + left.1.to_string(), + left.2.symbol.name.as_str(), + ) + .cmp(&( + right.0.as_str(), + right.1.to_string(), + right.2.symbol.name.as_str(), + )) + }); + matches.dedup_by(|left, right| { + left.0 == right.0 + && left.1 == right.1 + && left.2.guid == right.2.guid + && left.2.symbol.name == right.2.symbol.name + }); + Some(matches) + } + + fn analysis_function_guid(function: &Function) -> Option { + try_cached_function_guid(function).or_else(|| { + function + .lifted_il() + .ok() + .map(|lifted_il| function_guid(function, &lifted_il)) + }) + } + + fn available_mapped_address(guid: FunctionGUID, occupied_addresses: &HashSet) -> u64 { + let mut guid_prefix = [0u8; 8]; + guid_prefix.copy_from_slice(&guid.as_bytes()[..8]); + let mut address = u64::from_be_bytes(guid_prefix); + while occupied_addresses.contains(&address) { + address = address.wrapping_add(1); + } + address + } + + fn index_analysis_entities(node: &SimilaritySessionNode) -> AnalysisEntityIndex { + let mut entities: AnalysisEntityIndex = HashMap::new(); + for entity in &node.entities() { + let Some(function) = node.entity_function(entity) else { + continue; + }; + let Some(guid) = Self::analysis_function_guid(&function) else { + continue; + }; + entities.entry(guid).or_default().push(( + function.symbol().raw_name().to_string_lossy().into_owned(), + SimilarityEntityRef { + node_id: node.id(), + entity_id: entity, + }, + )); + } + entities + } + + fn find_analysis_entity( + entities: &AnalysisEntityIndex, + matched_function: &WarpFunction, + ) -> Option { + let matches = entities.get(&matched_function.guid)?; + matches + .iter() + .find(|(name, _)| *name == matched_function.symbol.name) + .or_else(|| matches.first()) + .map(|(_, entity)| *entity) + } + + fn create_mapped_target( + &self, + node: &SimilaritySessionNode, + container: String, + source: SourceId, + matched_function: &WarpFunction, + ) -> Option { + let key = ( + node.id(), + container, + source, + matched_function.guid, + matched_function.symbol.name.clone(), + ); + let mut mapped_targets = self.mapped_targets.write().ok()?; + if let Some(entity_id) = mapped_targets.get(&key) { + if node.entity(*entity_id).is_some() { + return Some(SimilarityEntityRef { + node_id: node.id(), + entity_id: *entity_id, + }); + } + mapped_targets.remove(&key); + } + + // TODO: At some point I may want to raise the barrier to entry for providers, this behavior + // TODO: may get in the way or progressing towards more core enabled functionality (e.g. porting types). + // Mapped entities let WARP name and apply matches which do not have a loaded + // analysis function. Their opaque address is identity storage and is never navigated. + let occupied_addresses = node + .entities() + .iter() + .filter_map(|entity| node.entity(entity).map(|info| info.address)) + .collect::>(); + let address = Self::available_mapped_address(matched_function.guid, &occupied_addresses); + + let entity_id = node.create_entity(SimilarityEntityInfo { + entity_type: SimilarityEntityType::SimilarityEntityFunction, + address, + name: matched_function.symbol.name.clone(), + }); + mapped_targets.insert(key, entity_id); + Some(SimilarityEntityRef { + node_id: node.id(), + entity_id, + }) + } + + fn related_entity_function( + node: &SimilaritySessionNode, + entity: SimilarityEntityRef, + ) -> Option> { + if entity.node_id == node.id() { + return node.entity_function(entity.entity_id); + } + for candidate in node.incoming_nodes().iter() { + if candidate.id() == entity.node_id { + return candidate.entity_function(entity.entity_id); + } + } + for candidate in node.outgoing_nodes().iter() { + if candidate.id() == entity.node_id { + return candidate.entity_function(entity.entity_id); + } + } + None + } + + fn render_blocks(function: &Function) -> Option> { + let lifted_il = function.lifted_il().ok()?; + let relocatable_regions = relocatable_regions(&function.view()); + Some( + sorted_basic_blocks(function) + .into_iter() + .map(|block| { + ( + basic_block_guid(&relocatable_regions, &block, &lifted_il), + block.start(), + block.end(), + ) + }) + .collect(), + ) + } + + fn unique_block_annotations( + blocks: &[(BasicBlockGUID, u64, u64)], + other_guids: &HashSet, + annotation_type: SimilarityAnnotationType, + ) -> Vec { + blocks + .iter() + .filter(|(guid, _, _)| !other_guids.contains(guid)) + .map(|(_, start, end)| SimilarityRangeAnnotation { + start: *start, + end: *end, + annotation_type, + }) + .collect() + } + + fn add_global_entity_results( + &self, + node: &SimilaritySessionNode, + entity: SimilarityEntityId, + results: &mut SimilarityProviderResults<'_>, + ) { + let Some(func) = node.entity_function(entity) else { + return; + }; + + // Function GUID lookup is an exact-match constraint, so every returned function has + // maximum similarity and confidence. + let matched_functions = + Self::find_global_functions_with_matching_guid(&func).unwrap_or_default(); + let mut target_data = Vec::with_capacity(matched_functions.len()); + let mut seen_targets = HashSet::with_capacity(matched_functions.len()); + let source = SimilarityEntityRef { + node_id: node.id(), + entity_id: entity, + }; + + for (container, catalog_source, matched_function) in matched_functions { + let target = + self.create_mapped_target(node, container, catalog_source, &matched_function); + let Some(target) = target else { continue }; + if !seen_targets.insert(target) { + continue; + } + let id = results.add_result(source, target, u8::MAX, u8::MAX); + if id != 0.into() { + target_data.push((target, TargetData::Catalog(matched_function))); + } + } + + if let Ok(mut stored_data) = self.target_data.write() { + stored_data.extend(target_data); + } + } + + fn add_edge_entity_results( + &self, + to: &SimilaritySessionNode, + source_id: SourceId, + source_entities: &AnalysisEntityIndex, + entity: SimilarityEntityId, + results: &mut SimilarityProviderResults<'_>, + ) { + let Some(func) = to.entity_function(entity) else { + return; + }; + let matched_functions = self + .find_prepared_functions_with_matching_guid(source_id, &func) + .unwrap_or_default(); + let mut target_data = Vec::with_capacity(matched_functions.len()); + let mut seen_targets = HashSet::with_capacity(matched_functions.len()); + let source = SimilarityEntityRef { + node_id: to.id(), + entity_id: entity, + }; + + for matched_function in matched_functions { + let Some(target) = Self::find_analysis_entity(source_entities, &matched_function) + else { + continue; + }; + if !seen_targets.insert(target) { + continue; + } + let id = results.add_result(source, target, u8::MAX, u8::MAX); + if id != 0.into() { + target_data.push((target, TargetData::Analysis(matched_function.symbol.name))); + } + } + if let Ok(mut stored_data) = self.target_data.write() { + stored_data.extend(target_data); + } + } + + fn render_functions( + context: &SimilarityRenderContext, + function: &Function, + entity: SimilarityEntityRef, + matched_function: &Function, + matched_entity: SimilarityEntityRef, + ) { + let function_blocks = Self::render_blocks(function); + let matched_blocks = Self::render_blocks(matched_function); + let function_renderer = DiffRenderer::new(); + let matched_renderer = DiffRenderer::new(); + + if let (Some(function_blocks), Some(matched_blocks)) = (function_blocks, matched_blocks) { + let function_guids = function_blocks + .iter() + .map(|(guid, _, _)| *guid) + .collect::>(); + let matched_guids = matched_blocks + .iter() + .map(|(guid, _, _)| *guid) + .collect::>(); + for annotation in Self::unique_block_annotations( + &function_blocks, + &matched_guids, + SimilarityAnnotationType::SimilarityAnnotationRemoved, + ) { + function_renderer.add_range_annotation(annotation); + } + for annotation in Self::unique_block_annotations( + &matched_blocks, + &function_guids, + SimilarityAnnotationType::SimilarityAnnotationAdded, + ) { + matched_renderer.add_range_annotation(annotation); + } + } + + function_renderer.render_function_for_entity(context, function, entity); + matched_renderer.render_function_for_entity(context, matched_function, matched_entity); + } +} + +impl SimilarityProvider for WarpSimilarityProvider { + fn visit_node( + &self, + node: &SimilaritySessionNode, + results: &mut SimilarityProviderResults<'_>, + completion: &SimilaritySessionCompletion, + ) -> bool { + let scheduled_entities = node.scheduled_entities().to_vec(); + // The session opens a node before invoking providers and retains it through its + // last dependent edge visit. + let view = node + .view() + .expect("similarity provider visited an inactive node"); + let processor = WarpFileProcessor::new(); + let source_path = SourcePath::new(view.file().file_path()); + let source_id = source_path.to_source_id(); + // Prepared node state is a complete snapshot. Rebuilding it avoids retaining stale + // function GUIDs when a sparse visit updates an existing function. + let functions = node + .entities() + .iter() + .filter_map(|entity| node.entity_function(entity)) + .collect::>(); + let result = match processor.process_view_with_functions_and_progress( + view.file().file_path(), + &view, + &functions, + |_| { + if completion.is_stop_requested() { + processor.state().cancel(); + } + }, + ) { + Ok(result) => result, + Err(error) => { + tracing::error!("Failed to prepare WARP similarity node: {:?}", error); + return false; + } + }; + if completion.is_stop_requested() { + return false; + } + + let Ok(mut container) = self.container.write() else { + return false; + }; + container + .sources + .insert(source_id, DiskContainerSource::new(source_path, result)); + drop(container); + + for entity in scheduled_entities { + if completion.is_stop_requested() { + return false; + } + self.add_global_entity_results(node, entity, results); + } + true + } + + fn visit_node_edge( + &self, + from: &SimilaritySessionNode, + to: &SimilaritySessionNode, + results: &mut SimilarityProviderResults<'_>, + completion: &SimilaritySessionCompletion, + ) -> bool { + let source = SourcePath::new(from.file().file_path()).to_source_id(); + let source_entities = Self::index_analysis_entities(from); + + for entity in &to.scheduled_entities() { + if completion.is_stop_requested() { + return false; + } + self.add_edge_entity_results(to, source, &source_entities, entity, results); + } + true + } + + fn result_name( + &self, + node: &SimilaritySessionNode, + _entity: SimilarityEntityId, + id: SimilarityResultId, + ) -> Option { + let result = node.result(id)?; + self.target_data + .read() + .ok()? + .get(&result.target) + .map(|data| data.name().to_string()) + } + + fn apply_result( + &self, + node: &SimilaritySessionNode, + entity: SimilarityEntityId, + id: SimilarityResultId, + ) -> SimilarityApplyStatus { + let Some(result) = node.result(id) else { + return SimilarityApplyStatus::SimilarityApplyFailed; + }; + let default_status = node.apply_target(entity, result.target); + if default_status == SimilarityApplyStatus::SimilarityApplySuccess { + return default_status; + } + let Some(function) = node.entity_function(entity) else { + return default_status; + }; + let matched_function = { + let target_data = match self.target_data.read() { + Ok(data) => data, + Err(_) => return default_status, + }; + match target_data.get(&result.target) { + Some(TargetData::Catalog(function)) => function.clone(), + _ => return default_status, + } + }; + + let view = function.view(); + let new_sym = to_bn_symbol_at_address(&view, &matched_function.symbol, function.start()); + view.define_auto_symbol(&new_sym); + SimilarityApplyStatus::SimilarityApplySuccess + } + + fn render_result( + &self, + node: &SimilaritySessionNode, + entity: SimilarityEntityId, + context: &SimilarityRenderContext, + result_id: SimilarityResultId, + ) { + let Some(result) = node.result(result_id) else { + return; + }; + let Some(function) = node.entity_function(entity) else { + return; + }; + let source = SimilarityEntityRef { + node_id: node.id(), + entity_id: entity, + }; + if let Some(matched_function) = Self::related_entity_function(node, result.target) { + Self::render_functions(context, &function, source, &matched_function, result.target); + } else { + DiffRenderer::new().render_function_for_entity(context, &function, source); + } + } +} diff --git a/plugins/warp/src/processor.rs b/plugins/warp/src/processor.rs index 666f8df32..5aa16f2e8 100644 --- a/plugins/warp/src/processor.rs +++ b/plugins/warp/src/processor.rs @@ -23,7 +23,7 @@ use binaryninja::binary_view::BinaryView; use binaryninja::function::Function as BNFunction; use binaryninja::project::file::ProjectFile; use binaryninja::project::Project; -use binaryninja::rc::{Guard, Ref}; +use binaryninja::rc::Ref; use crate::cache::cached_type_references; use crate::convert::platform_to_target; @@ -666,12 +666,66 @@ impl WarpFileProcessor { path: PathBuf, view: &BinaryView, ) -> Result, ProcessingError> { + let functions = view + .functions() + .iter() + .map(|function| function.to_owned()) + .collect::>(); + self.process_view_with_functions(path, view, &functions) + } + + pub fn process_view_with_functions( + &self, + path: PathBuf, + view: &BinaryView, + functions: &[Ref], + ) -> Result, ProcessingError> { + self.process_view_with_functions_and_progress(path, view, functions, |_| {}) + } + + pub fn process_view_with_functions_and_progress( + &self, + path: PathBuf, + view: &BinaryView, + functions: &[Ref], + progress: F, + ) -> Result, ProcessingError> + where + F: Fn(f64) + Sync, + { self.state .set_file_state(path.clone(), ProcessingFileState::Processing); + let result = + self.process_view_with_functions_and_progress_inner(view, functions, &progress); + let final_state = match &result { + Ok(_) => ProcessingFileState::Processed, + Err(_) => ProcessingFileState::Unprocessed, + }; + self.state.set_file_state(path, final_state); + if result.is_ok() { + progress(1.0); + } + result + } + + fn process_view_with_functions_and_progress_inner( + &self, + view: &BinaryView, + functions: &[Ref], + progress: &F, + ) -> Result, ProcessingError> + where + F: Fn(f64) + Sync, + { + progress(0.0); let mut chunks = Vec::new(); if self.file_data != IncludedDataField::Types { - let mut signature_chunks = self.create_signature_chunks(view)?; + let mut signature_chunks = self + .create_signature_chunks_for_functions_and_progress(functions, &|value| { + progress(value * 0.9) + })?; + self.check_cancelled()?; for (target, mut target_chunks) in signature_chunks.drain() { for signature_chunk in target_chunks.drain(..) { if signature_chunk.raw_functions().next().is_some() { @@ -685,8 +739,11 @@ impl WarpFileProcessor { } } } + self.check_cancelled()?; + progress(0.9); if self.file_data != IncludedDataField::Signatures { + self.check_cancelled()?; let type_chunk = self.create_type_chunk(view)?; if type_chunk.raw_types().next().is_some() { chunks.push(Chunk::new( @@ -696,9 +753,6 @@ impl WarpFileProcessor { } } - self.state - .set_file_state(path, ProcessingFileState::Processed); - Ok(WarpFile::new(WarpFileHeader::new(), chunks)) } @@ -709,20 +763,42 @@ impl WarpFileProcessor { &self, view: &BinaryView, ) -> Result>>, ProcessingError> { - let is_function_named = |f: &Guard| { + let functions = view + .functions() + .iter() + .map(|function| function.to_owned()) + .collect::>(); + self.create_signature_chunks_for_functions(&functions) + } + + pub fn create_signature_chunks_for_functions( + &self, + functions: &[Ref], + ) -> Result>>, ProcessingError> { + self.create_signature_chunks_for_functions_and_progress(functions, &|_| {}) + } + + fn create_signature_chunks_for_functions_and_progress( + &self, + functions: &[Ref], + progress: &F, + ) -> Result>>, ProcessingError> + where + F: Fn(f64) + Sync, + { + let is_function_named = |f: &&Ref| { self.included_functions == IncludedFunctionsField::All || f.defined_symbol().is_some() || f.has_user_annotations() }; - let is_function_tagged = |f: &Guard| { + let is_function_tagged = |f: &&Ref| { self.included_functions != IncludedFunctionsField::Selected || !f.function_tags(None, Some(INCLUDE_TAG_NAME)).is_empty() }; // TODO: is_function_blacklisted (use tag) // TODO: Move this background task to use the ProcessingState. - let view_functions = view.functions(); - let total_functions = view_functions.len(); + let total_functions = functions.len(); let done_functions = AtomicUsize::default(); let background_task = BackgroundTask::new( &format!("Generating signatures... ({}/{})", 0, total_functions), @@ -735,29 +811,37 @@ impl WarpFileProcessor { // a desired function is not in the created chunk. // TODO: Make this interruptable. with background_task.is_cancelled. let start = Instant::now(); - let built_functions: DashMap> = view_functions + let built_functions: DashMap> = functions .par_iter() - .inspect(|_| { - done_functions.fetch_add(1, Relaxed); + .map(|func| { + if self.state.is_cancelled() { + return None; + } + let result = if is_function_tagged(&func) + && is_function_named(&func) + && !func.analysis_skipped() + { + let target = platform_to_target(&func.platform()); + build_function( + func, + || func.lifted_il().ok(), + self.file_data == IncludedDataField::Symbols, + ) + .map(|built_function| (target, built_function)) + } else { + None + }; + let completed = done_functions.fetch_add(1, Relaxed) + 1; background_task.set_progress_text(&format!( "Generating signatures... ({}/{}) [{}s]", - done_functions.load(Relaxed), + completed, total_functions, start.elapsed().as_secs_f32() - )) - }) - .filter(is_function_tagged) - .filter(is_function_named) - .filter(|f| !f.analysis_skipped()) - .filter_map(|func| { - let target = platform_to_target(&func.platform()); - let built_function = build_function( - &func, - || func.lifted_il().ok(), - self.file_data == IncludedDataField::Symbols, - )?; - Some((target, built_function)) + )); + progress(completed as f64 / total_functions.max(1) as f64); + result }) + .filter_map(|result| result) .fold( DashMap::new, |acc: DashMap>, (target, function)| { @@ -771,6 +855,10 @@ impl WarpFileProcessor { }); acc }); + self.check_cancelled()?; + if total_functions == 0 { + progress(1.0); + } // Split into multiple chunks if a target has more than MAX_FUNCTIONS_PER_CHUNK functions. // We do this because otherwise some chunks may have too many flatbuffer tables for the verifier to handle.