From 7ee0fdafe3cbc0a8949012ab3a2b1493ed9663cd Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Wed, 15 Jul 2020 14:00:47 +1000
Subject: [PATCH] Add some client callbacks

---
 callbacks.go | 12 ++++++++++++
 client.go    |  3 +++
 config.go    |  2 ++
 peerconn.go  |  3 +++
 4 files changed, 20 insertions(+)
 create mode 100644 callbacks.go

diff --git a/callbacks.go b/callbacks.go
new file mode 100644
index 0000000000..e603997152
--- /dev/null
+++ b/callbacks.go
@@ -0,0 +1,12 @@
+package torrent
+
+import (
+	pp "github.com/anacrolix/torrent/peer_protocol"
+)
+
+// These are called synchronously, and do not pass ownership. The Client and other locks may still
+// be held. nil functions are not called.
+type Callbacks struct {
+	CompletedHandshake func(_ *PeerConn, infoHash InfoHash)
+	ReadMessage        func(*PeerConn, *pp.Message)
+}
diff --git a/client.go b/client.go
index fe7786ac4e..69ef1ef5d3 100644
--- a/client.go
+++ b/client.go
@@ -860,6 +860,9 @@ func (cl *Client) connBtHandshake(c *PeerConn, ih *metainfo.Hash) (ret metainfo.
 	c.PeerExtensionBytes = res.PeerExtensionBits
 	c.PeerID = res.PeerID
 	c.completedHandshake = time.Now()
+	if cb := cl.config.Callbacks.CompletedHandshake; cb != nil {
+		cb(c, res.Hash)
+	}
 	return
 }
 
diff --git a/config.go b/config.go
index ae9b4ea09b..c28ee12456 100644
--- a/config.go
+++ b/config.go
@@ -133,6 +133,8 @@ type ClientConfig struct {
 
 	DisableWebtorrent bool
 	DisableWebseeds   bool
+
+	Callbacks Callbacks
 }
 
 func (cfg *ClientConfig) SetListenAddr(addr string) *ClientConfig {
diff --git a/peerconn.go b/peerconn.go
index 1c2879ee30..aff37ff9c0 100644
--- a/peerconn.go
+++ b/peerconn.go
@@ -1059,6 +1059,9 @@ func (c *PeerConn) mainReadLoop() (err error) {
 			defer cl.lock()
 			err = decoder.Decode(&msg)
 		}()
+		if cb := cl.config.Callbacks.ReadMessage; cb != nil && err == nil {
+			cb(c, &msg)
+		}
 		if t.closed.IsSet() || c.closed.IsSet() {
 			return nil
 		}