ack to return overlay

go.mod

@@ -3,6 +3,7 @@ module github.com/ethersphere/bee
 go 1.13
 
 require (
+	github.com/btcsuite/btcd v0.20.1-beta
 	github.com/gogo/protobuf v1.3.1
 	github.com/gorilla/handlers v1.4.2
 	github.com/gorilla/mux v1.7.3

go.sum

@@ -20,6 +20,7 @@ github.com/btcsuite/btcd v0.0.0-20190523000118-16327141da8c/go.mod h1:3J08xEfcug
 github.com/btcsuite/btcd v0.0.0-20190824003749-130ea5bddde3/go.mod h1:3J08xEfcugPacsc34/LKRU2yO7YmuT8yt28J8k2+rrI=
 github.com/btcsuite/btcd v0.20.1-beta h1:Ik4hyJqN8Jfyv3S4AGBOmyouMsYE3EdYODkMbQjwPGw=
 github.com/btcsuite/btcd v0.20.1-beta/go.mod h1:wVuoA8VJLEcwgqHBwHmzLRazpKxTv13Px/pDuV7OomQ=
+github.com/btcsuite/btclog v0.0.0-20170628155309-84c8d2346e9f h1:bAs4lUbRJpnnkd9VhRV3jjAVU7DJVjMaK+IsvSeZvFo=
 github.com/btcsuite/btclog v0.0.0-20170628155309-84c8d2346e9f/go.mod h1:TdznJufoqS23FtqVCzL0ZqgP5MqXbb4fg/WgDys70nA=
 github.com/btcsuite/btcutil v0.0.0-20190207003914-4c204d697803/go.mod h1:+5NJ2+qvTyV9exUAL/rxXi3DcLg2Ts+ymUAY5y4NvMg=
 github.com/btcsuite/btcutil v0.0.0-20190425235716-9e5f4b9a998d/go.mod h1:+5NJ2+qvTyV9exUAL/rxXi3DcLg2Ts+ymUAY5y4NvMg=

pkg/p2p/libp2p/internal/handshake/handshake.go

@@ -6,7 +6,6 @@ package handshake
 
 import (
 	"fmt"
-
 	"github.com/ethersphere/bee/pkg/logging"
 	"github.com/ethersphere/bee/pkg/p2p"
 	"github.com/ethersphere/bee/pkg/p2p/libp2p/internal/handshake/pb"
@@ -35,7 +34,7 @@ func New(overlay string, networkID int32, logger logging.Logger) *Service {
 
 func (s *Service) Handshake(stream p2p.Stream) (i *Info, err error) {
 	w, r := protobuf.NewWriterAndReader(stream)
-	var resp pb.ShakeHand
+	var resp pb.ShakeHandAck
 	if err := w.WriteMsg(&pb.ShakeHand{
 		Address:   s.overlay,
 		NetworkID: s.networkID,
@@ -48,15 +47,15 @@ func (s *Service) Handshake(stream p2p.Stream) (i *Info, err error) {
 		return nil, fmt.Errorf("handshake read message: %w", err)
 	}
 
-	if err := w.WriteMsg(&pb.Ack{}); err != nil {
-		return nil, fmt.Errorf("handshake write message: %w", err)
+	if err := w.WriteMsg(&pb.Ack{Address: resp.ShakeHand.Address}); err != nil {
+		return nil, fmt.Errorf("ack write message: %w", err)
 	}
 
-	s.logger.Tracef("handshake read response: %s", resp.Address)
+	s.logger.Tracef("handshake read response: %s", resp.ShakeHand.Address)
 	return &Info{
-		Address:   resp.Address,
-		NetworkID: resp.NetworkID,
-		Light:     resp.Light,
+		Address:   resp.ShakeHand.Address,
+		NetworkID: resp.ShakeHand.NetworkID,
+		Light:     resp.ShakeHand.Light,
 	}, nil
 }
 
@@ -70,16 +69,19 @@ func (s *Service) Handle(stream p2p.Stream) (i *Info, err error) {
 	}
 
 	s.logger.Tracef("handshake handler received request %s", req.Address)
-	if err := w.WriteMsg(&pb.ShakeHand{
+	if err := w.WriteMsg(&pb.ShakeHandAck{
+		ShakeHand: &pb.ShakeHand{
 			Address:   s.overlay,
 			NetworkID: s.networkID,
+		},
+		Ack: &pb.Ack{Address: req.Address},
 	}); err != nil {
 		return nil, fmt.Errorf("handshake handler write message: %w", err)
 	}
 
 	var ack pb.Ack
 	if err := r.ReadMsg(&ack); err != nil {
-		return nil, fmt.Errorf("handshake handler read message: %w", err)
+		return nil, fmt.Errorf("ack read message: %w", err)
 	}
 
 	s.logger.Tracef("handshake handled response: %s", s.overlay)

pkg/p2p/libp2p/internal/handshake/handshake_test.go

@@ -18,23 +18,41 @@ import (
 type StreamMock struct {
 	readBuffer        *bytes.Buffer
 	writeBuffer       *bytes.Buffer
+	writeCounter      int
+	readCounter       int
 	readError         error
 	writeError        error
+	readErrCheckmark  int
+	writeErrCheckmark int
+}
+
+func (s *StreamMock) setReadErr(err error, checkmark int) {
+	s.readError = err
+	s.readErrCheckmark = checkmark
+
+}
+
+func (s *StreamMock) setWriteErr(err error, checkmark int) {
+	s.writeError = err
+	s.writeErrCheckmark = checkmark
+
 }
 
 func (s *StreamMock) Read(p []byte) (n int, err error) {
-	if s.readError != nil {
+	if s.readError != nil && s.readErrCheckmark <= s.readCounter {
 		return 0, s.readError
 	}
 
+	s.readCounter++
 	return s.readBuffer.Read(p)
 }
 
 func (s *StreamMock) Write(p []byte) (n int, err error) {
-	if s.writeError != nil {
+	if s.writeError != nil && s.writeErrCheckmark <= s.writeCounter {
 		return 0, s.writeError
 	}
 
+	s.writeCounter++
 	return s.writeBuffer.Write(p)
 }
 
@@ -64,10 +82,13 @@ func TestHandshake(t *testing.T) {
 		stream2 := &StreamMock{readBuffer: &buffer2, writeBuffer: &buffer1}
 
 		w, r := protobuf.NewWriterAndReader(stream2)
-		if err := w.WriteMsg(&pb.ShakeHand{
+		if err := w.WriteMsg(&pb.ShakeHandAck{
+			ShakeHand: &pb.ShakeHand{
 				Address:   expectedInfo.Address,
 				NetworkID: expectedInfo.NetworkID,
 				Light:     expectedInfo.Light,
+			},
+			Ack: &pb.Ack{},
 		}); err != nil {
 			t.Fatal(err)
 		}
@@ -86,10 +107,11 @@ func TestHandshake(t *testing.T) {
 		}
 	})
 
-	t.Run("ERROR - write error ", func(t *testing.T) {
+	t.Run("ERROR - shakehand write error ", func(t *testing.T) {
 		testErr := errors.New("test error")
 		expectedErr := fmt.Errorf("handshake write message: %w", testErr)
-		stream := &StreamMock{writeError: testErr}
+		stream := &StreamMock{}
+		stream.setWriteErr(testErr, 0)
 		res, err := handshakeService.Handshake(stream)
 		if err == nil || err.Error() != expectedErr.Error() {
 			t.Fatal("expected:", expectedErr, "got:", err)
@@ -100,10 +122,11 @@ func TestHandshake(t *testing.T) {
 		}
 	})
 
-	t.Run("ERROR - read error ", func(t *testing.T) {
+	t.Run("ERROR - shakehand read error ", func(t *testing.T) {
 		testErr := errors.New("test error")
 		expectedErr := fmt.Errorf("handshake read message: %w", testErr)
-		stream := &StreamMock{writeBuffer: &bytes.Buffer{}, readError: testErr}
+		stream := &StreamMock{writeBuffer: &bytes.Buffer{}}
+		stream.setReadErr(testErr, 0)
 		res, err := handshakeService.Handshake(stream)
 		if err == nil || err.Error() != expectedErr.Error() {
 			t.Fatal("expected:", expectedErr, "got:", err)
@@ -113,6 +136,44 @@ func TestHandshake(t *testing.T) {
 			t.Fatal("handshake returned non-nil res")
 		}
 	})
+
+	t.Run("ERROR - ack write error ", func(t *testing.T) {
+		testErr := errors.New("test error")
+		expectedErr := fmt.Errorf("ack write message: %w", testErr)
+
+		expectedInfo := Info{
+			Address:   "node2",
+			NetworkID: 1,
+			Light:     false,
+		}
+
+		var buffer1 bytes.Buffer
+		var buffer2 bytes.Buffer
+		stream1 := &StreamMock{readBuffer: &buffer1, writeBuffer: &buffer2}
+		stream1.setWriteErr(testErr, 1)
+		stream2 := &StreamMock{readBuffer: &buffer2, writeBuffer: &buffer1}
+
+		w, _ := protobuf.NewWriterAndReader(stream2)
+		if err := w.WriteMsg(&pb.ShakeHandAck{
+			ShakeHand: &pb.ShakeHand{
+				Address:   expectedInfo.Address,
+				NetworkID: expectedInfo.NetworkID,
+				Light:     expectedInfo.Light,
+			},
+			Ack: &pb.Ack{Address: info.Address},
+		}); err != nil {
+			t.Fatal(err)
+		}
+
+		res, err := handshakeService.Handshake(stream1)
+		if err == nil || err.Error() != expectedErr.Error() {
+			t.Fatal("expected:", expectedErr, "got:", err)
+		}
+
+		if res != nil {
+			t.Fatal("handshake returned non-nil res")
+		}
+	})
 }
 
 func TestHandle(t *testing.T) {
@@ -146,7 +207,7 @@ func TestHandle(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		if err := w.WriteMsg(&pb.Ack{}); err != nil {
+		if err := w.WriteMsg(&pb.Ack{Address: node2Info.Address}); err != nil {
 			t.Fatal(err)
 		}
 
@@ -160,12 +221,12 @@ func TestHandle(t *testing.T) {
 		}
 
 		_, r := protobuf.NewWriterAndReader(stream2)
-		var got pb.ShakeHand
+		var got pb.ShakeHandAck
 		if err := r.ReadMsg(&got); err != nil {
 			t.Fatal(err)
 		}
 
-		if nodeInfo != Info(got) {
+		if nodeInfo != Info(*got.ShakeHand) {
 			t.Fatalf("got %+v, expected %+v", got, node2Info)
 		}
 	})
@@ -173,7 +234,8 @@ func TestHandle(t *testing.T) {
 	t.Run("ERROR - read error ", func(t *testing.T) {
 		testErr := errors.New("test error")
 		expectedErr := fmt.Errorf("handshake handler read message: %w", testErr)
-		stream := &StreamMock{readError: testErr}
+		stream := &StreamMock{}
+		stream.setReadErr(testErr, 0)
 		res, err := handshakeService.Handle(stream)
 		if err == nil || err.Error() != expectedErr.Error() {
 			t.Fatal("expected:", expectedErr, "got:", err)
@@ -189,7 +251,7 @@ func TestHandle(t *testing.T) {
 		expectedErr := fmt.Errorf("handshake handler write message: %w", testErr)
 		var buffer bytes.Buffer
 		stream := &StreamMock{readBuffer: &buffer, writeBuffer: &buffer}
-
+		stream.setWriteErr(testErr, 1)
 		w, _ := protobuf.NewWriterAndReader(stream)
 		if err := w.WriteMsg(&pb.ShakeHand{
 			Address:   "node1",
@@ -199,7 +261,6 @@ func TestHandle(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		stream.writeError = testErr
 		res, err := handshakeService.Handle(stream)
 		if err == nil || err.Error() != expectedErr.Error() {
 			t.Fatal("expected:", expectedErr, "got:", err)
@@ -209,4 +270,38 @@ func TestHandle(t *testing.T) {
 			t.Fatal("handshake returned non-nil res")
 		}
 	})
+
+	t.Run("ERROR - ack read error ", func(t *testing.T) {
+		testErr := errors.New("test error")
+		expectedErr := fmt.Errorf("ack read message: %w", testErr)
+
+		node2Info := Info{
+			Address:   "node2",
+			NetworkID: 1,
+			Light:     false,
+		}
+
+		var buffer1 bytes.Buffer
+		var buffer2 bytes.Buffer
+		stream1 := &StreamMock{readBuffer: &buffer1, writeBuffer: &buffer2}
+		stream2 := &StreamMock{readBuffer: &buffer2, writeBuffer: &buffer1}
+		stream1.setReadErr(testErr, 1)
+		w, _ := protobuf.NewWriterAndReader(stream2)
+		if err := w.WriteMsg(&pb.ShakeHand{
+			Address:   node2Info.Address,
+			NetworkID: node2Info.NetworkID,
+			Light:     node2Info.Light,
+		}); err != nil {
+			t.Fatal(err)
+		}
+
+		res, err := handshakeService.Handle(stream1)
+		if err == nil || err.Error() != expectedErr.Error() {
+			t.Fatal("expected:", expectedErr, "got:", err)
+		}
+
+		if res != nil {
+			t.Fatal("handshake returned non-nil res")
+		}
+	})
 }

pkg/p2p/libp2p/internal/handshake/pb/handshake.proto

@@ -11,5 +11,11 @@ message ShakeHand {
     bool Light = 3;
 }
 
-message Ack {}
+message ShakeHandAck {
+    ShakeHand ShakeHand = 1;
+    Ack Ack = 2;
+}
 
+message Ack {
+    string Address = 1;
+}