From ec09ff0b7cf7f1a8118528a58a20104b96996398 Mon Sep 17 00:00:00 2001
From: Jan Nanista <jan.jakub.nanista@gmail.com>
Date: Mon, 11 Dec 2023 11:16:32 -0800
Subject: [PATCH] chore: Label return types for IEndpoint

---
 packages/protocol-utils-evm/src/endpoint/sdk.ts | 15 +++++++++------
 packages/protocol-utils/src/endpoint/types.ts   | 16 +++++++++++++---
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/packages/protocol-utils-evm/src/endpoint/sdk.ts b/packages/protocol-utils-evm/src/endpoint/sdk.ts
index bb590075e..3f64cc71f 100644
--- a/packages/protocol-utils-evm/src/endpoint/sdk.ts
+++ b/packages/protocol-utils-evm/src/endpoint/sdk.ts
@@ -4,21 +4,24 @@ import type { EndpointId } from '@layerzerolabs/lz-definitions'
 import { ignoreZero, makeZeroAddress, OmniSDK } from '@layerzerolabs/utils-evm'
 
 export class Endpoint extends OmniSDK implements IEndpoint {
-    async getDefaultReceiveLibrary(eid: EndpointId): Promise<string | undefined> {
+    async getDefaultReceiveLibrary(eid: EndpointId): Promise<Address | undefined> {
         return ignoreZero(await this.contract.contract.defaultReceiveLibrary(eid))
     }
 
-    async getSendLibrary(sender: Address, dstEid: EndpointId): Promise<string | undefined> {
+    async getSendLibrary(sender: Address, dstEid: EndpointId): Promise<Address | undefined> {
         return ignoreZero(await this.contract.contract.getSendLibrary(sender, dstEid))
     }
 
-    async getReceiveLibrary(receiver: Address, srcEid: EndpointId): Promise<[string | undefined, boolean]> {
+    async getReceiveLibrary(
+        receiver: Address,
+        srcEid: EndpointId
+    ): Promise<[address: Address | undefined, isDefault: boolean]> {
         return await this.contract.contract.getReceiveLibrary(receiver, srcEid)
     }
 
     async setDefaultReceiveLibrary(
         eid: EndpointId,
-        lib: string | null | undefined,
+        lib: Address | null | undefined,
         gracePeriod: number = 0
     ): Promise<OmniTransaction> {
         const data = this.contract.contract.interface.encodeFunctionData('setDefaultReceiveLibrary', [
@@ -33,7 +36,7 @@ export class Endpoint extends OmniSDK implements IEndpoint {
         }
     }
 
-    async getDefaultSendLibrary(eid: EndpointId): Promise<string | undefined> {
+    async getDefaultSendLibrary(eid: EndpointId): Promise<Address | undefined> {
         return ignoreZero(await this.contract.contract.defaultSendLibrary(eid))
     }
 
@@ -53,7 +56,7 @@ export class Endpoint extends OmniSDK implements IEndpoint {
         return this.contract.contract.isRegisteredLibrary(lib)
     }
 
-    async registerLibrary(lib: string): Promise<OmniTransaction> {
+    async registerLibrary(lib: Address): Promise<OmniTransaction> {
         const data = this.contract.contract.interface.encodeFunctionData('registerLibrary', [lib])
 
         return {
diff --git a/packages/protocol-utils/src/endpoint/types.ts b/packages/protocol-utils/src/endpoint/types.ts
index 1507a4973..4789f8dea 100644
--- a/packages/protocol-utils/src/endpoint/types.ts
+++ b/packages/protocol-utils/src/endpoint/types.ts
@@ -1,4 +1,11 @@
-import type { Address, OmniGraph, OmniPointBasedFactory, OmniTransaction, IOmniSDK } from '@layerzerolabs/utils'
+import type {
+    Address,
+    OmniGraph,
+    OmniPointBasedFactory,
+    OmniTransaction,
+    IOmniSDK,
+    Bytes32,
+} from '@layerzerolabs/utils'
 import type { EndpointId } from '@layerzerolabs/lz-definitions'
 
 export interface IEndpoint extends IOmniSDK {
@@ -15,8 +22,11 @@ export interface IEndpoint extends IOmniSDK {
     isRegisteredLibrary(lib: Address): Promise<boolean>
     registerLibrary(lib: Address): Promise<OmniTransaction>
 
-    getSendLibrary(sender: Address, dstEid: EndpointId): Promise<string | undefined>
-    getReceiveLibrary(receiver: Address, srcEid: EndpointId): Promise<[string | undefined, boolean]>
+    getSendLibrary(sender: Address, dstEid: EndpointId): Promise<Address | undefined>
+    getReceiveLibrary(
+        receiver: Address,
+        srcEid: EndpointId
+    ): Promise<[address: Bytes32 | undefined, isDefault: boolean]>
 }
 
 export interface EndpointEdgeConfig {