swarm: swap implementation with event subscription to p2p server

swarm/network/stream/delivery.go

@@ -19,8 +19,6 @@ package stream
 import (
 	"context"
 	"errors"
-	"math/big"
-	"time"
 
 	"fmt"
 
@@ -138,11 +136,6 @@ type RetrieveRequestMsg struct {
 	HopCount  uint8
 }
 
-//TODO: what is the correct price
-func (rrm *RetrieveRequestMsg) Price() *big.Int {
-	return big.NewInt(int64(4096))
-}
-
 func (d *Delivery) handleRetrieveRequestMsg(ctx context.Context, sp *Peer, req *RetrieveRequestMsg) error {
 	log.Trace("received request", "peer", sp.ID(), "hash", req.Addr)
 	handleRetrieveRequestMsgCount.Inc(1)

swarm/network/stream/peer.go

@@ -31,7 +31,6 @@ import (
 	"github.com/ethereum/go-ethereum/swarm/spancontext"
 	"github.com/ethereum/go-ethereum/swarm/state"
 	"github.com/ethereum/go-ethereum/swarm/storage"
-	"github.com/ethereum/go-ethereum/swarm/swap"
 	opentracing "github.com/opentracing/opentracing-go"
 )
 
@@ -54,8 +53,7 @@ var ErrMaxPeerServers = errors.New("max peer servers")
 
 // Peer is the Peer extension for the streaming protocol
 type Peer struct {
-	//*protocols.Peer
+	*protocols.Peer
-	*swap.SwapPeer
 	streamer *Registry
 	pq       *pq.PriorityQueue
 	serverMu sync.RWMutex
@@ -77,7 +75,7 @@ type WrappedPriorityMsg struct {
 // NewPeer is the constructor for Peer
 func NewPeer(peer *protocols.Peer, streamer *Registry) *Peer {
 	p := &Peer{
-		SwapPeer:     swap.NewSwapPeer(peer, streamer.swap),
+		Peer:         peer,
 		pq:           pq.New(int(PriorityQueue), PriorityQueueCap),
 		streamer:     streamer,
 		servers:      make(map[Stream]*server),

swarm/network/stream/stream.go

@@ -34,8 +34,6 @@ import (
 	"github.com/ethereum/go-ethereum/swarm/pot"
 	"github.com/ethereum/go-ethereum/swarm/state"
 	"github.com/ethereum/go-ethereum/swarm/storage"
-	"github.com/ethereum/go-ethereum/swarm/swap"
-	opentracing "github.com/opentracing/opentracing-go"
 )
 
 const (
@@ -104,8 +102,6 @@ func NewRegistry(localID enode.ID, delivery *Delivery, syncChunkStore storage.Sy
 	RegisterSwarmSyncerServer(streamer, syncChunkStore)
 	RegisterSwarmSyncerClient(streamer, syncChunkStore)
 
-	streamer.swap = swap
-
 	if options.DoSync {
 		// latestIntC function ensures that
 		//   - receiving from the in chan is not blocked by processing inside the for loop
@@ -385,7 +381,7 @@ func (r *Registry) Run(p *network.BzzPeer) error {
 		}
 	}
 
-	return sp.RunAccountedProtocol(sp.HandleMsg)
+	return sp.Run(sp.HandleMsg)
 }
 
 // updateSyncing subscribes to SYNC streams by iterating over the

swarm/swap/api.go

@@ -30,7 +30,7 @@ var (
 
 //This is the API definition to access swarm swap accounting data via RPC
 type API struct {
-	*Protocol
+	swap *Swap
 }
 
 //TODO: define metrics
@@ -39,13 +39,13 @@ type SwapMetrics struct {
 }
 
 //Create a new API instance
-func NewAPI(swap *Protocol) *API {
+func NewAPI(swap *Swap) *API {
-	return &API{Protocol: swap}
+	return &API{swap: swap}
 }
 
 //Get the balance for this node with a specific peer
 func (swapapi *API) BalanceWithPeer(ctx context.Context, peer discover.NodeID) (balance *big.Int, err error) {
-	balance = swapapi.swap.peers[peer].balance
+	balance = swapapi.swap.peers[peer]
 	if balance == nil {
 		err = ErrNoSuchPeerAccounting
 	}
@@ -58,8 +58,8 @@ func (swapapi *API) BalanceWithPeer(ctx context.Context, peer discover.NodeID) (
 //It assumes that if a disfavorable balance is represented as a negative value
 func (swapapi *API) Balance(ctx context.Context) (balance *big.Int, err error) {
 	balance = big.NewInt(0)
-	for _, peer := range swapapi.swap.peers {
+	for _, peerBalance := range swapapi.swap.peers {
-		balance.Add(balance, peer.balance)
+		balance.Add(balance, peerBalance)
 	}
 	return
 }

swarm/swap/protocol.go

@@ -21,7 +21,6 @@ import (
 	"fmt"
 	"sync"
 
-	"github.com/ethereum/go-ethereum/metrics"
 	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/p2p/discover"
 	"github.com/ethereum/go-ethereum/p2p/protocols"
@@ -42,29 +41,25 @@ const (
 type Protocol struct {
 	peersMu sync.RWMutex
 	peers   map[discover.NodeID]*Peer
-	swap    *Swap
 }
 
 //This is the peer representing a participant in this protocol
 type Peer struct {
 	*protocols.Peer
-	swapProtocol *Protocol
 }
 
 //Create a new protocol instance
-func NewSwapProtocol(swapAccount *Swap) *Protocol {
+func NewProtocol() *Protocol {
 	proto := &Protocol{
 		peers: make(map[discover.NodeID]*Peer),
-		swap:  swapAccount,
 	}
 	return proto
 }
 
 // NewPeer is the constructor for the protocol Peer
-func NewPeer(peer *protocols.Peer, swap *Protocol) *Peer {
+func NewPeer(peer *protocols.Peer) *Peer {
 	p := &Peer{
 		Peer: peer,
-		swapProtocol: swap,
 	}
 	return p
 }
@@ -79,17 +74,19 @@ type IssueChequeMsg struct {
 //redeem a cheque, which means it kicks off the process to cash it in.
 //In this case, this message is sent to peer A
 type RedeemChequeMsg struct {
+	Cheque *Cheque
 }
 
 /////////////////////////////////////////////////////////////////////
 // SECTION: node.Service interface
 /////////////////////////////////////////////////////////////////////
-func (p *Protocol) Start(srv *p2p.Server) error {
+func (s *Swap) Start(srv *p2p.Server) error {
+	s.registerForEvents(srv)
 	log.Debug("Started swap")
 	return nil
 }
 
-func (p *Protocol) Stop() error {
+func (s *Swap) Stop() error {
 	log.Info("swap shutting down")
 	return nil
 }
@@ -104,23 +101,23 @@ var swapSpec = &protocols.Spec{
 	},
 }
 
-func (p *Protocol) Protocols() []p2p.Protocol {
+func (s *Swap) Protocols() []p2p.Protocol {
 	return []p2p.Protocol{
 		{
 			Name:    swapSpec.Name,
 			Version: swapSpec.Version,
 			Length:  swapSpec.Length(),
-			Run:     p.Run,
+			Run:     s.run,
 		},
 	}
 }
 
-func (p *Protocol) APIs() []rpc.API {
+func (s *Swap) APIs() []rpc.API {
 	apis := []rpc.API{
 		{
 			Namespace: "swap",
 			Version:   "1.0",
-			Service:   NewAPI(p),
+			Service:   NewAPI(s),
 			Public:    true,
 		},
 	}
@@ -130,55 +127,52 @@ func (p *Protocol) APIs() []rpc.API {
 /////////////////////////////////////////////////////////////////////
 // SECTION: p2p.protocol interface
 /////////////////////////////////////////////////////////////////////
-func (swap *Protocol) Run(peer *p2p.Peer, rw p2p.MsgReadWriter) error {
+func (s *Swap) run(peer *p2p.Peer, rw p2p.MsgReadWriter) error {
 	p := protocols.NewPeer(peer, rw, swapSpec)
-	sp := NewPeer(p, swap)
+	sp := NewPeer(p)
-	swap.setPeer(sp)
+	s.protocol.setPeer(sp)
-	defer swap.deletePeer(sp)
+	defer s.protocol.deletePeer(sp)
-	defer swap.Close()
+	defer s.protocol.Close()
 	return sp.Run(sp.handleSwapMsg)
 }
 
-func (swap *Protocol) NodeInfo() interface{} {
+func (s *Swap) NodeInfo() interface{} {
 	return nil
 }
 
-func (swap *Protocol) PeerInfo(id discover.NodeID) interface{} {
+func (s *Swap) PeerInfo(id discover.NodeID) interface{} {
 	return nil
 }
 
 //------------------------------------------------------------------------------------------
-func (swap *Protocol) Close() {
+func (proto *Protocol) Close() {
 }
 
-func (swap *Protocol) getPeer(peerId discover.NodeID) *Peer {
+func (proto *Protocol) getPeer(peerId discover.NodeID) *Peer {
-	swap.peersMu.RLock()
+	proto.peersMu.RLock()
-	defer swap.peersMu.RUnlock()
+	defer proto.peersMu.RUnlock()
 
-	return swap.peers[peerId]
+	return proto.peers[peerId]
 }
 
-func (swap *Protocol) setPeer(peer *Peer) {
+func (proto *Protocol) setPeer(peer *Peer) {
-	swap.peersMu.Lock()
+	proto.peersMu.Lock()
-	defer swap.peersMu.Unlock()
+	defer proto.peersMu.Unlock()
 
-	swap.peers[peer.ID()] = peer
+	proto.peers[peer.ID()] = peer
-	metrics.GetOrRegisterGauge("swap.peers", nil).Update(int64(len(swap.peers)))
 }
 
-func (swap *Protocol) deletePeer(peer *Peer) {
+func (proto *Protocol) deletePeer(peer *Peer) {
-	swap.peersMu.Lock()
+	proto.peersMu.Lock()
-	defer swap.peersMu.Unlock()
+	defer proto.peersMu.Unlock()
 
-	delete(swap.peers, peer.ID())
+	delete(proto.peers, peer.ID())
-	metrics.GetOrRegisterGauge("swap.peers", nil).Update(int64(len(swap.peers)))
-	swap.peersMu.Unlock()
 }
 
-func (swap *Protocol) peersCount() (c int) {
+func (proto *Protocol) peersCount() (c int) {
-	swap.peersMu.Lock()
+	proto.peersMu.Lock()
-	c = len(swap.peers)
+	c = len(proto.peers)
-	swap.peersMu.Unlock()
+	proto.peersMu.Unlock()
 	return
 }
 
@@ -198,13 +192,13 @@ func (p *Peer) handleSwapMsg(ctx context.Context, msg interface{}) error {
 }
 
 //A IssueChequeMsg has been received
-func (sp *Peer) handleIssueChequeMsg(ctx context.Context, msg interface{}) (err error) {
+func (p *Peer) handleIssueChequeMsg(ctx context.Context, msg interface{}) (err error) {
 	log.Debug("SwapProtocolPeer: handleIssueChequeMsg")
 	return err
 }
 
 //A RedeemChequeMsg has been received
-func (sp *Peer) handleRedeemChequeMsg(ctx context.Context, msg interface{}) (err error) {
+func (p *Peer) handleRedeemChequeMsg(ctx context.Context, msg interface{}) (err error) {
 	log.Debug("SwapProtocolPeer: handleRedeemChequeMsg")
 	return err
 }

swarm/swap/swap.go

@@ -26,7 +26,6 @@ import (
 
 	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/p2p/discover"
-	"github.com/ethereum/go-ethereum/p2p/protocols"
 	"github.com/ethereum/go-ethereum/swarm/log"
 	"github.com/ethereum/go-ethereum/swarm/network/stream"
 	"github.com/ethereum/go-ethereum/swarm/state"
@@ -42,8 +41,8 @@ var (
 	payAt  = big.NewInt(-4096 * 10000) // threshold that triggers payment {request} (bytes)
 	dropAt = big.NewInt(-4096 * 12000) // threshold that triggers disconnect (bytes)
 
-	ErrInsufficientFunds = errors.New("Insufficient funds")
 	ErrNotAccountedMsg   = errors.New("Message does not need accounting")
+	ErrInsufficientFunds = errors.New("Insufficient funds")
 )
 
 const (
@@ -60,63 +59,80 @@ type Swap struct {
 	chequeManager *ChequeManager
 	stateStore    state.Store
 	lock          sync.RWMutex
-	peers         map[discover.NodeID]*SwapPeer
+	peers         map[discover.NodeID]*big.Int
-}
+	p2pServer     *p2p.Server
-
+	protocol      *Protocol
-//Protocols which want to send and handle priced messages will need to use
-//this peer instead of protocols.Peer, which is embedded
-type SwapPeer struct {
-	*protocols.Peer
-	lock        sync.RWMutex
-	swapAccount *Swap
-	handlerFunc func(context.Context, interface{}) error
-	balance     *big.Int
-	storeID     string
 }
 
+//PriceOracle is responsible to maintain the price matrix for accounted messages,
+//to get its price and to evaluate if a message is accountable or not
 type PriceOracle interface {
 	IsAccountedMsg(event *p2p.PeerEvent) bool
-	GetPriceForMsg(event *p2p.PeerEvent) *PriceTag
+	GetPriceForMsg(event *p2p.PeerEvent) (*big.Int, EntryDirection)
 }
 
 //This defines if a price will be debited or credited to an account
 type EntryDirection bool
 
+//For some messages the sender pays (e.g. RetrieveRequestMsg),
+//for others the receiver (ChunkDeliveryMsg)
 const (
 	ChargeSender   EntryDirection = true
 	ChargeReceiver EntryDirection = false
 )
 
+//An accountable message needs some meta information attached to it
+//in order to evaluate the correct price
 type PriceTag struct {
 	Price     *big.Int
 	SizeBased bool
 	Direction EntryDirection
 }
 
+//The default price oracle
 type DefaultPriceOracle struct {
-	priceMatrix map[string]map[int]*PriceTag
+	priceMatrix map[string]map[uint64]*PriceTag
 }
 
+//Load a default price matrix
 func (dpo *DefaultPriceOracle) LoadPriceMatrix() {
-	priceMatrix := dpo.loadDefaultPriceMatrix()
+	dpo.priceMatrix = dpo.loadDefaultPriceMatrix()
 }
 
-func (dpo *DefaultPriceOracle) loadDefaultPriceMatrix() map[string]map[int]*PriceTag {
+//Set up a default price matrix
-	priceMatrix := make(map[string]map[int]*big.Int)
+//This statically builds the price matrix according to previous knowledge
+//of which messages need accounting
+//The matrix can be queried by
+// * the protocol string, returns a sub-map
+// * the sub-map by the message code, which returns the PriceTag
+func (dpo *DefaultPriceOracle) loadDefaultPriceMatrix() map[string]map[uint64]*PriceTag {
+	//setup matrix
+	priceMatrix := make(map[string]map[uint64]*PriceTag)
 
+	//define accounted messages from the stream protocol
 	streamProtocol := stream.Spec.Name
+	priceMatrix[streamProtocol] = make(map[uint64]*PriceTag)
 
-	priceMatrix[streamProtocol] = make(map[int]*PriceTag)
+	//get indexes of accounted messages in order to populate the map
+	retrieveRequestMsgIndex, ok := stream.Spec.GetCode(stream.RetrieveRequestMsg{})
+	if !ok {
+		log.Crit("setting up price matrix failed; expected message code not found in Spec!")
+		return nil
+	}
+	deliveryMsgIndex, ok := stream.Spec.GetCode(stream.RetrieveRequestMsg{})
+	if !ok {
+		log.Crit("setting up price matrix failed; expected message code not found in Spec!")
+		return nil
+	}
 
-	retrieveRequestMsgIndex := dpo.getIndexForMsgType(stream.RetrieveRequestMsg, stream.Spec.Messages)
+	//assign the price tag to the message
-	deliveryMsgIndex := dpo.getIndexForMsgType(stream.ChunkDeliveryMsg, stream.Spec.Messages)
+	priceMatrix[streamProtocol][uint64(retrieveRequestMsgIndex)] = &PriceTag{
-
-	priceMatrix[streamProtocol][retrieveRequestMsgIndex] = &PriceTag{
 		Price:     big.NewInt(10),
 		SizeBased: false,
 		Direction: ChargeSender,
 	}
-	priceMatrix[streamProtocol][deliveryMsgIndex] = &PriceTag{
+
+	priceMatrix[streamProtocol][uint64(deliveryMsgIndex)] = &PriceTag{
 		Price:     big.NewInt(100),
 		SizeBased: true,
 		Direction: ChargeReceiver,
@@ -126,59 +142,57 @@ func (dpo *DefaultPriceOracle) loadDefaultPriceMatrix() map[string]map[int]*Pric
 
 }
 
-func (dpo *DefaultPriceOracle) getIndexForMsgType(iface interface{}, types []interface{}) {
+//Check if a message needs accounting
-	for i, v := range types {
-		if iface == v {
-			return i
-		}
-	}
-	return -1
-}
-
 func (dpo *DefaultPriceOracle) IsAccountedMsg(event *p2p.PeerEvent) bool {
-	if protoPriceMap, ok := dpo.priceMatrix[event.Protocol]; !ok {
+	var protoPriceMap map[uint64]*PriceTag
+	var ok bool
+	if protoPriceMap, ok = dpo.priceMatrix[event.Protocol]; !ok {
 		return false
 	}
-	if msgCode, ok := protoPriceMap[event.MsgCode]; !ok {
+	if _, ok = protoPriceMap[*event.MsgCode]; !ok {
 		return false
 	}
 	return true
 }
 
+//Get the actual price of a message
+//If the message is size-based, it returns the calculated price
 func (dpo *DefaultPriceOracle) GetPriceForMsg(event *p2p.PeerEvent) (*big.Int, EntryDirection) {
 	if dpo.IsAccountedMsg(event) {
-		priceTag := dpo.priceMatrix[event.Protocol][event.MsgCode]
+		priceTag := dpo.priceMatrix[event.Protocol][*event.MsgCode]
-		price := priceTag.Price
+		price := &big.Int{}
+		price.Set(priceTag.Price)
 		if priceTag.SizeBased {
-			price = price * event.MsgSize
+			price.Mul(priceTag.Price, big.NewInt(int64(*event.MsgSize)))
 		}
 		return price, priceTag.Direction
 	}
 	return nil, false
 }
 
-//A message which needs accounting needs to implement this interface
+//This swap implementation works by listening to message events on the p2p server.
-type PricedMsg interface {
+//It then handles the event received, filtering for messages and evaluating
-	Price() *big.Int
+//if it needs accounting
-}
+func (s *Swap) registerForEvents(srv *p2p.Server) {
-
-func (s *Swap) RegisterForEvents(srv *p2p.Server) {
 	go func() {
 		events := make(chan *p2p.PeerEvent)
-		sub := server.SubscribeEvents(events)
+		sub := srv.SubscribeEvents(events)
 		defer sub.Unsubscribe()
 
 		for {
 			select {
 			case event := <-events:
-				go handleMsgEvent(event)
+				go s.handleMsgEvent(event)
-			case <-sub.Err():
+			case err := <-sub.Err():
+				log.Error(err.Error())
 				return
 			}
 		}
 	}()
 }
 
+//Handle the message.
+//Determine if it needs accounting, and if yes, account for it
 func (s *Swap) handleMsgEvent(event *p2p.PeerEvent) {
 	if !s.priceOracle.IsAccountedMsg(event) {
 		return
@@ -186,6 +200,9 @@ func (s *Swap) handleMsgEvent(event *p2p.PeerEvent) {
 	s.AccountMsgForPeer(event)
 }
 
+//Do the accounting
+//Depending on the charging type of the message (set in the PriceTag),
+//it will charge the sender or receiver
 func (s *Swap) AccountMsgForPeer(event *p2p.PeerEvent) {
 	price, direction := s.priceOracle.GetPriceForMsg(event)
 	if price == nil {
@@ -194,231 +211,137 @@ func (s *Swap) AccountMsgForPeer(event *p2p.PeerEvent) {
 	}
 
 	if direction == ChargeSender {
+		//we are sending a ChargeSender message, thus debit us and credit remote
 		if event.Type == p2p.PeerEventTypeMsgSend {
-			s.chargeSender(event, price)
+			s.chargeLocal(event, price)
+			//we are receiving a ChargeSender message, so credit us and debit remote
 		} else if event.Type == p2p.PeerEventTypeMsgRecv {
-			s.chargeReceiver(event, price)
+			s.chargeRemote(event, price)
 		}
 	} else {
+		//we are receiving a ChargeReceiver message, thus debit us and credit remote
 		if event.Type == p2p.PeerEventTypeMsgRecv {
-			s.chargeReceiver(event, price)
+			s.chargeLocal(event, price)
+			//we are sending a ChargeReceiver message, thus credit us and debit remote
 		} else if event.Type == p2p.PeerEventTypeMsgSend {
-			s.chargeSender(event, price)
+			s.chargeRemote(event, price)
 		}
 	}
-
 }
 
-//Handler for received messages
+//Debit us and credit remote
-func (sp *SwapPeer) RunAccountedProtocol(protocolHandler func(ctx context.Context, msg interface{}) error) error {
+func (s *Swap) chargeLocal(event *p2p.PeerEvent, amount *big.Int) {
-	//the `peer.Run` function is a loop, so in order to pre-/post-process a message with accounting,
+	s.lock.Lock()
-	//we need to save the actual handler
+	defer s.lock.Unlock()
-	sp.handlerFunc = protocolHandler
-
-	//then run the handler loop function
-	return sp.Run(sp.handle)
-}
-
-//get a peer's balance
-func (swap *Swap) GetPeerBalance(peer discover.NodeID) *big.Int {
-	if p, ok := swap.peers[peer]; ok {
-		return p.balance
-	}
-	return nil
-}
-
-//Handle a received message; this is the handler loop function.
-//Check if it needs accounting, and if yes, apply accounting logic:
-//Check for sufficient funds, perform operation, then account
-func (sp *SwapPeer) handle(ctx context.Context, msg interface{}) error {
-	var err error
-	var price *big.Int
-
-	//the message is one which needs accounting...
-	//only account if swapAccount != nil (== swap is disabled)
-	if _, ok := msg.(PricedMsg); ok && sp.swapAccount != nil {
-		//..so first check if there are enough funds for the operation available
-		//(for crediting, this means if we are not essentially "overdrafting", or crossing the threshold)
-		price, err = sp.checkAvailableFunds(ctx, msg, CreditEntry)
-		//if not (or some other error occured), return error
-		if err != nil {
-			//also, if the error is indeed insufficient funds, then disconnect the peer
-			if err == ErrInsufficientFunds {
-				log.Error("Insufficient funds, dropping peer")
-				sp.Drop(err)
-			}
-			return err
-		}
-		//at this point we know there are sufficient funds, so process the message
-		err = sp.handlerFunc(ctx, msg)
-		if err == nil {
-			//and if no errors occurred, finally book the entry
-			sp.AccountMsgForPeer(ctx, msg, price, CreditEntry)
-		}
-	} else {
-		//this message doesn't need accounting, so just process it
-		err = sp.handlerFunc(ctx, msg)
-	}
-	return err
-}
-
-//Send a message
-//Check if it needs accounting, and if yes, apply accounting logic:
-//Check for sufficient funds, perform operation, then account
-func (sp *SwapPeer) Send(ctx context.Context, msg interface{}) error {
-	var err error
-	var price *big.Int
-
-	//the message is one which needs accounting...
-	//only account if swapAccount != nil (== swap is disabled)
-	if _, ok := msg.(PricedMsg); ok && sp.swapAccount != nil {
-		//..so first check if there are enough funds for the operation available
-		price, err = sp.checkAvailableFunds(ctx, msg, DebitEntry)
-		//if not (or some other error occured), return error
-		if err != nil {
-			//also, if the error is indeed insufficient funds, then disconnect the peer
-			if err == ErrInsufficientFunds {
-				log.Error("Insufficient funds, dropping peer")
-				sp.Drop(err)
-			}
-			return err
-		}
-		//at this point we know there are sufficient funds, so process the message
-		err = sp.Peer.Send(ctx, msg)
-		if err == nil {
-			//and if no errors occurred, finally book the entry
-			sp.AccountMsgForPeer(ctx, msg, price, DebitEntry)
-		}
-	} else {
-		//this message doesn't need accounting, so just process it
-		err = sp.Peer.Send(ctx, msg)
-	}
-	return err
-}
-
-//check that the operation has enough funds available
-func (sp *SwapPeer) checkAvailableFunds(ctx context.Context, msg interface{}, direction EntryDirection) (*big.Int, error) {
-	sp.lock.Lock()
-	defer sp.lock.Unlock()
-
-	if accounted, ok := msg.(PricedMsg); ok {
-		price := accounted.Price()
-		//local node is being credited (in its favor), so check upper limit
-		if direction == CreditEntry {
-			//TODO: is there a check needed here?
-			//It should actually have been done on the client side, the debitor!
-			//creditor could theoretically go over payAt, but if well done,
-			//should have been checked on the client side so this shouldn't happen?
-			checkBalance := &big.Int{}
-			checkBalance.Add(sp.balance, price)
-			//(checkBalance *Int) CmpAbs(payAt)
-			// -1 if |checkBalance|  <  |payAt|
-			//  0 if |checkBalance|  == |payAt|
-			// +1 if |checkBalance|  >  |payAt|
-			if checkBalance.CmpAbs(payAt) == 1 {
-				return nil, ErrInsufficientFunds
-			}
-		} else if direction == DebitEntry {
-			//NOTE: ErrInsufficientFunds should only be returned
-			//if the dropAt is exceeded, but should be ignored for payAt,
-			//as there is a "clemency" margin between triggering the check
-			//and actually disconnecting the peer
-
-			//(checkBalance *Int) Cmp(dropAt)
-			// -1 if checkBalance  <  dropAt
-			//  0 if checkBalance  ==	dropAt
-			// +1 if checkBalance  >  dropAt
-			checkBalance := &big.Int{}
-			checkBalance.Sub(sp.balance, price.Abs(price))
-			if checkBalance.Cmp(dropAt) == -1 {
-				return nil, ErrInsufficientFunds
-			}
-		}
-		return price, nil
-	}
-	return nil, ErrNotAccountedMsg
-}
-
-//The balance is accounted from the point of view of the local node
-//Thus, we credit the balance and increase it when the amount is in favor of the local node
-//We debit the balance and decrease it when the amount is in favor of the remote peer
-func (sp *SwapPeer) AccountMsgForPeer(ctx context.Context, msg interface{}, price *big.Int, direction EntryDirection) {
-	if _, ok := msg.(PricedMsg); ok {
-		sp.lock.Lock()
-		defer sp.lock.Unlock()
 	//local node is being credited (in its favor), so its balance increases
-		if direction == CreditEntry {
+	if s.peers[event.Peer] == nil {
-			//NOTE: do we need to check for sufficient funds again?
+		s.peers[event.Peer] = &big.Int{}
-			//operations are not atomic/transactional, so balance may have changed in the meanwhile!
+		s.stateStore.Get(event.Peer.String(), s.peers[event.Peer])
-			sp.balance.Add(sp.balance, price)
-			//local node is being debited (in favor of remote peer), so its balance decreases
-		} else if direction == DebitEntry {
-			sp.balance.Sub(sp.balance, price)
 	}
+	peerBalance := s.peers[event.Peer]
+	peerBalance.Sub(peerBalance, amount)
+	//local node is being debited (in favor of remote peer), so its balance decreases
 	//TODO: save to store here? init store?
-		sp.swapAccount.stateStore.Put(sp.storeID, sp.balance)
+	s.stateStore.Put(event.Peer.String(), peerBalance)
-		//(sp.balance *Int) Cmp(payAt)
+	//(balance *Int) Cmp(payAt)
-		// -1 if sp.balance <  payAt
+	// -1 if balance <  payAt
-		//  0 if sp.balance == payAt
+	//  0 if balance == payAt
-		// +1 if sp.balance >  payAt
+	// +1 if balance >  payAt
-		if sp.balance.Cmp(payAt) == -1 {
+	if peerBalance.Cmp(payAt) == -1 {
-			err := sp.issueCheque(ctx)
+		ctx := context.TODO()
+		err := s.issueCheque(ctx, event.Peer)
 		if err != nil {
 			//TODO: special error handling, as at this point the accounting has been done
 			//but the cheque could not be sent?
 			log.Warn("Payment threshold exceeded, but error sending cheque!", "err", err)
 		}
 	}
-		if sp.balance.Cmp(dropAt) == -1 {
+	if peerBalance.Cmp(dropAt) == -1 {
-			sp.Drop(ErrInsufficientFunds)
+		peer := s.protocol.getPeer(event.Peer)
+		if peer != nil {
+			peer.Drop(ErrInsufficientFunds)
+			//this little hack allows for tests to verify that this error occurred
+			event.Error = ErrInsufficientFunds.Error()
 		}
-		log.Debug(fmt.Sprintf("balance for peer %s: %s", sp.ID(), sp.balance.String()))
 	}
+	log.Debug(fmt.Sprintf("balance for peer %s: %s", event.Peer.String(), peerBalance.String()))
+}
+
+//Credit us and debit remote
+func (s *Swap) chargeRemote(event *p2p.PeerEvent, amount *big.Int) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	//local node is being credited (in its favor), so its balance increases
+	if s.peers[event.Peer] == nil {
+		s.peers[event.Peer] = &big.Int{}
+	}
+	peerBalance := s.peers[event.Peer]
+	peerBalance.Add(peerBalance, amount)
+	//local node is being credited(in favor of local peer), so its balance increases
+	//TODO: save to store here? init store?
+	s.stateStore.Put(event.Peer.String(), peerBalance)
+	//(balance *Int) Cmp(payAt)
+	// -1 if balance <  payAt
+	//  0 if balance == payAt
+	// +1 if balance >  payAt
+	if peerBalance.Cmp(payAt) == -1 {
+		ctx := context.TODO()
+		err := s.issueCheque(ctx, event.Peer)
+		if err != nil {
+			//TODO: special error handling, as at this point the accounting has been done
+			//but the cheque could not be sent?
+			log.Warn("Payment threshold exceeded, but error sending cheque!", "err", err)
+		}
+	}
+	if peerBalance.Cmp(dropAt) == -1 {
+		peer := s.protocol.getPeer(event.Peer)
+		if peer != nil {
+			//this little hack allows for tests to verify that this error occurred
+			peer.Drop(ErrInsufficientFunds)
+		}
+	}
+	log.Debug(fmt.Sprintf("balance for peer %s: %s", event.Peer.String(), peerBalance.String()))
+}
+
+//get a peer's balance
+func (swap *Swap) GetPeerBalance(peer discover.NodeID) *big.Int {
+	swap.lock.RLock()
+	defer swap.lock.RUnlock()
+	if p, ok := swap.peers[peer]; ok {
+		return p
+	}
+	return nil
 }
 
 //Issue a cheque for the remote peer. Happens if we are indebted with the peer
 //and crossed the payment threshold
-func (sp *SwapPeer) issueCheque(ctx context.Context) error {
+func (s *Swap) issueCheque(ctx context.Context, id discover.NodeID) error {
 	amount := &big.Int{}
-	cheque := sp.swapAccount.chequeManager.CreateCheque(sp.ID(), amount.Abs(payAt))
+	cheque := s.chequeManager.CreateCheque(id, amount.Abs(payAt))
 	msg := IssueChequeMsg{
 		Cheque: cheque,
 	}
 	//TODO: This should now be via the actual SwapProtocol
-	return sp.Send(ctx, msg)
+	p := s.protocol.getPeer(id)
-}
+	if p == nil {
-
+		return fmt.Errorf("wanting to send to non-connected peer!")
-//Create a new swap accounted peer
-func NewSwapPeer(peer *protocols.Peer, swap *Swap) *SwapPeer {
-	sp := &SwapPeer{
-		Peer:        peer,
-		swapAccount: swap,
-		storeID:     peer.String()[:24] + "-swap",
 	}
-	//swap is not enabled
+	return p.Send(ctx, msg)
-	if swap != nil {
-		//check if there is one already in the stateStore and load it
-		balance := &big.Int{}
-		swap.stateStore.Get(peer.String()[:24]+"-swap", &balance)
-		sp.balance = balance
-		swap.lock.Lock()
-		defer swap.lock.Unlock()
-		swap.peers[peer.ID()] = sp
-	}
-	return sp
 }
 
 // New - swap constructor
-func New(stateStore state.Store, srv *p2p.Server) (swap *Swap, err error) {
+func New(stateStore state.Store) (swap *Swap) {
+
+	priceOracle := &DefaultPriceOracle{}
 
 	swap = &Swap{
 		chequeManager: NewChequeManager(stateStore),
 		stateStore:    stateStore,
-		peers:         make(map[discover.NodeID]*SwapPeer),
+		peers:         make(map[discover.NodeID]*big.Int),
-		priceOracle:   &DefaultPriceOracle{},
+		priceOracle:   priceOracle,
+		protocol:      NewProtocol(),
 	}
-
+	priceOracle.LoadPriceMatrix()
-	swap.RegisterForEvents(srv)
 
 	return
 }

swarm/swap/swap_test.go

@@ -1,4 +1,4 @@
-// Copyright 2018 The go-ethereum Authors
+// Copyreeeeight 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
@@ -54,6 +54,8 @@ var testSpec = &protocols.Spec{
 		testExceedsPayAtMsg{},
 		testExceedsDropAtMsg{},
 		testCheapMsg{},
+		testSizeBasedMsg{},
+		testChargeRecvMsg{},
 	},
 }
 
@@ -79,6 +81,8 @@ func (d *dummyRW) ReadMsg() (p2p.Msg, error) {
 type testExceedsPayAtMsg struct{}
 type testExceedsDropAtMsg struct{}
 type testCheapMsg struct{}
+type testSizeBasedMsg struct{}
+type testChargeRecvMsg struct{}
 
 //this message is just one unit more expensive than the payment threshold
 func (tmsg *testExceedsPayAtMsg) Price() *big.Int {
@@ -99,6 +103,16 @@ func (tmsg *testCheapMsg) Price() *big.Int {
 	return big.NewInt(100)
 }
 
+//a message which needs to be charged based on price
+func (tmsg *testSizeBasedMsg) Price() *big.Int {
+	return big.NewInt(222)
+}
+
+//a message which needs to be charged based on price
+func (tmsg *testChargeRecvMsg) Price() *big.Int {
+	return big.NewInt(999)
+}
+
 func init() {
 	flag.Parse()
 
@@ -124,20 +138,32 @@ func TestExceedsPayAt(t *testing.T) {
 	swap, testDir := createTestSwap(t)
 	defer os.RemoveAll(testDir)
 
-	//create a dummy peer
+	swap.priceOracle = setupPriceMatrix()
-	testPeer := newDummyPeer()
-	sp := NewSwapPeer(testPeer, swap)
-	sp.handlerFunc = dummyMsgHandler
 
-	//send the message
+	node := createAndStartSvcNode(swap, t)
-	ctx := context.Background()
+	defer node.Stop()
-	err := sp.Send(ctx, &testExceedsPayAtMsg{})
+	defer os.RemoveAll(node.DataDir())
-	if err != nil {
+
-		t.Fatal("Unecpected error on sending message", "err", err)
+	peerID := adapters.RandomNodeConfig().ID
+
+	code, ok := testSpec.GetCode(&testExceedsPayAtMsg{})
+	if !ok {
+		t.Fatal("test message not found in spec")
 	}
 
+	event := &p2p.PeerEvent{
+		Type:     p2p.PeerEventTypeMsgSend,
+		Protocol: testSpec.Name,
+		Peer:     peerID,
+		MsgCode:  &code,
+		MsgSize:  new(uint32),
+		Error:    "",
+	}
+
+	swap.handleMsgEvent(event)
+
 	//check that a cheque is present
-	cheques := sp.swapAccount.chequeManager.openDebitCheques[sp.ID()]
+	cheques := swap.chequeManager.openDebitCheques[peerID]
 	if cheques == nil {
 		t.Fatal("Expected cheques for this peer to be present, but are nil")
 	}
@@ -157,44 +183,83 @@ func TestExceedsPayAt(t *testing.T) {
 //tests that a message is being sent which crosses the drop limit
 //in that case, we should receive a InsufficientFunds error
 func TestExceedsDropAt(t *testing.T) {
+	//create a test swap account
 	swap, testDir := createTestSwap(t)
 	defer os.RemoveAll(testDir)
 
-	testPeer := newDummyPeer()
+	swap.priceOracle = setupPriceMatrix()
-	sp := NewSwapPeer(testPeer, swap)
-	sp.handlerFunc = dummyMsgHandler
 
-	ctx := context.Background()
+	node := createAndStartSvcNode(swap, t)
-	err := sp.Send(ctx, &testExceedsDropAtMsg{})
+	defer node.Stop()
-	if err != ErrInsufficientFunds {
+	defer os.RemoveAll(node.DataDir())
-		t.Fatal("Expected test to fail with insufficient funds, but it didn't")
+
+	//peerID := adapters.RandomNodeConfig().ID
+	peer := newDummyPeer()
+	swap.protocol.setPeer(peer.Peer)
+
+	code, ok := testSpec.GetCode(&testExceedsDropAtMsg{})
+	if !ok {
+		t.Fatal("test message not found in spec")
+	}
+
+	event := &p2p.PeerEvent{
+		Type:     p2p.PeerEventTypeMsgSend,
+		Protocol: testSpec.Name,
+		Peer:     peer.ID(),
+		MsgCode:  &code,
+		MsgSize:  new(uint32),
+		Error:    "",
+	}
+
+	swap.handleMsgEvent(event)
+	if event.Error != ErrInsufficientFunds.Error() {
+		t.Fatal("Excpected this test to fail with insufficient funds, but it did not")
 	}
 }
 
 //send a message with cost,
 //then check that the balance has the expected amount
 func TestSendCheapMessage(t *testing.T) {
+	fmt.Println("send")
 	swap, testDir := createTestSwap(t)
 	defer os.RemoveAll(testDir)
 
-	testPeer := newDummyPeer()
+	swap.priceOracle = setupPriceMatrix()
-	sp := NewSwapPeer(testPeer, swap)
+
+	node := createAndStartSvcNode(swap, t)
+	defer node.Stop()
+	defer os.RemoveAll(node.DataDir())
+
+	//peerID := adapters.RandomNodeConfig().ID
+	peer := newDummyPeer()
+	swap.protocol.setPeer(peer.Peer)
 	//set an arbitrary test balance value
 	testBalance := big.NewInt(1234567890)
-	sp.balance = testBalance
+	swap.peers[peer.ID()] = testBalance
 
-	//send the message
+	code, ok := testSpec.GetCode(&testCheapMsg{})
-	msg := &testCheapMsg{}
+	if !ok {
-	ctx := context.Background()
+		t.Fatal("test message not found in spec")
-	err := sp.Send(ctx, msg)
-	if err != nil {
-		t.Fatal("Unexpected error sending message")
 	}
 
+	event := &p2p.PeerEvent{
+		Type:     p2p.PeerEventTypeMsgSend,
+		Protocol: testSpec.Name,
+		Peer:     peer.ID(),
+		MsgCode:  &code,
+		MsgSize:  new(uint32),
+		Error:    "",
+	}
+
+	swap.handleMsgEvent(event)
+
+	peerBalance := swap.peers[peer.ID()]
+	msg := &testCheapMsg{}
+
 	//check the new balance
-	if sp.balance.Cmp(testBalance.Sub(testBalance, msg.Price())) != 0 {
+	if peerBalance.Cmp(testBalance.Sub(testBalance, msg.Price())) != 0 {
 		t.Fatal(fmt.Sprintf("Unexpected balance value after sending cheap message test. Expected balance: %s, balance is: %s",
-			testBalance.Sub(testBalance, msg.Price()).String(), sp.balance.String()))
+			testBalance.Sub(testBalance, msg.Price()).String(), peerBalance.String()))
 	}
 }
 
@@ -208,35 +273,142 @@ func TestRestoreBalanceFromStateStore(t *testing.T) {
 	swap, testDir := createTestSwap(t)
 	defer os.RemoveAll(testDir)
 
-	//create the dummy p2p protocol Peer
+	swap.priceOracle = setupPriceMatrix()
-	testPeer := newDummyPeer()
+
-	//create the "source" swap peer
+	node := createAndStartSvcNode(swap, t)
-	sp := NewSwapPeer(testPeer, swap)
+	defer node.Stop()
+	defer os.RemoveAll(node.DataDir())
+
+	//peerID := adapters.RandomNodeConfig().ID
+	peer := newDummyPeer()
+	swap.protocol.setPeer(peer.Peer)
 	//create a reference an arbitrary balance
 	testBalance := big.NewInt(1234567890)
 	//assign the same value to the peer
-	sp.balance = big.NewInt(1234567890)
+	swap.peers[peer.ID()] = big.NewInt(1234567890)
 
-	//send a message, should trigger saving to stateStore
+	code, ok := testSpec.GetCode(&testCheapMsg{})
-	msg := &testCheapMsg{}
+	if !ok {
-	ctx := context.Background()
+		t.Fatal("test message not found in spec")
-	err := sp.Send(ctx, msg)
-	if err != nil {
-		log.Error(err.Error())
-		t.Fatal("Unexpected error sending message")
 	}
 
-	//create a new peer with same underlying protocols.Peer
+	event := &p2p.PeerEvent{
-	//this will try to load the balance from the stateStore,
+		Type:     p2p.PeerEventTypeMsgSend,
-	//as it is the same discover.NodeID
+		Protocol: testSpec.Name,
-	sp2 := NewSwapPeer(testPeer, swap)
+		Peer:     peer.ID(),
+		MsgCode:  &code,
+		MsgSize:  new(uint32),
+		Error:    "",
+	}
+
+	swap.handleMsgEvent(event)
+
+	peerBalance := swap.peers[peer.ID()]
+	expectedBalance := &big.Int{}
+	swap.stateStore.Get(peer.ID().String(), expectedBalance)
 
 	//compare the balances
-	expectedBalance := &big.Int{}
+	expectedBalance.Sub(testBalance, (&testCheapMsg{}).Price())
-	expectedBalance.Sub(testBalance, msg.Price())
+	if peerBalance.Cmp(expectedBalance) != 0 {
-	if sp2.balance.Cmp(expectedBalance) != 0 {
 		t.Fatal(fmt.Sprintf("Unexpected balance value after sending cheap message test. Expected balance: %s, balance is: %s",
-			expectedBalance.String(), sp2.balance.String()))
+			expectedBalance.String(), peerBalance.String()))
+	}
+}
+
+//send a message with cost,
+//then check that the balance has the expected amount
+//the message is charged size based
+func TestSendSizeBasedMsg(t *testing.T) {
+	swap, testDir := createTestSwap(t)
+	defer os.RemoveAll(testDir)
+
+	swap.priceOracle = setupPriceMatrix()
+
+	node := createAndStartSvcNode(swap, t)
+	defer node.Stop()
+	defer os.RemoveAll(node.DataDir())
+
+	peer := newDummyPeer()
+	swap.protocol.setPeer(peer.Peer)
+	//set an arbitrary test balance value
+	testBalance := big.NewInt(1234567890)
+	swap.peers[peer.ID()] = testBalance
+
+	code, ok := testSpec.GetCode(&testSizeBasedMsg{})
+	if !ok {
+		t.Fatal("test message not found in spec")
+	}
+
+	size := new(uint32)
+	*size = 500
+
+	event := &p2p.PeerEvent{
+		Type:     p2p.PeerEventTypeMsgSend,
+		Protocol: testSpec.Name,
+		Peer:     peer.ID(),
+		MsgCode:  &code,
+		MsgSize:  size,
+		Error:    "",
+	}
+
+	swap.handleMsgEvent(event)
+
+	peerBalance := swap.peers[peer.ID()]
+	msg := &testSizeBasedMsg{}
+	expectedBalance := &big.Int{}
+	expectedBalance.Mul(msg.Price(), big.NewInt(int64(*size)))
+
+	//check the new balance
+	if peerBalance.Cmp(testBalance.Sub(testBalance, expectedBalance)) != 0 {
+		t.Fatal(fmt.Sprintf("Unexpected balance value after sending cheap message test. Expected balance: %s, balance is: %s",
+			expectedBalance.String(), peerBalance.String()))
+	}
+}
+
+//charge as being the receiver of a receiver-pays message
+//as this test for simplicity only sends messages,
+//it means that the balance is changed in positive,
+//instead of negative as with all other tests
+func TestChargeAsReceiver(t *testing.T) {
+	swap, testDir := createTestSwap(t)
+	defer os.RemoveAll(testDir)
+
+	swap.priceOracle = setupPriceMatrix()
+
+	node := createAndStartSvcNode(swap, t)
+	defer node.Stop()
+	defer os.RemoveAll(node.DataDir())
+
+	//peerID := adapters.RandomNodeConfig().ID
+	peer := newDummyPeer()
+	swap.protocol.setPeer(peer.Peer)
+	//set an arbitrary test balance value
+	testBalance := big.NewInt(1234567890)
+	swap.peers[peer.ID()] = testBalance
+
+	code, ok := testSpec.GetCode(&testChargeRecvMsg{})
+	if !ok {
+		t.Fatal("test message not found in spec")
+	}
+
+	event := &p2p.PeerEvent{
+		Type:     p2p.PeerEventTypeMsgSend,
+		Protocol: testSpec.Name,
+		Peer:     peer.ID(),
+		MsgCode:  &code,
+		MsgSize:  new(uint32),
+		Error:    "",
+	}
+
+	swap.handleMsgEvent(event)
+
+	peerBalance := swap.peers[peer.ID()]
+	msg := &testChargeRecvMsg{}
+
+	//check the new balance
+	if peerBalance.Cmp(testBalance.Add(testBalance, msg.Price())) != 0 {
+		t.Fatal(fmt.Sprintf("Unexpected balance value after sending cheap message test. Expected balance: %s, balance is: %s",
+			testBalance.Add(testBalance, msg.Price()).String(), peerBalance.String()))
 	}
 }
 
@@ -251,61 +423,97 @@ func createTestSwap(t *testing.T) (*Swap, string) {
 	if err2 != nil {
 		t.Fatal(err2)
 	}
-	swap, err3 := New(stateStore)
+	swap := New(stateStore)
-	if err3 != nil {
-		t.Fatal(err3)
-	}
 	return swap, dir
 }
 
+func setupPriceMatrix() PriceOracle {
+
+	priceOracle := &DefaultPriceOracle{}
+	priceOracle.priceMatrix = make(map[string]map[uint64]*PriceTag)
+	priceOracle.priceMatrix[testSpec.Name] = make(map[uint64]*PriceTag)
+	protoMatrix := priceOracle.priceMatrix[testSpec.Name]
+	protoMatrix[0] = &PriceTag{
+		Direction: ChargeSender,
+		Price:     (&testExceedsPayAtMsg{}).Price(),
+	}
+
+	protoMatrix[1] = &PriceTag{
+		Direction: ChargeSender,
+		Price:     (&testExceedsDropAtMsg{}).Price(),
+	}
+
+	protoMatrix[2] = &PriceTag{
+		Direction: ChargeSender,
+		Price:     (&testCheapMsg{}).Price(),
+	}
+
+	protoMatrix[3] = &PriceTag{
+		Direction: ChargeSender,
+		Price:     (&testSizeBasedMsg{}).Price(),
+		SizeBased: true,
+	}
+
+	protoMatrix[4] = &PriceTag{
+		Direction: ChargeReceiver,
+		Price:     (&testChargeRecvMsg{}).Price(),
+	}
+
+	return priceOracle
+}
+
 //dummy message handler (needed or we will have a panic in the accounting)
 func dummyMsgHandler(ctx context.Context, msg interface{}) error {
 	return nil
 }
 
-//tests some basic things over RPC
+func createAndStartSvcNode(swap *Swap, t *testing.T) *node.Node {
-func TestSwapRPC(t *testing.T) {
+	stack, err := newServiceNode(p2pPort, 0, 0)
-	swap, testDir := createTestSwap(t)
-	defer os.RemoveAll(testDir)
-
-	// create the two nodes
-	stack_one, err := newServiceNode(p2pPort, 0, 0)
 	if err != nil {
 		t.Fatal("Create servicenode #1 fail", "err", err)
 	}
 
-	instance := NewSwapProtocol(swap)
-	// wrapper function for servicenode to start the service
 	swapsvc := func(ctx *node.ServiceContext) (node.Service, error) {
-		return &API{
+		return swap, nil
-			Protocol: instance,
-		}, nil
 	}
 
-	// register adds the service to the services the servicenode starts when started
+	err = stack.Register(swapsvc)
-	err = stack_one.Register(swapsvc)
 	if err != nil {
 		t.Fatal("Register service in servicenode #1 fail", "err", err)
 	}
+
 	// start the nodes
-	err = stack_one.Start()
+	err = stack.Start()
 	if err != nil {
 		t.Fatal("servicenode #1 start failed", "err", err)
 	}
 
+	return stack
+}
+
+//tests some basic things over RPC
+func TestSwapRPC(t *testing.T) {
+
+	swap, testDir := createTestSwap(t)
+	defer os.RemoveAll(testDir)
+
+	stack := createAndStartSvcNode(swap, t)
+	defer stack.Stop()
+	defer os.RemoveAll(stack.DataDir())
+
 	// connect to the servicenode RPCs
-	rpcclient_one, err := rpc.Dial(filepath.Join(stack_one.DataDir(), ipcpath))
+	rpcclient, err := rpc.Dial(filepath.Join(stack.DataDir(), ipcpath))
 	if err != nil {
-		t.Fatal("connect to servicenode #1 IPC fail", "err", err)
+		t.Fatal("connect to servicenode IPC fail", "err", err)
 	}
-	defer os.RemoveAll(stack_one.DataDir())
+	defer os.RemoveAll(stack.DataDir())
 
 	var balance *big.Int
-	err = rpcclient_one.Call(&balance, "swap_balance")
+	err = rpcclient.Call(&balance, "swap_balance")
 	if err != nil {
-		t.Fatal("servicenode #1 RPC failed", "err", err)
+		t.Fatal("servicenode RPC failed", "err", err)
 	}
-	log.Debug("servicenode #1 balance", "balance-1", balance)
+	log.Debug("servicenode balance", "balance", balance)
 
 	if balance.Cmp(big.NewInt(0)) != 0 {
 		t.Fatal("Expected balance to be 0 but it is not")
@@ -321,32 +529,30 @@ func TestSwapRPC(t *testing.T) {
 	fakeBalance1 := big.NewInt(fake1)
 	fakeBalance2 := big.NewInt(fake2)
 
-	swap.peers[id1] = NewSwapPeer(dummyPeer1, swap)
+	swap.peers[id1] = fakeBalance1
-	swap.peers[id2] = NewSwapPeer(dummyPeer2, swap)
+	swap.peers[id2] = fakeBalance2
-	swap.peers[id1].balance = fakeBalance1
-	swap.peers[id2].balance = fakeBalance2
 
-	err = rpcclient_one.Call(&balance, "swap_balanceWithPeer", id1)
+	err = rpcclient.Call(&balance, "swap_balanceWithPeer", id1)
 	if err != nil {
-		t.Fatal("servicenode #1 RPC failed", "err", err)
+		t.Fatal("servicenode RPC failed", "err", err)
 	}
 	log.Debug("balance1", "balance-1", balance)
 	if balance.Cmp(fakeBalance1) != 0 {
 		t.Fatal(fmt.Sprintf("Expected balance %s to be equal to fake balance %s, but it is not", balance.String(), fakeBalance1.String()))
 	}
 
-	err = rpcclient_one.Call(&balance, "swap_balanceWithPeer", id2)
+	err = rpcclient.Call(&balance, "swap_balanceWithPeer", id2)
 	if err != nil {
-		t.Fatal("servicenode #1 RPC failed", "err", err)
+		t.Fatal("servicenode RPC failed", "err", err)
 	}
 	log.Debug("balance2", "balance-2", balance)
 	if balance.Cmp(fakeBalance2) != 0 {
 		t.Fatal(fmt.Sprintf("Expected balance %s to be equal to fake balance %s, but it is not", balance.String(), fakeBalance2.String()))
 	}
 
-	err = rpcclient_one.Call(&balance, "swap_balance")
+	err = rpcclient.Call(&balance, "swap_balance")
 	if err != nil {
-		t.Fatal("servicenode #1 RPC failed", "err", err)
+		t.Fatal("servicenode RPC failed", "err", err)
 	}
 	log.Debug("balance", "balance", balance)
 
@@ -356,10 +562,24 @@ func TestSwapRPC(t *testing.T) {
 	}
 }
 
+type dummyPeer struct {
+	*Peer
+	testFunc func(error)
+}
+
+func (dp *dummyPeer) Drop(err error) {
+	fmt.Println("DD")
+	dp.testFunc(err)
+}
+
 //creates a dummy protocols.Peer with dummy MsgReadWriter
-func newDummyPeer() *protocols.Peer {
+func newDummyPeer() *dummyPeer {
 	id := adapters.RandomNodeConfig().ID
-	return protocols.NewPeer(p2p.NewPeer(id, "testPeer", nil), &dummyRW{}, testSpec)
+	protoPeer := protocols.NewPeer(p2p.NewPeer(id, "testPeer", nil), &dummyRW{}, testSpec)
+	dummy := &dummyPeer{
+		Peer: NewPeer(protoPeer),
+	}
+	return dummy
 }
 
 //creates a p2p.Service node stub

swarm/swap_test.go

@@ -35,6 +35,7 @@ func TestSwapNetworkSymmetricFileUpload(t *testing.T) {
 	sim := simulation.New(map[string]simulation.ServiceFunc{
 		"swarm": func(ctx *adapters.ServiceContext, bucket *sync.Map) (s node.Service, cleanup func(), err error) {
 			config := api.NewConfig()
+			config.Port = strconv.Itoa(8500 + rand.Intn(9999))
 
 			dir, err := ioutil.TempDir("", "swap-network-test-node")
 			if err != nil {
@@ -63,6 +64,7 @@ func TestSwapNetworkSymmetricFileUpload(t *testing.T) {
 			if err != nil {
 				return nil, cleanup, err
 			}
+
 			bucket.Store(bucketKeySwarm, swarm)
 			log.Info("new swarm", "bzzKey", config.BzzKey, "baseAddr", fmt.Sprintf("%x", swarm.bzz.BaseAddr()))
 			return swarm, cleanup, nil

swarm/swarm.go

@@ -215,6 +215,10 @@ func NewSwarm(config *api.Config, mockStore *mock.NodeStore) (self *Swarm, err e
 
 	self.bzz = network.NewBzz(bzzconfig, to, stateStore, stream.Spec, self.streamer.Run)
 
+	if config.SwapEnabled {
+		self.swap = swap.New(stateStore)
+	}
+
 	// Pss = postal service over swarm (devp2p over bzz)
 	self.ps, err = pss.NewPss(to, config.Pss)
 	if err != nil {
@@ -370,6 +374,10 @@ func (self *Swarm) Start(srv *p2p.Server) error {
 	}
 	log.Info("Swarm network started", "bzzaddr", fmt.Sprintf("%x", self.bzz.Hive.BaseAddr()))
 
+	if self.swap != nil {
+		self.swap.Start(srv)
+	}
+
 	if self.ps != nil {
 		self.ps.Start(srv)
 	}
@@ -448,6 +456,10 @@ func (self *Swarm) Protocols() (protos []p2p.Protocol) {
 	if self.ps != nil {
 		protos = append(protos, self.ps.Protocols()...)
 	}
+
+	if self.swap != nil {
+		protos = append(protos, self.swap.Protocols()...)
+	}
 	return
 }