Fix leaked server connections (#2596)

Fix leaked server connections

close unused connections in the client lib
close deprecated connection in the server lib
The Server Picker is reusable in the guard if we want in the future. So we can support the server address changes.

---------

Co-authored-by: Maycon Santos <mlsmaycon@gmail.com>

* Add logging

---------

Co-authored-by: Maycon Santos <mlsmaycon@gmail.com>

relay/server/peer.go

@@ -115,6 +115,7 @@ func (p *Peer) Write(b []byte) (int, error) {
 // connection.
 func (p *Peer) CloseGracefully(ctx context.Context) {
 	p.connMu.Lock()
+	defer p.connMu.Unlock()
 	err := p.writeWithTimeout(ctx, messages.MarshalCloseMsg())
 	if err != nil {
 		p.log.Errorf("failed to send close message to peer: %s", p.String())
@@ -124,8 +125,15 @@ func (p *Peer) CloseGracefully(ctx context.Context) {
 	if err != nil {
 		p.log.Errorf("failed to close connection to peer: %s", err)
 	}
+}
 
+func (p *Peer) Close() {
+	p.connMu.Lock()
 	defer p.connMu.Unlock()
+
+	if err := p.conn.Close(); err != nil {
+		p.log.Errorf("failed to close connection to peer: %s", err)
+	}
 }
 
 // String returns the peer ID
@@ -167,6 +175,7 @@ func (p *Peer) handleHealthcheckEvents(ctx context.Context, hc *healthcheck.Send
 			if err != nil {
 				p.log.Errorf("failed to close connection to peer: %s", err)
 			}
+			p.log.Info("peer connection closed due healthcheck timeout")
 			return
 		case <-ctx.Done():
 			return

relay/server/store.go

@@ -19,10 +19,14 @@ func NewStore() *Store {
 }
 
 // AddPeer adds a peer to the store
-// todo: consider to close peer conn if the peer already exists
 func (s *Store) AddPeer(peer *Peer) {
 	s.peersLock.Lock()
 	defer s.peersLock.Unlock()
+	odlPeer, ok := s.peers[peer.String()]
+	if ok {
+		odlPeer.Close()
+	}
+
 	s.peers[peer.String()] = peer
 }
 

relay/server/store_test.go

@@ -2,13 +2,57 @@ package server
 
 import (
 	"context"
+	"net"
 	"testing"
+	"time"
 
 	"go.opentelemetry.io/otel"
 
 	"github.com/netbirdio/netbird/relay/metrics"
 )
 
+type mockConn struct {
+}
+
+func (m mockConn) Read(b []byte) (n int, err error) {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m mockConn) Write(b []byte) (n int, err error) {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m mockConn) Close() error {
+	return nil
+}
+
+func (m mockConn) LocalAddr() net.Addr {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m mockConn) RemoteAddr() net.Addr {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m mockConn) SetDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m mockConn) SetReadDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m mockConn) SetWriteDeadline(t time.Time) error {
+	//TODO implement me
+	panic("implement me")
+}
+
 func TestStore_DeletePeer(t *testing.T) {
 	s := NewStore()
 
@@ -27,8 +71,9 @@ func TestStore_DeleteDeprecatedPeer(t *testing.T) {
 
 	m, _ := metrics.NewMetrics(context.Background(), otel.Meter(""))
 
-	p1 := NewPeer(m, []byte("peer_id"), nil, nil)
-	p2 := NewPeer(m, []byte("peer_id"), nil, nil)
+	conn := &mockConn{}
+	p1 := NewPeer(m, []byte("peer_id"), conn, nil)
+	p2 := NewPeer(m, []byte("peer_id"), conn, nil)
 
 	s.AddPeer(p1)
 	s.AddPeer(p2)