Prevent AddPeerInfo to override existing server information

be1-go/hub/standard_hub/hub_state/Peers.go

@@ -1,6 +1,7 @@
 package hub_state
 
 import (
+	"golang.org/x/xerrors"
 	"popstellar/message/query/method"
 	"sync"
 
@@ -26,10 +27,19 @@ func NewPeers() Peers {
 }
 
 // AddPeerInfo adds a peer's info to the table
-func (p *Peers) AddPeerInfo(socketId string, info method.ServerInfo) {
+func (p *Peers) AddPeerInfo(socketId string, info method.ServerInfo) error {
 	p.Lock()
 	defer p.Unlock()
+
+	currentInfo, ok := p.peersInfo[socketId]
+	if ok {
+		return xerrors.Errorf(
+			"cannot add %s because peersInfo[%s] already contains %s",
+			info, socketId, currentInfo)
+	}
+
 	p.peersInfo[socketId] = info
+	return nil
 }
 
 // AddPeerGreeted adds a peer's socket ID to the slice of peers greeted

be1-go/hub/standard_hub/message_handling.go

@@ -451,7 +451,10 @@ func (h *Hub) handleGreetServer(socket socket.Socket, byteMessage []byte) error
 	}
 
 	// store information about the server
-	h.peers.AddPeerInfo(socket.ID(), greetServer.Params)
+	err = h.peers.AddPeerInfo(socket.ID(), greetServer.Params)
+	if err != nil {
+		return xerrors.Errorf("failed to add peer info: %v", err)
+	}
 
 	if h.peers.IsPeerGreeted(socket.ID()) {
 		return nil

be1-go/hub/standard_hub/mod_test.go

@@ -1877,6 +1877,76 @@ func Test_Handle_GreetServer_Already_Greeted(t *testing.T) {
 	require.Nil(t, sock.msg)
 }
 
+// Test that receiving multiple greet server messages from the same source will
+// not override the existing server information and that an error is raised
+func Test_Handle_GreetServer_Already_Received(t *testing.T) {
+	keypair := generateKeyPair(t)
+
+	hub, err := NewHub(keypair.public, "", "", nolog, nil)
+	require.NoError(t, err)
+
+	serverInfo1 := method.ServerInfo{
+		PublicKey:     "",
+		ServerAddress: "ws://localhost:9003/server",
+		ClientAddress: "ws://localhost:9002/client",
+	}
+
+	serverInfo2 := method.ServerInfo{
+		PublicKey:     "",
+		ServerAddress: "ws://localhost:9005/server",
+		ClientAddress: "ws://localhost:9004/client",
+	}
+
+	serverGreet1 := method.GreetServer{
+		Base: query.Base{
+			JSONRPCBase: jsonrpc.JSONRPCBase{
+				JSONRPC: "2.0",
+			},
+			Method: query.MethodGreetServer,
+		},
+		Params: serverInfo1,
+	}
+
+	serverGreet2 := method.GreetServer{
+		Base: query.Base{
+			JSONRPCBase: jsonrpc.JSONRPCBase{
+				JSONRPC: "2.0",
+			},
+			Method: query.MethodGreetServer,
+		},
+		Params: serverInfo2,
+	}
+
+	sock := &fakeSocket{id: "fakeID"}
+
+	msg1, err := json.Marshal(serverGreet1)
+	require.NoError(t, err)
+
+	msg2, err := json.Marshal(serverGreet2)
+	require.NoError(t, err)
+
+	err = hub.handleMessageFromServer(&socket.IncomingMessage{
+		Socket:  sock,
+		Message: msg1,
+	})
+	require.NoError(t, err)
+	require.NoError(t, sock.err)
+
+	// check that handling GreetServer from the same source twice throw an error
+	err = hub.handleMessageFromServer(&socket.IncomingMessage{
+		Socket:  sock,
+		Message: msg2,
+	})
+	require.Error(t, err)
+	require.Error(t, sock.err)
+
+	// check that the peersInfo were not modified by the second GreetServer
+	peersInfo := hub.GetPeersInfo()
+	require.Len(t, peersInfo, 1)
+	require.Equal(t, serverInfo1, peersInfo[0])
+	require.NotEqual(t, serverInfo2, peersInfo[0])
+}
+
 // -----------------------------------------------------------------------------
 // Utility functions