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 }