From 571871e35af1453f4d72f546c441239666f45f38 Mon Sep 17 00:00:00 2001
From: rkonfj <rkonfj@gmail.com>
Date: Mon, 29 Apr 2024 21:24:32 +0800
Subject: [PATCH] p2p: add SharedKey func

---
 p2p/conn.go                                 | 8 ++++++++
 secure/aescbc/aescbc.go                     | 4 ++++
 secure/chacha20poly1305/chacha20poly1305.go | 4 ++++
 secure/symm_algo.go                         | 1 +
 4 files changed, 17 insertions(+)

diff --git a/p2p/conn.go b/p2p/conn.go
index cb9e56e..acd5989 100644
--- a/p2p/conn.go
+++ b/p2p/conn.go
@@ -240,6 +240,14 @@ func (c *PeerPacketConn) runControlEventLoop(wsConn *disco.WSConn, udpConn *disc
 	}
 }
 
+// SharedKey get the key shared with the peer
+func (c *PeerPacketConn) SharedKey(peerID peer.ID) ([]byte, error) {
+	if c.cfg.SymmAlgo == nil {
+		return nil, errors.New("get shared key from plain conn")
+	}
+	return c.cfg.SymmAlgo.SecretKey()(peerID.String())
+}
+
 // ListenPacket listen the p2p network for read/write packets
 func ListenPacket(peermap *peermap.Peermap, opts ...Option) (*PeerPacketConn, error) {
 	id := make([]byte, 16)
diff --git a/secure/aescbc/aescbc.go b/secure/aescbc/aescbc.go
index 507aa89..e94d80a 100644
--- a/secure/aescbc/aescbc.go
+++ b/secure/aescbc/aescbc.go
@@ -124,6 +124,10 @@ func (s *AESCBC) Decrypt(b []byte, pubKey string) ([]byte, error) {
 	return PKCS7UnPadding(plainBytes)
 }
 
+func (s *AESCBC) SecretKey() secure.ProvideSecretKey {
+	return s.provideSecretKey
+}
+
 func (s *AESCBC) ensureChiperBlock(pubKey string) (cipher.Block, error) {
 	s.mut.RLock()
 	block, ok := s.cipher.Get(pubKey)
diff --git a/secure/chacha20poly1305/chacha20poly1305.go b/secure/chacha20poly1305/chacha20poly1305.go
index 2932032..da1a28a 100644
--- a/secure/chacha20poly1305/chacha20poly1305.go
+++ b/secure/chacha20poly1305/chacha20poly1305.go
@@ -62,6 +62,10 @@ func (s *Chacha20Poly1305) Decrypt(data []byte, pubKey string) ([]byte, error) {
 	return plain, nil
 }
 
+func (s *Chacha20Poly1305) SecretKey() secure.ProvideSecretKey {
+	return s.provideSecretKey
+}
+
 func (s *Chacha20Poly1305) ensureChiperAEAD(pubKey string) (cipher.AEAD, error) {
 	s.mut.RLock()
 	aead, ok := s.cipher.Get(pubKey)
diff --git a/secure/symm_algo.go b/secure/symm_algo.go
index 987a093..0b2710c 100644
--- a/secure/symm_algo.go
+++ b/secure/symm_algo.go
@@ -5,4 +5,5 @@ type ProvideSecretKey func(pubKey string) ([]byte, error)
 type SymmAlgo interface {
 	Encrypt(data []byte, pubKey string) ([]byte, error)
 	Decrypt(data []byte, pubKey string) ([]byte, error)
+	SecretKey() ProvideSecretKey
 }