1 files changed, 30 insertions, 0 deletions
diff --git a/p2p/protocols/protocol.go b/p2p/protocols/protocol.go
index 615f74b56..7dddd852f 100644
--- a/p2p/protocols/protocol.go
+++ b/p2p/protocols/protocol.go
@@ -122,6 +122,16 @@ type WrappedMsg struct {
 	Payload []byte
 }
 
+//For accounting, the design is to allow the Spec to describe which and how its messages are priced
+//To access this functionality, we provide a Hook interface which will call accounting methods
+//NOTE: there could be more such (horizontal) hooks in the future
+type Hook interface {
+	//A hook for sending messages
+	Send(peer *Peer, size uint32, msg interface{}) error
+	//A hook for receiving messages
+	Receive(peer *Peer, size uint32, msg interface{}) error
+}
+
 // Spec is a protocol specification including its name and version as well as
 // the types of messages which are exchanged
 type Spec struct {
@@ -141,6 +151,9 @@ type Spec struct {
 	// each message must have a single unique data type
 	Messages []interface{}
 
+	//hook for accounting (could be extended to multiple hooks in the future)
+	Hook Hook
+
 	initOnce sync.Once
 	codes    map[reflect.Type]uint64
 	types    map[uint64]reflect.Type
@@ -274,6 +287,15 @@ func (p *Peer) Send(ctx context.Context, msg interface{}) error {
 		Payload: r,
 	}
 
+	//if the accounting hook is set, call it
+	if p.spec.Hook != nil {
+		err := p.spec.Hook.Send(p, wmsg.Size, msg)
+		if err != nil {
+			p.Drop(err)
+			return err
+		}
+	}
+
 	code, found := p.spec.GetCode(msg)
 	if !found {
 		return errorf(ErrInvalidMsgType, "%v", code)
@@ -336,6 +358,14 @@ func (p *Peer) handleIncoming(handle func(ctx context.Context, msg interface{})
 		return errorf(ErrDecode, "<= %v: %v", msg, err)
 	}
 
+	//if the accounting hook is set, call it
+	if p.spec.Hook != nil {
+		err := p.spec.Hook.Receive(p, wmsg.Size, val)
+		if err != nil {
+			return err
+		}
+	}
+
 	// call the registered handler callbacks
 	// a registered callback take the decoded message as argument as an interface
 	// which the handler is supposed to cast to the appropriate type