Merge pull request #64 from loopholelabs/add-custom-payload-event-sup…

…port

Add support for custom payloads on device authority transfer

cmd/drafter-mounter/main.go

@@ -262,7 +262,7 @@ func main() {
 			OnRemoteDeviceExposed: func(remoteDeviceID uint32, path string) {
 				log.Println("Exposed remote device", remoteDeviceID, "at", path)
 			},
-			OnRemoteDeviceAuthorityReceived: func(remoteDeviceID uint32) {
+			OnRemoteDeviceAuthorityReceived: func(remoteDeviceID uint32, customPayload []byte) {
 				log.Println("Received authority for remote device", remoteDeviceID)
 			},
 			OnRemoteDeviceMigrationCompleted: func(remoteDeviceID uint32) {
@@ -481,6 +481,18 @@ l:
 						}
 					},
 
+					OnBeforeSendDeviceAuthority: func(deviceID uint32, remote bool) []byte {
+						var customPayload []byte
+
+						if remote {
+							log.Println("Sending authority for remote device", deviceID, "with custom payload", customPayload)
+						} else {
+							log.Println("Sending authority for local device", deviceID, "with custom payload", customPayload)
+						}
+
+						return customPayload
+					},
+
 					OnDeviceSent: func(deviceID uint32, remote bool) {
 						if remote {
 							log.Println("Sent remote device", deviceID)

cmd/drafter-peer/main.go

@@ -359,7 +359,7 @@ func main() {
 			OnRemoteDeviceExposed: func(remoteDeviceID uint32, path string) {
 				log.Println("Exposed remote device", remoteDeviceID, "at", path)
 			},
-			OnRemoteDeviceAuthorityReceived: func(remoteDeviceID uint32) {
+			OnRemoteDeviceAuthorityReceived: func(remoteDeviceID uint32, customPayload []byte) {
 				log.Println("Received authority for remote device", remoteDeviceID)
 			},
 			OnRemoteDeviceMigrationCompleted: func(remoteDeviceID uint32) {
@@ -626,6 +626,18 @@ func main() {
 				log.Println("Suspend:", time.Since(before))
 			},
 
+			OnBeforeSendDeviceAuthority: func(deviceID uint32, remote bool) []byte {
+				var customPayload []byte
+
+				if remote {
+					log.Println("Sending authority for remote device", deviceID, "with custom payload", customPayload)
+				} else {
+					log.Println("Sending authority for local device", deviceID, "with custom payload", customPayload)
+				}
+
+				return customPayload
+			},
+
 			OnDeviceSent: func(deviceID uint32, remote bool) {
 				if remote {
 					log.Println("Sent remote device", deviceID)

pkg/mounter/migrate_from.go

@@ -37,7 +37,7 @@ type MigrateFromAndMountDevice struct {
 type MigrateFromHooks struct {
 	OnRemoteDeviceReceived           func(remoteDeviceID uint32, name string)
 	OnRemoteDeviceExposed            func(remoteDeviceID uint32, path string)
-	OnRemoteDeviceAuthorityReceived  func(remoteDeviceID uint32)
+	OnRemoteDeviceAuthorityReceived  func(remoteDeviceID uint32, customPayload []byte)
 	OnRemoteDeviceMigrationCompleted func(remoteDeviceID uint32)
 
 	OnRemoteAllDevicesReceived     func()
@@ -258,12 +258,12 @@ func MigrateFromAndMount(
 								}
 
 							case byte(registry.EventCustomTransferAuthority):
-								if receivedButNotReadyRemoteDevices.Add(-1) <= 0 {
-									signalAllRemoteDevicesReady()
+								if hook := hooks.OnRemoteDeviceAuthorityReceived; hook != nil {
+									hook(index, e.CustomPayload)
 								}
 
-								if hook := hooks.OnRemoteDeviceAuthorityReceived; hook != nil {
-									hook(index)
+								if receivedButNotReadyRemoteDevices.Add(-1) <= 0 {
+									signalAllRemoteDevicesReady()
 								}
 							}
 

pkg/mounter/migrate_to.go

@@ -20,6 +20,12 @@ import (
 type MounterMigrateToHooks struct {
 	OnBeforeGetDirtyBlocks func(deviceID uint32, remote bool)
 
+	// This is called before a device authority is transferred;
+	// the return value will be sent as the custom payload for
+	// the registry.EventCustomTransferAuthority event and can be
+	// received in OnRemoteDeviceAuthorityReceived on the destination
+	OnBeforeSendDeviceAuthority func(deviceID uint32, remote bool) []byte
+
 	OnDeviceSent                       func(deviceID uint32, remote bool)
 	OnDeviceAuthoritySent              func(deviceID uint32, remote bool)
 	OnDeviceInitialMigrationProgress   func(deviceID uint32, remote bool, ready int, total int)
@@ -359,9 +365,15 @@ func (migratableMounter *MigratableMounter) MigrateTo(
 				}
 			}
 
+			var customPayload []byte
+			if hook := hooks.OnBeforeSendDeviceAuthority; hook != nil {
+				customPayload = hook(uint32(index), input.prev.prev.prev.remote)
+			}
+
 			if err := to.SendEvent(&packets.Event{
-				Type:       packets.EventCustom,
-				CustomType: byte(registry.EventCustomTransferAuthority),
+				Type:          packets.EventCustom,
+				CustomType:    byte(registry.EventCustomTransferAuthority),
+				CustomPayload: customPayload,
 			}); err != nil {
 				panic(errors.Join(ErrCouldNotSendTransferAuthorityEvent, err))
 			}

pkg/peer/migrate_from.go

@@ -283,12 +283,12 @@ func (peer *Peer[L, R, G]) MigrateFrom(
 								}
 
 							case byte(registry.EventCustomTransferAuthority):
-								if receivedButNotReadyRemoteDevices.Add(-1) <= 0 {
-									signalAllRemoteDevicesReady()
+								if hook := hooks.OnRemoteDeviceAuthorityReceived; hook != nil {
+									hook(index, e.CustomPayload)
 								}
 
-								if hook := hooks.OnRemoteDeviceAuthorityReceived; hook != nil {
-									hook(index)
+								if receivedButNotReadyRemoteDevices.Add(-1) <= 0 {
+									signalAllRemoteDevicesReady()
 								}
 							}
 

pkg/peer/migrate_to.go

@@ -27,6 +27,12 @@ type MigrateToHooks struct {
 	OnBeforeSuspend func()
 	OnAfterSuspend  func()
 
+	// This is called before a device authority is transferred;
+	// the return value will be sent as the custom payload for
+	// the registry.EventCustomTransferAuthority event and can be
+	// received in OnRemoteDeviceAuthorityReceived on the destination
+	OnBeforeSendDeviceAuthority func(deviceID uint32, remote bool) []byte
+
 	OnDeviceSent                       func(deviceID uint32, remote bool)
 	OnDeviceAuthoritySent              func(deviceID uint32, remote bool)
 	OnDeviceInitialMigrationProgress   func(deviceID uint32, remote bool, ready int, total int)
@@ -386,9 +392,15 @@ func (migratablePeer *MigratablePeer[L, R, G]) MigrateTo(
 				}
 			}
 
+			var customPayload []byte
+			if hook := hooks.OnBeforeSendDeviceAuthority; hook != nil {
+				customPayload = hook(uint32(index), input.prev.prev.prev.remote)
+			}
+
 			if err := to.SendEvent(&packets.Event{
-				Type:       packets.EventCustom,
-				CustomType: byte(registry.EventCustomTransferAuthority),
+				Type:          packets.EventCustom,
+				CustomType:    byte(registry.EventCustomTransferAuthority),
+				CustomPayload: customPayload,
 			}); err != nil {
 				panic(errors.Join(mounter.ErrCouldNotSendTransferAuthorityEvent, err))
 			}