swarm/swap: locks and bugfixing

p2p/protocols/accounting.go

@@ -16,6 +16,10 @@
 
 package protocols
 
+import (
+	"sync"
+)
+
 type PriceOracle interface {
 	Price(uint32, interface{}) (EntryDirection, uint64)
 	Accountable(interface{}) bool
@@ -38,6 +42,7 @@ const (
 type AccountingHook struct {
 	BalanceManager
 	PriceOracle
+	lock sync.RWMutex //lock the balances
 }
 
 func NewAccountingHook(mgr BalanceManager, po PriceOracle) *AccountingHook {
@@ -49,6 +54,8 @@ func NewAccountingHook(mgr BalanceManager, po PriceOracle) *AccountingHook {
 }
 
 func (ah *AccountingHook) Send(peer *Peer, size uint32, msg interface{}) error {
+	ah.lock.Lock()
+	defer ah.lock.Unlock()
 	var err error
 	if !ah.PriceOracle.Accountable(msg) {
 		return nil
@@ -63,6 +70,8 @@ func (ah *AccountingHook) Send(peer *Peer, size uint32, msg interface{}) error {
 }
 
 func (ah *AccountingHook) Receive(peer *Peer, size uint32, msg interface{}) error {
+	ah.lock.Lock()
+	defer ah.lock.Unlock()
 	var err error
 	if !ah.PriceOracle.Accountable(msg) {
 		return nil

swarm/network/stream/stream.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"math"
+	"reflect"
 	"sync"
 	"time"
 
@@ -193,8 +194,10 @@ func NewRegistry(localID enode.ID, delivery *Delivery, syncChunkStore storage.Sy
 func (r *Registry) setupSpec() {
 	r.createSpec()
 	r.createPriceOracle()
+	if !reflect.ValueOf(r.balanceMgr).IsNil() {
 		r.spec.Hook = protocols.NewAccountingHook(r.balanceMgr, r.priceOracle)
 	}
+}
 
 // RegisterClient registers an incoming streamer constructor
 func (r *Registry) RegisterClientFunc(stream string, f func(*Peer, string, bool) (Client, error)) {

swarm/swap_test.go

@@ -319,7 +319,7 @@ func TestSwapNetworkAsymmetricFileUpload(t *testing.T) {
 		// or until the timeout is reached.
 		for {
 			if retrieve(sim, files, &checkStatusM, &nodeStatusM, &totalFoundCount) == 0 {
-				return nil
+				break
 			}
 		}