From 345e696a7d68225b553c8947fe222c115ff6d5c1 Mon Sep 17 00:00:00 2001
From: Alex Stokes <r.alex.stokes@gmail.com>
Date: Tue, 3 Sep 2019 13:14:04 -0700
Subject: [PATCH] Add option to inline "short" public keys for peer IDs

Added to remain interoperable w/ the Go implementation
---
 libp2p/peer/id.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/libp2p/peer/id.py b/libp2p/peer/id.py
index f75563c..4d19ee9 100644
--- a/libp2p/peer/id.py
+++ b/libp2p/peer/id.py
@@ -11,6 +11,32 @@ from libp2p.crypto.keys import PublicKey
 # intended for debugging, logging, etc.
 FRIENDLY_IDS = True
 
+# NOTE: On inlining...
+# See: https://github.com/libp2p/specs/issues/138
+# NOTE: enabling to be interoperable w/ the Go implementation
+ENABLE_INLINING = True
+MAX_INLINE_KEY_LENGTH = 42
+
+IDENTITY_MULTIHASH_CODE = 0x00
+
+if ENABLE_INLINING:
+
+    class IdentityHash:
+        _digest: bytes
+
+        def __init__(self) -> None:
+            self._digest = bytearray()
+
+        def update(self, input: bytes) -> None:
+            self._digest += input
+
+        def digest(self) -> bytes:
+            return self._digest
+
+    multihash.FuncReg.register(
+        IDENTITY_MULTIHASH_CODE, "identity", hash_new=lambda: IdentityHash()
+    )
+
 
 class ID:
     _bytes: bytes
@@ -68,6 +94,8 @@ class ID:
     def from_pubkey(cls, key: PublicKey) -> "ID":
         serialized_key = key.serialize()
         algo = multihash.Func.sha2_256
+        if ENABLE_INLINING and len(serialized_key) <= MAX_INLINE_KEY_LENGTH:
+            algo = IDENTITY_MULTIHASH_CODE
         mh_digest = multihash.digest(serialized_key, algo)
         return cls(mh_digest.encode())