src/darwin/Framework/CHIP/MTROTAImageTransferHandler.mm - third_party/github/project-chip/connectedhomeip - Git at Google

 /**
  *
  *    Copyright (c) 2024 Project CHIP Authors
  *
  *    Licensed under the Apache License, Version 2.0 (the "License");
  *    you may not use this file except in compliance with the License.
  *    You may obtain a copy of the License at
  *
  *        http://www.apache.org/licenses/LICENSE-2.0
  *
  *    Unless required by applicable law or agreed to in writing, software
  *    distributed under the License is distributed on an "AS IS" BASIS,
  *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  *    See the License for the specific language governing permissions and
  *    limitations under the License.
  */

 #import "MTROTAImageTransferHandler.h"
 #import "MTRDeviceControllerFactory_Internal.h"
 #import "MTRDeviceController_Internal.h"
 #import "MTRError_Internal.h"
 #import "MTROTAUnsolicitedBDXMessageHandler.h"
 #import "NSStringSpanConversion.h"

 using namespace chip;
 using namespace chip::bdx;
 using namespace chip::app;

 constexpr uint32_t kMaxBdxBlockSize = 1024;

 // Timeout for the BDX transfer session. The OTA Spec mandates this should be >= 5 minutes.
 constexpr System::Clock::Timeout kBdxTimeout = System::Clock::Seconds16(5 * 60);

 constexpr bdx::TransferRole kBdxRole = bdx::TransferRole::kSender;

 // An ARC-managed object that lets us do weak references to a MTROTAImageTransferHandler
 // (which is a C++ object).
 @interface MTROTAImageTransferHandlerWrapper : NSObject
 - (instancetype)initWithMTROTAImageTransferHandler:(MTROTAImageTransferHandler *)otaImageTransferHandler;

 // Non-atomic property since it's read or written only on the Matter queue.
 @property (nonatomic, nullable, readwrite, assign) MTROTAImageTransferHandler * otaImageTransferHandler;
 @end

 @implementation MTROTAImageTransferHandlerWrapper

 - (instancetype)initWithMTROTAImageTransferHandler:(MTROTAImageTransferHandler *)otaImageTransferHandler
 {
     if (self = [super init]) {
         _otaImageTransferHandler = otaImageTransferHandler;
     }
     return self;
 }
 @end

 MTROTAImageTransferHandler::MTROTAImageTransferHandler(chip::System::Layer * layer)
 {
     assertChipStackLockedByCurrentThread();

     mSystemLayer = layer;

     MTROTAUnsolicitedBDXMessageHandler::GetInstance()->OnTransferHandlerCreated(this);
     mOTAImageTransferHandlerWrapper = [[MTROTAImageTransferHandlerWrapper alloc] initWithMTROTAImageTransferHandler:this];
 }

 CHIP_ERROR MTROTAImageTransferHandler::Init(Messaging::ExchangeContext * exchangeCtx)
 {
     assertChipStackLockedByCurrentThread();

     mPeer = GetPeerScopedNodeId(exchangeCtx);
     auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mPeer.GetFabricIndex()];
     VerifyOrReturnError(controller != nil, CHIP_ERROR_INCORRECT_STATE);

     mDelegate = controller.otaProviderDelegate;
     mDelegateNotificationQueue = controller.otaProviderDelegateQueue;

     // We should have already checked that this controller supports OTA.
     VerifyOrReturnError(mDelegate != nil, CHIP_ERROR_INCORRECT_STATE);
     VerifyOrReturnError(mDelegateNotificationQueue != nil, CHIP_ERROR_INCORRECT_STATE);

     BitFlags<bdx::TransferControlFlags> flags(bdx::TransferControlFlags::kReceiverDrive);

     return AsyncResponder::Init(mSystemLayer, exchangeCtx, kBdxRole, flags, kMaxBdxBlockSize, kBdxTimeout);
 }

 MTROTAImageTransferHandler::~MTROTAImageTransferHandler()
 {
     assertChipStackLockedByCurrentThread();

     if (mNeedToCallTransferSessionEnd) {
         // TODO: Store the actual error involved in error cases, so we can pass the right thing here.
         InvokeTransferSessionEndCallback(CHIP_ERROR_INTERNAL);
     }

     MTROTAUnsolicitedBDXMessageHandler::GetInstance()->OnTransferHandlerDestroyed(this);
     mOTAImageTransferHandlerWrapper.otaImageTransferHandler = nullptr;
 }

 CHIP_ERROR MTROTAImageTransferHandler::OnTransferSessionBegin(const TransferSession::OutputEventType eventType)
 {
     assertChipStackLockedByCurrentThread();

     uint16_t fdl = 0;
     auto fd = mTransfer.GetFileDesignator(fdl);
     VerifyOrReturnError(fdl <= bdx::kMaxFileDesignatorLen, CHIP_ERROR_INCORRECT_STATE);
     CharSpan fileDesignatorSpan(Uint8::to_const_char(fd), fdl);

     auto fileDesignator = AsString(fileDesignatorSpan);
     if (fileDesignator == nil) {
         return CHIP_ERROR_INCORRECT_STATE;
     }

     auto offset = @(mTransfer.GetStartOffset());

     auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mPeer.GetFabricIndex()];
     VerifyOrReturnError(controller != nil, CHIP_ERROR_INCORRECT_STATE);

     MTROTAImageTransferHandlerWrapper * __weak weakWrapper = mOTAImageTransferHandlerWrapper;

     auto completionHandler = ^(NSError * _Nullable error) {
         [controller
             asyncDispatchToMatterQueue:^() {
                 assertChipStackLockedByCurrentThread();

                 // Check if the OTA image transfer handler is still valid. If not, return from the completion handler.
                 MTROTAImageTransferHandlerWrapper * strongWrapper = weakWrapper;
                 if (!strongWrapper || !strongWrapper.otaImageTransferHandler) {
                     return;
                 }

                 if (error != nil) {
                     CHIP_ERROR err = [MTRError errorToCHIPErrorCode:error];
                     LogErrorOnFailure(err);
                     NotifyEventHandled(eventType, err);
                     return;
                 }

                 // bdx::TransferSession will automatically reject a transfer if there are no
                 // common supported control modes. It will also default to the smaller
                 // block size.
                 TransferSession::TransferAcceptData acceptData;
                 acceptData.ControlMode = bdx::TransferControlFlags::kReceiverDrive;
                 acceptData.MaxBlockSize = mTransfer.GetTransferBlockSize();
                 acceptData.StartOffset = mTransfer.GetStartOffset();
                 acceptData.Length = mTransfer.GetTransferLength();

                 CHIP_ERROR err = mTransfer.AcceptTransfer(acceptData);
                 LogErrorOnFailure(err);
                 NotifyEventHandled(eventType, err);
             }
                           errorHandler:^(NSError *) {
                               // Not much we can do here, but if the controller is shut down we
                               // should already have been destroyed anyway.
                           }];
     };

     auto nodeId = @(mPeer.GetNodeId());

     auto strongDelegate = mDelegate;
     auto delegateQueue = mDelegateNotificationQueue;
     if (strongDelegate == nil || delegateQueue == nil) {
         LogErrorOnFailure(CHIP_ERROR_INCORRECT_STATE);
         return CHIP_ERROR_INCORRECT_STATE;
     }

     dispatch_async(delegateQueue, ^{
         if ([strongDelegate respondsToSelector:@selector
                             (handleBDXTransferSessionBeginForNodeID:controller:fileDesignator:offset:completion:)]) {
             [strongDelegate handleBDXTransferSessionBeginForNodeID:nodeId
                                                         controller:controller
                                                     fileDesignator:fileDesignator
                                                             offset:offset
                                                         completion:completionHandler];
         } else {
             [strongDelegate handleBDXTransferSessionBeginForNodeID:nodeId
                                                         controller:controller
                                                     fileDesignator:fileDesignator
                                                             offset:offset
                                                  completionHandler:completionHandler];
         }
     });
     mNeedToCallTransferSessionEnd = true;
     return CHIP_NO_ERROR;
 }

 void MTROTAImageTransferHandler::InvokeTransferSessionEndCallback(CHIP_ERROR error)
 {
     assertChipStackLockedByCurrentThread();

     mNeedToCallTransferSessionEnd = false;

     auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mPeer.GetFabricIndex()];
     VerifyOrReturn(controller != nil);
     auto nodeId = @(mPeer.GetNodeId());

     auto strongDelegate = mDelegate;
     auto delegateQueue = mDelegateNotificationQueue;
     if (strongDelegate == nil || delegateQueue == nil) {
         LogErrorOnFailure(CHIP_ERROR_INCORRECT_STATE);
         return;
     }

     auto nsError = [MTRError errorForCHIPErrorCode:error];
     if ([strongDelegate respondsToSelector:@selector(handleBDXTransferSessionEndForNodeID:controller:error:)]) {
         dispatch_async(delegateQueue, ^{
             [strongDelegate handleBDXTransferSessionEndForNodeID:nodeId
                                                       controller:controller
                                                            error:nsError];
         });
     }
 }

 CHIP_ERROR MTROTAImageTransferHandler::OnTransferSessionEnd(const TransferSession::OutputEventType eventType)
 {
     assertChipStackLockedByCurrentThread();

     CHIP_ERROR error = CHIP_NO_ERROR;
     if (eventType == TransferSession::OutputEventType::kTransferTimeout) {
         error = CHIP_ERROR_TIMEOUT;
     } else if (eventType != TransferSession::OutputEventType::kAckEOFReceived) {
         error = CHIP_ERROR_INTERNAL;
     }

     InvokeTransferSessionEndCallback(error);

     if (error == CHIP_NO_ERROR) {
         NotifyEventHandled(eventType, error);
     }
     return error;
 }

 CHIP_ERROR MTROTAImageTransferHandler::OnBlockQuery(const TransferSession::OutputEventType eventType, uint64_t bytesToSkip)
 {
     assertChipStackLockedByCurrentThread();

     auto blockSize = @(mTransfer.GetTransferBlockSize());
     auto blockIndex = @(mTransfer.GetNextBlockNum());

     auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mPeer.GetFabricIndex()];
     VerifyOrReturnError(controller != nil, CHIP_ERROR_INCORRECT_STATE);

     MTROTAImageTransferHandlerWrapper * __weak weakWrapper = mOTAImageTransferHandlerWrapper;

     auto completionHandler = ^(NSData * _Nullable data, BOOL isEOF) {
         [controller
             asyncDispatchToMatterQueue:^() {
                 assertChipStackLockedByCurrentThread();

                 // Check if the OTA image transfer handler is still valid. If not, return from the completion handler.
                 MTROTAImageTransferHandlerWrapper * strongWrapper = weakWrapper;
                 if (!strongWrapper || !strongWrapper.otaImageTransferHandler) {
                     return;
                 }

                 if (data == nil) {
                     NotifyEventHandled(eventType, CHIP_ERROR_INCORRECT_STATE);
                     return;
                 }

                 TransferSession::BlockData blockData;
                 blockData.Data = static_cast<const uint8_t *>([data bytes]);
                 blockData.Length = static_cast<size_t>([data length]);
                 blockData.IsEof = isEOF;

                 CHIP_ERROR err = mTransfer.PrepareBlock(blockData);
                 LogErrorOnFailure(err);
                 NotifyEventHandled(eventType, err);
             }
                           errorHandler:^(NSError *) {
                               // Not much we can do here, but if the controller is shut down we
                               // should already have been destroyed anyway.
                           }];
     };

     // TODO Handle MaxLength

     auto nodeId = @(mPeer.GetNodeId());

     auto strongDelegate = mDelegate;
     auto delegateQueue = mDelegateNotificationQueue;
     if (strongDelegate == nil || delegateQueue == nil) {
         LogErrorOnFailure(CHIP_ERROR_INCORRECT_STATE);
         return CHIP_ERROR_INCORRECT_STATE;
     }

     dispatch_async(delegateQueue, ^{
         if ([strongDelegate respondsToSelector:@selector(handleBDXQueryForNodeID:
                                                                       controller:blockSize:blockIndex:bytesToSkip:completion:)]) {
             [strongDelegate handleBDXQueryForNodeID:nodeId
                                          controller:controller
                                           blockSize:blockSize
                                          blockIndex:blockIndex
                                         bytesToSkip:@(bytesToSkip)
                                          completion:completionHandler];
         } else {
             [strongDelegate handleBDXQueryForNodeID:nodeId
                                          controller:controller
                                           blockSize:blockSize
                                          blockIndex:blockIndex
                                         bytesToSkip:@(bytesToSkip)
                                   completionHandler:completionHandler];
         }
     });
     return CHIP_NO_ERROR;
 }

 void MTROTAImageTransferHandler::HandleTransferSessionOutput(TransferSession::OutputEvent & event)
 {
     VerifyOrReturn(mDelegate != nil);

     TransferSession::OutputEventType eventType = event.EventType;

     ChipLogError(BDX, "OutputEvent type: %s", event.ToString(eventType));

     CHIP_ERROR err = CHIP_NO_ERROR;
     switch (event.EventType) {
     case TransferSession::OutputEventType::kInitReceived:
         err = OnTransferSessionBegin(eventType);
         if (err != CHIP_NO_ERROR) {
             LogErrorOnFailure(err);
             NotifyEventHandled(eventType, err);
         }
         break;
     case TransferSession::OutputEventType::kStatusReceived:
         ChipLogError(BDX, "Got StatusReport %x", static_cast<uint16_t>(event.statusData.statusCode));
         [[fallthrough]];
     case TransferSession::OutputEventType::kAckEOFReceived:
     case TransferSession::OutputEventType::kInternalError:
     case TransferSession::OutputEventType::kTransferTimeout:
         err = OnTransferSessionEnd(eventType);
         if (err != CHIP_NO_ERROR) {
             LogErrorOnFailure(err);
             NotifyEventHandled(eventType, err);
         }
         break;
     case TransferSession::OutputEventType::kQueryWithSkipReceived:
     case TransferSession::OutputEventType::kQueryReceived: {
         uint64_t bytesToSkip = (eventType == TransferSession::OutputEventType::kQueryWithSkipReceived ? event.bytesToSkip.BytesToSkip : 0);
         err = OnBlockQuery(eventType, bytesToSkip);
         if (err != CHIP_NO_ERROR) {
             LogErrorOnFailure(err);
             NotifyEventHandled(eventType, err);
         }
         break;
     }
     case TransferSession::OutputEventType::kNone:
     case TransferSession::OutputEventType::kAckReceived:
     case TransferSession::OutputEventType::kAcceptReceived:
     case TransferSession::OutputEventType::kBlockReceived:
     default:
         // Should never happen since this implements an OTA provider in BDX Receiver role.
         chipDie();
         break;
     }
 }

 void MTROTAImageTransferHandler::DestroySelf()
 {
     assertChipStackLockedByCurrentThread();

     delete this;
 }

 ScopedNodeId MTROTAImageTransferHandler::GetPeerScopedNodeId(Messaging::ExchangeContext * ec)
 {
     auto sessionHandle = ec->GetSessionHandle();

     if (sessionHandle->IsSecureSession()) {
         return sessionHandle->AsSecureSession()->GetPeer();
     } else if (sessionHandle->IsGroupSession()) {
         return sessionHandle->AsIncomingGroupSession()->GetPeer();
     }
     return ScopedNodeId();
 }

 CHIP_ERROR MTROTAImageTransferHandler::OnMessageReceived(
     Messaging::ExchangeContext * ec, const PayloadHeader & payloadHeader, System::PacketBufferHandle && payload)
 {
     assertChipStackLockedByCurrentThread();

     ChipLogProgress(BDX, "MTROTAImageTransferHandler: OnMessageReceived: message " ChipLogFormatMessageType " protocol " ChipLogFormatProtocolId,
         payloadHeader.GetMessageType(), ChipLogValueProtocolId(payloadHeader.GetProtocolID()));

     VerifyOrReturnError(ec != nullptr, CHIP_ERROR_INCORRECT_STATE);
     CHIP_ERROR err;

     // If we receive a ReceiveInit message, then we prepare for transfer.
     //
     // If init succeeds, or is not needed, we send the message to the AsyncTransferFacilitator for processing.
     if (payloadHeader.HasMessageType(MessageType::ReceiveInit)) {
         err = Init(ec);
         if (err != CHIP_NO_ERROR) {
             ChipLogError(Controller, "OnMessageReceived: Failed to prepare for transfer for BDX: %" CHIP_ERROR_FORMAT, err.Format());
             return err;
         }
     }

     // Send the message to the AsyncFacilitator to drive the BDX session state machine.
     AsyncTransferFacilitator::OnMessageReceived(ec, payloadHeader, std::move(payload));
     return err;
 }
	/**
	*
	* Copyright (c) 2024 Project CHIP Authors
	*
	* Licensed under the Apache License, Version 2.0 (the "License");
	* you may not use this file except in compliance with the License.
	* You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	#import "MTROTAImageTransferHandler.h"
	#import "MTRDeviceControllerFactory_Internal.h"
	#import "MTRDeviceController_Internal.h"
	#import "MTRError_Internal.h"
	#import "MTROTAUnsolicitedBDXMessageHandler.h"
	#import "NSStringSpanConversion.h"

	using namespace chip;
	using namespace chip::bdx;
	using namespace chip::app;

	constexpr uint32_t kMaxBdxBlockSize = 1024;

	// Timeout for the BDX transfer session. The OTA Spec mandates this should be >= 5 minutes.
	constexpr System::Clock::Timeout kBdxTimeout = System::Clock::Seconds16(5 * 60);

	constexpr bdx::TransferRole kBdxRole = bdx::TransferRole::kSender;

	// An ARC-managed object that lets us do weak references to a MTROTAImageTransferHandler
	// (which is a C++ object).
	@interface MTROTAImageTransferHandlerWrapper : NSObject
	- (instancetype)initWithMTROTAImageTransferHandler:(MTROTAImageTransferHandler *)otaImageTransferHandler;

	// Non-atomic property since it's read or written only on the Matter queue.
	@property (nonatomic, nullable, readwrite, assign) MTROTAImageTransferHandler * otaImageTransferHandler;
	@end

	@implementation MTROTAImageTransferHandlerWrapper

	- (instancetype)initWithMTROTAImageTransferHandler:(MTROTAImageTransferHandler *)otaImageTransferHandler
	{
	if (self = [super init]) {
	_otaImageTransferHandler = otaImageTransferHandler;
	}
	return self;
	}
	@end

	MTROTAImageTransferHandler::MTROTAImageTransferHandler(chip::System::Layer * layer)
	{
	assertChipStackLockedByCurrentThread();

	mSystemLayer = layer;

	MTROTAUnsolicitedBDXMessageHandler::GetInstance()->OnTransferHandlerCreated(this);
	mOTAImageTransferHandlerWrapper = [[MTROTAImageTransferHandlerWrapper alloc] initWithMTROTAImageTransferHandler:this];
	}

	CHIP_ERROR MTROTAImageTransferHandler::Init(Messaging::ExchangeContext * exchangeCtx)
	{
	assertChipStackLockedByCurrentThread();

	mPeer = GetPeerScopedNodeId(exchangeCtx);
	auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mPeer.GetFabricIndex()];
	VerifyOrReturnError(controller != nil, CHIP_ERROR_INCORRECT_STATE);

	mDelegate = controller.otaProviderDelegate;
	mDelegateNotificationQueue = controller.otaProviderDelegateQueue;

	// We should have already checked that this controller supports OTA.
	VerifyOrReturnError(mDelegate != nil, CHIP_ERROR_INCORRECT_STATE);
	VerifyOrReturnError(mDelegateNotificationQueue != nil, CHIP_ERROR_INCORRECT_STATE);

	BitFlags<bdx::TransferControlFlags> flags(bdx::TransferControlFlags::kReceiverDrive);

	return AsyncResponder::Init(mSystemLayer, exchangeCtx, kBdxRole, flags, kMaxBdxBlockSize, kBdxTimeout);
	}

	MTROTAImageTransferHandler::~MTROTAImageTransferHandler()
	{
	assertChipStackLockedByCurrentThread();

	if (mNeedToCallTransferSessionEnd) {
	// TODO: Store the actual error involved in error cases, so we can pass the right thing here.
	InvokeTransferSessionEndCallback(CHIP_ERROR_INTERNAL);
	}

	MTROTAUnsolicitedBDXMessageHandler::GetInstance()->OnTransferHandlerDestroyed(this);
	mOTAImageTransferHandlerWrapper.otaImageTransferHandler = nullptr;
	}

	CHIP_ERROR MTROTAImageTransferHandler::OnTransferSessionBegin(const TransferSession::OutputEventType eventType)
	{
	assertChipStackLockedByCurrentThread();

	uint16_t fdl = 0;
	auto fd = mTransfer.GetFileDesignator(fdl);
	VerifyOrReturnError(fdl <= bdx::kMaxFileDesignatorLen, CHIP_ERROR_INCORRECT_STATE);
	CharSpan fileDesignatorSpan(Uint8::to_const_char(fd), fdl);

	auto fileDesignator = AsString(fileDesignatorSpan);
	if (fileDesignator == nil) {
	return CHIP_ERROR_INCORRECT_STATE;
	}

	auto offset = @(mTransfer.GetStartOffset());

	auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mPeer.GetFabricIndex()];
	VerifyOrReturnError(controller != nil, CHIP_ERROR_INCORRECT_STATE);

	MTROTAImageTransferHandlerWrapper * __weak weakWrapper = mOTAImageTransferHandlerWrapper;

	auto completionHandler = ^(NSError * _Nullable error) {
	[controller
	asyncDispatchToMatterQueue:^() {
	assertChipStackLockedByCurrentThread();

	// Check if the OTA image transfer handler is still valid. If not, return from the completion handler.
	MTROTAImageTransferHandlerWrapper * strongWrapper = weakWrapper;
	if (!strongWrapper \|\| !strongWrapper.otaImageTransferHandler) {
	return;
	}

	if (error != nil) {
	CHIP_ERROR err = [MTRError errorToCHIPErrorCode:error];
	LogErrorOnFailure(err);
	NotifyEventHandled(eventType, err);
	return;
	}

	// bdx::TransferSession will automatically reject a transfer if there are no
	// common supported control modes. It will also default to the smaller
	// block size.
	TransferSession::TransferAcceptData acceptData;
	acceptData.ControlMode = bdx::TransferControlFlags::kReceiverDrive;
	acceptData.MaxBlockSize = mTransfer.GetTransferBlockSize();
	acceptData.StartOffset = mTransfer.GetStartOffset();
	acceptData.Length = mTransfer.GetTransferLength();

	CHIP_ERROR err = mTransfer.AcceptTransfer(acceptData);
	LogErrorOnFailure(err);
	NotifyEventHandled(eventType, err);
	}
	errorHandler:^(NSError *) {
	// Not much we can do here, but if the controller is shut down we
	// should already have been destroyed anyway.
	}];
	};

	auto nodeId = @(mPeer.GetNodeId());

	auto strongDelegate = mDelegate;
	auto delegateQueue = mDelegateNotificationQueue;
	if (strongDelegate == nil \|\| delegateQueue == nil) {
	LogErrorOnFailure(CHIP_ERROR_INCORRECT_STATE);
	return CHIP_ERROR_INCORRECT_STATE;
	}

	dispatch_async(delegateQueue, ^{
	if ([strongDelegate respondsToSelector:@selector
	(handleBDXTransferSessionBeginForNodeID:controller:fileDesignator:offset:completion:)]) {
	[strongDelegate handleBDXTransferSessionBeginForNodeID:nodeId
	controller:controller
	fileDesignator:fileDesignator
	offset:offset
	completion:completionHandler];
	} else {
	[strongDelegate handleBDXTransferSessionBeginForNodeID:nodeId
	controller:controller
	fileDesignator:fileDesignator
	offset:offset
	completionHandler:completionHandler];
	}
	});
	mNeedToCallTransferSessionEnd = true;
	return CHIP_NO_ERROR;
	}

	void MTROTAImageTransferHandler::InvokeTransferSessionEndCallback(CHIP_ERROR error)
	{
	assertChipStackLockedByCurrentThread();

	mNeedToCallTransferSessionEnd = false;

	auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mPeer.GetFabricIndex()];
	VerifyOrReturn(controller != nil);
	auto nodeId = @(mPeer.GetNodeId());

	auto strongDelegate = mDelegate;
	auto delegateQueue = mDelegateNotificationQueue;
	if (strongDelegate == nil \|\| delegateQueue == nil) {
	LogErrorOnFailure(CHIP_ERROR_INCORRECT_STATE);
	return;
	}

	auto nsError = [MTRError errorForCHIPErrorCode:error];
	if ([strongDelegate respondsToSelector:@selector(handleBDXTransferSessionEndForNodeID:controller:error:)]) {
	dispatch_async(delegateQueue, ^{
	[strongDelegate handleBDXTransferSessionEndForNodeID:nodeId
	controller:controller
	error:nsError];
	});
	}
	}

	CHIP_ERROR MTROTAImageTransferHandler::OnTransferSessionEnd(const TransferSession::OutputEventType eventType)
	{
	assertChipStackLockedByCurrentThread();

	CHIP_ERROR error = CHIP_NO_ERROR;
	if (eventType == TransferSession::OutputEventType::kTransferTimeout) {
	error = CHIP_ERROR_TIMEOUT;
	} else if (eventType != TransferSession::OutputEventType::kAckEOFReceived) {
	error = CHIP_ERROR_INTERNAL;
	}

	InvokeTransferSessionEndCallback(error);

	if (error == CHIP_NO_ERROR) {
	NotifyEventHandled(eventType, error);
	}
	return error;
	}

	CHIP_ERROR MTROTAImageTransferHandler::OnBlockQuery(const TransferSession::OutputEventType eventType, uint64_t bytesToSkip)
	{
	assertChipStackLockedByCurrentThread();

	auto blockSize = @(mTransfer.GetTransferBlockSize());
	auto blockIndex = @(mTransfer.GetNextBlockNum());

	auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mPeer.GetFabricIndex()];
	VerifyOrReturnError(controller != nil, CHIP_ERROR_INCORRECT_STATE);

	MTROTAImageTransferHandlerWrapper * __weak weakWrapper = mOTAImageTransferHandlerWrapper;

	auto completionHandler = ^(NSData * _Nullable data, BOOL isEOF) {
	[controller
	asyncDispatchToMatterQueue:^() {
	assertChipStackLockedByCurrentThread();

	// Check if the OTA image transfer handler is still valid. If not, return from the completion handler.
	MTROTAImageTransferHandlerWrapper * strongWrapper = weakWrapper;
	if (!strongWrapper \|\| !strongWrapper.otaImageTransferHandler) {
	return;
	}

	if (data == nil) {
	NotifyEventHandled(eventType, CHIP_ERROR_INCORRECT_STATE);
	return;
	}

	TransferSession::BlockData blockData;
	blockData.Data = static_cast<const uint8_t *>([data bytes]);
	blockData.Length = static_cast<size_t>([data length]);
	blockData.IsEof = isEOF;

	CHIP_ERROR err = mTransfer.PrepareBlock(blockData);
	LogErrorOnFailure(err);
	NotifyEventHandled(eventType, err);
	}
	errorHandler:^(NSError *) {
	// Not much we can do here, but if the controller is shut down we
	// should already have been destroyed anyway.
	}];
	};

	// TODO Handle MaxLength

	auto nodeId = @(mPeer.GetNodeId());

	auto strongDelegate = mDelegate;
	auto delegateQueue = mDelegateNotificationQueue;
	if (strongDelegate == nil \|\| delegateQueue == nil) {
	LogErrorOnFailure(CHIP_ERROR_INCORRECT_STATE);
	return CHIP_ERROR_INCORRECT_STATE;
	}

	dispatch_async(delegateQueue, ^{
	if ([strongDelegate respondsToSelector:@selector(handleBDXQueryForNodeID:
	controller:blockSize:blockIndex:bytesToSkip:completion:)]) {
	[strongDelegate handleBDXQueryForNodeID:nodeId
	controller:controller
	blockSize:blockSize
	blockIndex:blockIndex
	bytesToSkip:@(bytesToSkip)
	completion:completionHandler];
	} else {
	[strongDelegate handleBDXQueryForNodeID:nodeId
	controller:controller
	blockSize:blockSize
	blockIndex:blockIndex
	bytesToSkip:@(bytesToSkip)
	completionHandler:completionHandler];
	}
	});
	return CHIP_NO_ERROR;
	}

	void MTROTAImageTransferHandler::HandleTransferSessionOutput(TransferSession::OutputEvent & event)
	{
	VerifyOrReturn(mDelegate != nil);

	TransferSession::OutputEventType eventType = event.EventType;

	ChipLogError(BDX, "OutputEvent type: %s", event.ToString(eventType));

	CHIP_ERROR err = CHIP_NO_ERROR;
	switch (event.EventType) {
	case TransferSession::OutputEventType::kInitReceived:
	err = OnTransferSessionBegin(eventType);
	if (err != CHIP_NO_ERROR) {
	LogErrorOnFailure(err);
	NotifyEventHandled(eventType, err);
	}
	break;
	case TransferSession::OutputEventType::kStatusReceived:
	ChipLogError(BDX, "Got StatusReport %x", static_cast<uint16_t>(event.statusData.statusCode));
	[[fallthrough]];
	case TransferSession::OutputEventType::kAckEOFReceived:
	case TransferSession::OutputEventType::kInternalError:
	case TransferSession::OutputEventType::kTransferTimeout:
	err = OnTransferSessionEnd(eventType);
	if (err != CHIP_NO_ERROR) {
	LogErrorOnFailure(err);
	NotifyEventHandled(eventType, err);
	}
	break;
	case TransferSession::OutputEventType::kQueryWithSkipReceived:
	case TransferSession::OutputEventType::kQueryReceived: {
	uint64_t bytesToSkip = (eventType == TransferSession::OutputEventType::kQueryWithSkipReceived ? event.bytesToSkip.BytesToSkip : 0);
	err = OnBlockQuery(eventType, bytesToSkip);
	if (err != CHIP_NO_ERROR) {
	LogErrorOnFailure(err);
	NotifyEventHandled(eventType, err);
	}
	break;
	}
	case TransferSession::OutputEventType::kNone:
	case TransferSession::OutputEventType::kAckReceived:
	case TransferSession::OutputEventType::kAcceptReceived:
	case TransferSession::OutputEventType::kBlockReceived:
	default:
	// Should never happen since this implements an OTA provider in BDX Receiver role.
	chipDie();
	break;
	}
	}

	void MTROTAImageTransferHandler::DestroySelf()
	{
	assertChipStackLockedByCurrentThread();

	delete this;
	}

	ScopedNodeId MTROTAImageTransferHandler::GetPeerScopedNodeId(Messaging::ExchangeContext * ec)
	{
	auto sessionHandle = ec->GetSessionHandle();

	if (sessionHandle->IsSecureSession()) {
	return sessionHandle->AsSecureSession()->GetPeer();
	} else if (sessionHandle->IsGroupSession()) {
	return sessionHandle->AsIncomingGroupSession()->GetPeer();
	}
	return ScopedNodeId();
	}

	CHIP_ERROR MTROTAImageTransferHandler::OnMessageReceived(
	Messaging::ExchangeContext * ec, const PayloadHeader & payloadHeader, System::PacketBufferHandle && payload)
	{
	assertChipStackLockedByCurrentThread();

	ChipLogProgress(BDX, "MTROTAImageTransferHandler: OnMessageReceived: message " ChipLogFormatMessageType " protocol " ChipLogFormatProtocolId,
	payloadHeader.GetMessageType(), ChipLogValueProtocolId(payloadHeader.GetProtocolID()));

	VerifyOrReturnError(ec != nullptr, CHIP_ERROR_INCORRECT_STATE);
	CHIP_ERROR err;

	// If we receive a ReceiveInit message, then we prepare for transfer.
	//
	// If init succeeds, or is not needed, we send the message to the AsyncTransferFacilitator for processing.
	if (payloadHeader.HasMessageType(MessageType::ReceiveInit)) {
	err = Init(ec);
	if (err != CHIP_NO_ERROR) {
	ChipLogError(Controller, "OnMessageReceived: Failed to prepare for transfer for BDX: %" CHIP_ERROR_FORMAT, err.Format());
	return err;
	}
	}

	// Send the message to the AsyncFacilitator to drive the BDX session state machine.
	AsyncTransferFacilitator::OnMessageReceived(ec, payloadHeader, std::move(payload));
	return err;
	}