fix: netbirdio#35 peer Registration Race when client connects to the …

…signal server
pulsastrix · Jun 17, 2021 · 9538693 · 9538693
1 parent f3547e2
commit 9538693
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 2 deletions.
diff --git a/go.sum b/go.sum
@@ -220,6 +220,8 @@ github.com/pion/udp v0.1.1 h1:8UAPvyqmsxK8oOjloDk4wUt63TzFe9WEJkg5lChlj7o=
 github.com/pion/udp v0.1.1/go.mod h1:6AFo+CMdKQm7UiA0eUPA8/eVCTx8jBIITLZHc9DWX5M=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
+github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/posener/complete v1.1.1/go.mod h1:em0nMJCgc9GFtwrmVmEMR/ZL6WyhyjMBndrE9hABlRI=
@@ -360,6 +362,7 @@ golang.org/x/sys v0.0.0-20181026203630-95b1ffbd15a5/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20181107165924-66b7b1311ac8/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20190228124157-a34e9553db1e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190411185658-b44545bcd369/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

diff --git a/signal/client.go b/signal/client.go
@@ -109,15 +109,22 @@ func (c *Client) connect(key string, msgHandler func(msg *proto.Message) error)
 	// add key fingerprint to the request header to be identified on the server side
 	md := metadata.New(map[string]string{proto.HeaderId: key})
 	ctx := metadata.NewOutgoingContext(c.ctx, md)
-	//ctx, cancel := context.WithCancel(ctx)
-	//defer cancel()
 
 	stream, err := c.realClient.ConnectStream(ctx)
 
 	c.stream = stream
 	if err != nil {
 		return err
 	}
+	// blocks
+	header, err := c.stream.Header()
+	if err != nil {
+		return err
+	}
+	registered := header.Get(proto.HeaderRegistered)
+	if len(registered) == 0 {
+		return fmt.Errorf("didn't receive a registration header from the Signal server whille connecting to the streams")
+	}
 	//connection established we are good to use the stream
 	c.connWg.Done()
 

diff --git a/signal/proto/constants.go b/signal/proto/constants.go
@@ -2,3 +2,4 @@ package proto
 
 // protocol constants, field names that can be used by both client and server
 const HeaderId = "x-wiretrustee-peer-id"
+const HeaderRegistered = "x-wiretrustee-peer-registered"
diff --git a/signal/signal.go b/signal/signal.go
@@ -52,6 +52,13 @@ func (s *Server) ConnectStream(stream proto.SignalExchange_ConnectStreamServer)
 		return err
 	}
 
+	//needed to confirm that the peer has been registered so that the client can proceed
+	header := metadata.Pairs(proto.HeaderRegistered, "1")
+	err = stream.SendHeader(header)
+	if err != nil {
+		return err
+	}
+
 	log.Infof("peer [%s] has successfully connected", p.Id)
 	for {
 		msg, err := stream.Recv()
Original file line number	Diff line number	Diff line change
Expand Up		@@ -2,3 +2,4 @@ package proto

		// protocol constants, field names that can be used by both client and server
		const HeaderId = "x-wiretrustee-peer-id"
		const HeaderRegistered = "x-wiretrustee-peer-registered"