diff --git a/p2p/protocols/accounting.go b/p2p/protocols/accounting.go
new file mode 100644
index 0000000000..95713ca6f6
--- /dev/null
+++ b/p2p/protocols/accounting.go
@@ -0,0 +1,159 @@
+// Copyright 2018 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see .
+
+package protocols
+
+import "github.com/ethereum/go-ethereum/metrics"
+
+//define some metrics
+var (
+ //NOTE: these metrics just define the interfaces and are currently *NOT persisted* over sessions
+ //All metrics are cumulative
+
+ //total amount of units credited
+ mBalanceCredit = metrics.NewRegisteredCounterForced("account.balance.credit", nil)
+ //total amount of units debited
+ mBalanceDebit = metrics.NewRegisteredCounterForced("account.balance.debit", nil)
+ //total amount of bytes credited
+ mBytesCredit = metrics.NewRegisteredCounterForced("account.bytes.credit", nil)
+ //total amount of bytes debited
+ mBytesDebit = metrics.NewRegisteredCounterForced("account.bytes.debit", nil)
+ //total amount of credited messages
+ mMsgCredit = metrics.NewRegisteredCounterForced("account.msg.credit", nil)
+ //total amount of debited messages
+ mMsgDebit = metrics.NewRegisteredCounterForced("account.msg.debit", nil)
+ //how many times local node had to drop remote peers
+ mPeerDrops = metrics.NewRegisteredCounterForced("account.peerdrops", nil)
+ //how many times local node overdrafted and dropped
+ mSelfDrops = metrics.NewRegisteredCounterForced("account.selfdrops", nil)
+ //how many cheques have been issued
+ //mChequesIssued = metrics.NewRegisteredCounterForced("account.cheques.issued", nil)
+ //how many cheques have been received
+ //mChequesReceived = metrics.NewRegisteredCounterForced("account.cheques.received", nil)
+)
+
+//Prices defines how prices are being passed on to the accounting instance
+type Prices interface {
+ //Return the Price for a message
+ Price(interface{}) *Price
+}
+
+type Payer bool
+
+const (
+ Sender = Payer(true)
+ Receiver = Payer(false)
+)
+
+//Price represents the costs of a message
+type Price struct {
+ Value uint64 //
+ PerByte bool //True if the price is per byte or for unit
+ Payer Payer
+}
+
+//ForSender gives back the price for sending a message
+func (p *Price) For(payer Payer, size uint32) int64 {
+ price := p.Value
+ if p.PerByte {
+ price *= uint64(size)
+ }
+ if p.Payer == payer {
+ return 0 - int64(price)
+ }
+ return int64(price)
+}
+
+//Balance is the actual accounting instance
+//Balance defines the operations needed for accounting
+//Implementations internally maintain the balance for every peer
+type Balance interface {
+ //Adds amount to the local balance with remote node `peer`;
+ //positive amount = credit local node
+ //negative amount = debit local node
+ Add(amount int64, peer *Peer) error
+}
+
+//Accounting implements the Hook interface
+//It interfaces to the balances through the Balance interface,
+//while interfacing with protocols and its prices through the Prices interface
+type Accounting struct {
+ Balance //interface to accounting logic
+ Prices //interface to prices logic
+}
+
+func NewAccounting(balance Balance, po Prices) *Accounting {
+ ah := &Accounting{
+ Prices: po,
+ Balance: balance,
+ }
+ return ah
+}
+
+//Implement Hook.Send
+func (ah *Accounting) Send(peer *Peer, size uint32, msg interface{}) error {
+ //get the price for a message (through the protocol spec)
+ price := ah.Price(msg)
+ //this message doesn't need accounting
+ if price == nil {
+ return nil
+ }
+ //evaluate the price for sending messages
+ costToLocalNode := price.For(Sender, size)
+ //do the accounting
+ err := ah.Add(costToLocalNode, peer)
+ //record metrics
+ ah.doMetrics(costToLocalNode, size, err)
+ return err
+}
+
+//Implement Hook.Receive
+func (ah *Accounting) Receive(peer *Peer, size uint32, msg interface{}) error {
+ //get the price for a message (through the protocol spec)
+ price := ah.Price(msg)
+ //this message doesn't need accounting
+ if price == nil {
+ return nil
+ }
+ //evaluate the price for receiving messages
+ costToLocalNode := price.For(Receiver, size)
+ //do the accounting
+ err := ah.Add(costToLocalNode, peer)
+ //record metrics
+ ah.doMetrics(costToLocalNode, size, err)
+ return err
+}
+
+//record some metrics
+func (ah *Accounting) doMetrics(price int64, size uint32, err error) {
+ /*
+ if price > 0 {
+ mBalanceCredit.Inc(int64(price))
+ mBytesCredit.Inc(int64(size))
+ mMsgCredit.Inc(1)
+ if err != nil {
+ mPeerDrops.Inc(1)
+ }
+ } else {
+ mBalanceDebit.Inc(int64(price))
+ mBytesDebit.Inc(int64(size))
+ mMsgDebit.Inc(1)
+ if err != nil {
+ mSelfDrops.Inc(1)
+ }
+ }
+ */
+}
diff --git a/p2p/protocols/accounting_test.go b/p2p/protocols/accounting_test.go
new file mode 100644
index 0000000000..fade8bcb52
--- /dev/null
+++ b/p2p/protocols/accounting_test.go
@@ -0,0 +1,356 @@
+// Copyright 2018 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see .
+
+package protocols
+
+import (
+ "bytes"
+ "context"
+ "testing"
+ "time"
+
+ "github.com/ethereum/go-ethereum/p2p"
+ "github.com/ethereum/go-ethereum/p2p/simulations/adapters"
+ "github.com/ethereum/go-ethereum/rlp"
+)
+
+//dummy Balance implementation
+type dummyBalance struct {
+ amount int64
+ peer *Peer
+}
+
+//dummy Prices implementation
+type dummyPrices struct{}
+
+//a dummy message which needs size based accounting
+//sender pays
+type perBytesMsgSenderPays struct {
+ Content string
+}
+
+//a dummy message which needs size based accounting
+//receiver pays
+type perBytesMsgReceiverPays struct {
+ Content string
+}
+
+//a dummy message which is paid for per unit
+//sender pays
+type perUnitMsgSenderPays struct{}
+
+//receiver pays
+type perUnitMsgReceiverPays struct{}
+
+//a dummy message which has zero as its price
+type zeroPriceMsg struct{}
+
+//a dummy message which has no accounting
+type nilPriceMsg struct{}
+
+//return the price for the defined messages
+func (d *dummyPrices) Price(msg interface{}) *Price {
+ switch msg.(type) {
+ //size based message cost, receiver pays
+ case *perBytesMsgReceiverPays:
+ return &Price{
+ PerByte: true,
+ Value: uint64(100),
+ Payer: Receiver,
+ }
+ //size based message cost, sender pays
+ case *perBytesMsgSenderPays:
+ return &Price{
+ PerByte: true,
+ Value: uint64(100),
+ Payer: Sender,
+ }
+ //unitary cost, receiver pays
+ case *perUnitMsgReceiverPays:
+ return &Price{
+ PerByte: false,
+ Value: uint64(99),
+ Payer: Receiver,
+ }
+ //unitary cost, sender pays
+ case *perUnitMsgSenderPays:
+ return &Price{
+ PerByte: false,
+ Value: uint64(99),
+ Payer: Sender,
+ }
+ case *zeroPriceMsg:
+ return &Price{
+ PerByte: false,
+ Value: uint64(0),
+ Payer: Sender,
+ }
+ case *nilPriceMsg:
+ return nil
+ }
+ return nil
+}
+
+//dummy accounting implementation, only stores values for later check
+func (d *dummyBalance) Add(amount int64, peer *Peer) error {
+ d.amount = amount
+ d.peer = peer
+ return nil
+}
+
+type testCase struct {
+ msg interface{}
+ size uint32
+ sendResult int64
+ recvResult int64
+}
+
+//lowest level unit test
+func TestBalance(t *testing.T) {
+ //create instances
+ balance := &dummyBalance{}
+ prices := &dummyPrices{}
+ //create the spec
+ spec := createTestSpec()
+ //create the accounting hook for the spec
+ acc := NewAccounting(balance, prices)
+ //create a peer
+ id := adapters.RandomNodeConfig().ID
+ p := p2p.NewPeer(id, "testPeer", nil)
+ peer := NewPeer(p, &dummyRW{}, spec)
+ //price depends on size, receiver pays
+ msg := &perBytesMsgReceiverPays{Content: "testBalance"}
+ size, _ := rlp.EncodeToBytes(msg)
+
+ testCases := []testCase{
+ {
+ msg,
+ uint32(len(size)),
+ int64(len(size) * 100),
+ int64(len(size) * -100),
+ },
+ {
+ &perBytesMsgSenderPays{Content: "testBalance"},
+ uint32(len(size)),
+ int64(len(size) * -100),
+ int64(len(size) * 100),
+ },
+ {
+ &perUnitMsgSenderPays{},
+ 0,
+ int64(-99),
+ int64(99),
+ },
+ {
+ &perUnitMsgReceiverPays{},
+ 0,
+ int64(99),
+ int64(-99),
+ },
+ {
+ &zeroPriceMsg{},
+ 0,
+ int64(0),
+ int64(0),
+ },
+ {
+ &nilPriceMsg{},
+ 0,
+ int64(0),
+ int64(0),
+ },
+ }
+ checkAccountingTestCases(t, testCases, acc, peer, balance, true)
+ checkAccountingTestCases(t, testCases, acc, peer, balance, false)
+}
+
+func TestBalanceWithPeer(t *testing.T) {
+ //create instances
+ balance := &dummyBalance{}
+ prices := &dummyPrices{}
+ //create the spec
+ spec := createTestSpec()
+ //create the accounting hook for the spec
+ spec.Hook = NewAccounting(balance, prices)
+ //create a peer
+ id := adapters.RandomNodeConfig().ID
+ p := p2p.NewPeer(id, "testPeer", nil)
+ peer := NewPeer(p, &dummyRW{}, spec)
+ //price depends on size, receiver pays
+ msg := &perBytesMsgReceiverPays{Content: "testBalance"}
+ size, _ := rlp.EncodeToBytes(msg)
+ testCases := []testCase{
+ {
+ msg,
+ uint32(len(size)),
+ int64(len(size) * 100),
+ int64(len(size) * -100),
+ },
+ {
+ &perBytesMsgSenderPays{Content: "testBalance"},
+ uint32(len(size)),
+ int64(len(size) * -100),
+ int64(len(size) * 100),
+ },
+ {
+ &perUnitMsgSenderPays{},
+ 0,
+ int64(-99),
+ int64(99),
+ },
+ {
+ &perUnitMsgReceiverPays{},
+ 0,
+ int64(99),
+ int64(-99),
+ },
+ {
+ &zeroPriceMsg{},
+ 0,
+ int64(0),
+ int64(0),
+ },
+ {
+ &nilPriceMsg{},
+ 0,
+ int64(0),
+ int64(0),
+ },
+ }
+ checkPeerTestCases(t, testCases, peer, balance, spec, true)
+ checkPeerTestCases(t, testCases, peer, balance, spec, false)
+}
+
+func TestAccountedMsgSimulation(t *testing.T) {
+
+ /*
+ using this type of simulation introduces circular dependency
+ and drags the swap package into this package
+
+ sim := simulation.New(map[string]simulation.ServiceFunc{
+ "dummyService": func(ctx *adapters.ServiceContext, bucket *sync.Map) (s node.Service, cleanup func(), err error) {
+ },
+ })
+ defer sim.Close()
+ */
+}
+
+func checkAccountingTestCases(t *testing.T, cases []testCase, acc *Accounting, peer *Peer, balance *dummyBalance, send bool) {
+ for _, c := range cases {
+ var err error
+ var expectedResult int64
+ //reset balance before every check
+ balance.amount = 0
+ if send {
+ err = acc.Send(peer, c.size, c.msg)
+ expectedResult = c.sendResult
+ } else {
+ err = acc.Receive(peer, c.size, c.msg)
+ expectedResult = c.recvResult
+ }
+
+ checkResults(t, err, balance, peer, expectedResult)
+ }
+}
+
+func checkPeerTestCases(t *testing.T, cases []testCase, peer *Peer, balance *dummyBalance, spec *Spec, send bool) {
+ for _, c := range cases {
+ var err error
+ var expectedResult int64
+ //reset balance before every check
+ balance.amount = 0
+ var ctx context.Context
+ if send {
+ err = peer.Send(ctx, c.msg)
+ expectedResult = c.sendResult
+ } else {
+ peer.rw.(*dummyRW).msg = c.msg
+ peer.rw.(*dummyRW).code, _ = spec.GetCode(c.msg)
+ err = peer.handleIncoming(func(ctx context.Context, msg interface{}) error {
+ return nil
+ })
+ expectedResult = c.recvResult
+ }
+
+ checkResults(t, err, balance, peer, expectedResult)
+ }
+}
+
+func checkResults(t *testing.T, err error, balance *dummyBalance, peer *Peer, result int64) {
+ if err != nil {
+ t.Fatal(err)
+ }
+ if balance.peer != peer {
+ t.Fatalf("expected Add to be called with peer %v, got %v", peer, balance.peer)
+ }
+ if balance.amount != result {
+ t.Fatalf("Expected balance to be %d but is %d", result, balance.amount)
+ }
+}
+
+//dummy implementation of a MsgReadWriter
+//this allows for quick and easy unit tests without
+//having to build up the complete protocol
+type dummyRW struct {
+ msg interface{}
+ size uint32
+ code uint64
+}
+
+func (d *dummyRW) WriteMsg(msg p2p.Msg) error {
+ return nil
+}
+
+func (d *dummyRW) ReadMsg() (p2p.Msg, error) {
+ enc := bytes.NewReader(d.getDummyMsg())
+ return p2p.Msg{
+ Code: d.code,
+ Size: d.size,
+ Payload: enc,
+ ReceivedAt: time.Now(),
+ }, nil
+}
+
+func (d *dummyRW) getDummyMsg() []byte {
+ r, _ := rlp.EncodeToBytes(d.msg)
+ var b bytes.Buffer
+ wmsg := WrappedMsg{
+ Context: b.Bytes(),
+ Size: uint32(len(r)),
+ Payload: r,
+ }
+ rr, _ := rlp.EncodeToBytes(wmsg)
+ d.size = uint32(len(rr))
+ return rr
+}
+
+//create a test spec
+func createTestSpec() *Spec {
+ spec := &Spec{
+ Name: "test",
+ Version: 42,
+ MaxMsgSize: 10 * 1024,
+ Messages: []interface{}{
+ perBytesMsgReceiverPays{},
+ perBytesMsgSenderPays{},
+ perUnitMsgReceiverPays{},
+ perUnitMsgSenderPays{},
+ zeroPriceMsg{},
+ nilPriceMsg{},
+ },
+ }
+ return spec
+}
diff --git a/p2p/protocols/protocol.go b/p2p/protocols/protocol.go
index 615f74b569..b6329aa2c0 100644
--- a/p2p/protocols/protocol.go
+++ b/p2p/protocols/protocol.go
@@ -32,6 +32,7 @@ import (
"bufio"
"bytes"
"context"
+ "errors"
"fmt"
"io"
"reflect"
@@ -122,6 +123,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 +152,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 +288,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 +359,19 @@ 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 {
+ if wmsg.Size != uint32(len(wmsg.Payload)) {
+ errMsg := "Advertised message size and payload length don't match"
+ log.Warn(errMsg)
+ p.Drop(errors.New(errMsg))
+ }
+ 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
diff --git a/p2p/protocols/protocol_test.go b/p2p/protocols/protocol_test.go
index 4755db3e62..e26c2aab84 100644
--- a/p2p/protocols/protocol_test.go
+++ b/p2p/protocols/protocol_test.go
@@ -185,6 +185,162 @@ func runProtoHandshake(t *testing.T, proto *protoHandshake, errs ...error) {
}
}
+type dummyHook struct {
+ peer *Peer
+ size uint32
+ msg interface{}
+ send bool
+ err error
+ waitC chan struct{}
+}
+
+type dummyMsg struct {
+ Content string
+}
+
+func (d *dummyHook) Send(peer *Peer, size uint32, msg interface{}) error {
+ d.peer = peer
+ d.size = size
+ d.msg = msg
+ d.send = true
+ return d.err
+}
+
+func (d *dummyHook) Receive(peer *Peer, size uint32, msg interface{}) error {
+ d.peer = peer
+ d.size = size
+ d.msg = msg
+ d.send = false
+ d.waitC <- struct{}{}
+ return d.err
+}
+
+func TestProtocolHook(t *testing.T) {
+ testHook := &dummyHook{
+ waitC: make(chan struct{}, 1),
+ }
+ spec := &Spec{
+ Name: "test",
+ Version: 42,
+ MaxMsgSize: 10 * 1024,
+ Messages: []interface{}{
+ dummyMsg{},
+ },
+ Hook: testHook,
+ }
+
+ runFunc := func(p *p2p.Peer, rw p2p.MsgReadWriter) error {
+ peer := NewPeer(p, rw, spec)
+ ctx := context.TODO()
+ peer.Send(ctx, &dummyMsg{
+ Content: "handshake"})
+
+ handle := func(ctx context.Context, msg interface{}) error {
+ return nil
+ }
+
+ return peer.Run(handle)
+ }
+
+ conf := adapters.RandomNodeConfig()
+ tester := p2ptest.NewProtocolTester(t, conf.ID, 2, runFunc)
+ err := tester.TestExchanges(p2ptest.Exchange{
+ Expects: []p2ptest.Expect{
+ {
+ Code: 0,
+ Msg: &dummyMsg{Content: "handshake"},
+ Peer: tester.Nodes[0].ID(),
+ },
+ },
+ })
+ if err != nil {
+ t.Fatal(err)
+ }
+ if testHook.msg == nil || testHook.msg.(*dummyMsg).Content != "handshake" {
+ t.Fatal("Expected msg to be set, but it is not")
+ }
+ if testHook.send != true {
+ t.Fatal("Expected a send message, but it is not")
+ }
+ if testHook.peer == nil || testHook.peer.ID() != tester.Nodes[0].ID() {
+ t.Fatal("Expected peer ID to be set correctly, but it is not")
+ }
+ if testHook.size != 11 { //11 is the length of the encoded message
+ t.Fatalf("Expected size to be %d, but it is %d ", 1, testHook.size)
+ }
+
+ err = tester.TestExchanges(p2ptest.Exchange{
+ Triggers: []p2ptest.Trigger{
+ {
+ Code: 0,
+ Msg: &dummyMsg{Content: "response"},
+ Peer: tester.Nodes[1].ID(),
+ },
+ },
+ })
+
+ <-testHook.waitC
+
+ if err != nil {
+ t.Fatal(err)
+ }
+ if testHook.msg == nil || testHook.msg.(*dummyMsg).Content != "response" {
+ t.Fatal("Expected msg to be set, but it is not")
+ }
+ if testHook.send != false {
+ t.Fatal("Expected a send message, but it is not")
+ }
+ if testHook.peer == nil || testHook.peer.ID() != tester.Nodes[1].ID() {
+ t.Fatal("Expected peer ID to be set correctly, but it is not")
+ }
+ if testHook.size != 10 { //11 is the length of the encoded message
+ t.Fatalf("Expected size to be %d, but it is %d ", 1, testHook.size)
+ }
+
+ testHook.err = fmt.Errorf("dummy error")
+ err = tester.TestExchanges(p2ptest.Exchange{
+ Triggers: []p2ptest.Trigger{
+ {
+ Code: 0,
+ Msg: &dummyMsg{Content: "response"},
+ Peer: tester.Nodes[1].ID(),
+ },
+ },
+ })
+
+ <-testHook.waitC
+
+ time.Sleep(100 * time.Millisecond)
+ err = tester.TestDisconnected(&p2ptest.Disconnect{tester.Nodes[1].ID(), testHook.err})
+ if err != nil {
+ t.Fatalf("Expected a specific disconnect error, but got different one: %v", err)
+ }
+
+}
+
+//We need to test that if the hook is not defined, then message infrastructure
+//(send,receive) still works
+func TestNoHook(t *testing.T) {
+ //create a test spec
+ spec := createTestSpec()
+ //a random node
+ id := adapters.RandomNodeConfig().ID
+ //a peer
+ p := p2p.NewPeer(id, "testPeer", nil)
+ rw := &dummyRW{}
+ peer := NewPeer(p, rw, spec)
+ ctx := context.TODO()
+ msg := &perBytesMsgSenderPays{Content: "testBalance"}
+ //send a message
+ peer.Send(ctx, msg)
+ //simulate receiving a message
+ rw.msg = msg
+ peer.handleIncoming(func(ctx context.Context, msg interface{}) error {
+ return nil
+ })
+ //all should just work and not result in any error
+}
+
func TestProtoHandshakeVersionMismatch(t *testing.T) {
runProtoHandshake(t, &protoHandshake{41, "420"}, errorf(ErrHandshake, errorf(ErrHandler, "(msg code 0): 41 (!= 42)").Error()))
}