diff options
Diffstat (limited to 'p2p/protocols/protocol.go')
| -rw-r--r-- | p2p/protocols/protocol.go | 30 |
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 |