github.com/mysteriumnetwork/node@v0.0.0-20240516044423-365054f76801/communication/nats/connector.go (about)

     1  /*
     2   * Copyright (C) 2020 The "MysteriumNetwork/node" Authors.
     3   *
     4   * This program is free software: you can redistribute it and/or modify
     5   * it under the terms of the GNU General Public License as published by
     6   * the Free Software Foundation, either version 3 of the License, or
     7   * (at your option) any later version.
     8   *
     9   * This program is distributed in the hope that it will be useful,
    10   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    11   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    12   * GNU General Public License for more details.
    13   *
    14   * You should have received a copy of the GNU General Public License
    15   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    16   */
    17  
    18  package nats
    19  
    20  import (
    21  	"context"
    22  	"net"
    23  	"net/url"
    24  
    25  	nats_lib "github.com/nats-io/nats.go"
    26  	"github.com/pkg/errors"
    27  	"github.com/rs/zerolog/log"
    28  
    29  	"github.com/mysteriumnetwork/node/firewall"
    30  	"github.com/mysteriumnetwork/node/requests"
    31  	"github.com/mysteriumnetwork/node/requests/resolver"
    32  )
    33  
    34  // BrokerConnector establishes new connections to NATS servers and handles reconnects.
    35  type BrokerConnector struct {
    36  	// resolveContext specifies the resolve function for doing custom DNS lookup.
    37  	// If ResolveContext is nil, then the transport dials using package net.
    38  	resolveContext resolver.ResolveContext
    39  
    40  	dialer requests.DialContext
    41  }
    42  
    43  // NewBrokerConnector creates a new BrokerConnector.
    44  func NewBrokerConnector(dialer requests.DialContext, resolveContext resolver.ResolveContext) *BrokerConnector {
    45  	return &BrokerConnector{
    46  		resolveContext: resolveContext,
    47  		dialer:         dialer,
    48  	}
    49  }
    50  
    51  func (b *BrokerConnector) resolveServers(serverURLs []*url.URL) ([]*url.URL, error) {
    52  	if b.resolveContext == nil {
    53  		return serverURLs, nil
    54  	}
    55  
    56  	for _, serverURL := range serverURLs {
    57  		ctx, cancel := context.WithTimeout(context.Background(), nats_lib.DefaultTimeout)
    58  		defer cancel()
    59  
    60  		addrs, err := b.resolveContext(ctx, "tcp", serverURL.Host)
    61  		if err != nil {
    62  			return nil, errors.Wrapf(err, `failed to resolve NATS server "%s"`, serverURL.Hostname())
    63  		}
    64  
    65  		cacheBrokerDNS(serverURL.Host, addrs)
    66  
    67  		for _, addr := range addrs {
    68  			serverURLResolved := *serverURL
    69  			serverURLResolved.Host = addr
    70  			serverURLs = append(serverURLs, &serverURLResolved)
    71  		}
    72  	}
    73  
    74  	return serverURLs, nil
    75  }
    76  
    77  // Connect establishes a new connection to the broker(s).
    78  func (b *BrokerConnector) Connect(serverURLs ...*url.URL) (Connection, error) {
    79  	log.Debug().Msgf("Connecting to NATS servers: %v", serverURLs)
    80  
    81  	serverURLs, err := b.resolveServers(serverURLs)
    82  	if err != nil {
    83  		return nil, err
    84  	}
    85  
    86  	servers := make([]string, len(serverURLs))
    87  	for i, serverURL := range serverURLs {
    88  		servers[i] = serverURL.String()
    89  	}
    90  
    91  	removeFirewallRule, err := firewall.AllowURLAccess(servers...)
    92  	if err != nil {
    93  		return nil, errors.Wrapf(err, `failed to allow NATS servers "%v" in firewall`, servers)
    94  	}
    95  
    96  	conn, err := newConnection(b.dialer, servers...)
    97  	if err != nil {
    98  		return nil, err
    99  	}
   100  
   101  	if err := conn.Open(); err != nil {
   102  		return nil, err
   103  	}
   104  
   105  	conn.onClose = removeFirewallRule
   106  
   107  	return conn, nil
   108  }
   109  
   110  func cacheBrokerDNS(server string, addrs []string) {
   111  	host, _, err := net.SplitHostPort(server)
   112  	if err != nil {
   113  		log.Warn().Msgf("Failed to parse broker address: %v", server)
   114  		return
   115  	}
   116  
   117  	cacheAddrs := []string{}
   118  
   119  	for _, addr := range addrs {
   120  		if server == addr {
   121  			continue
   122  		}
   123  
   124  		ip, _, err := net.SplitHostPort(addr)
   125  		if err != nil {
   126  			log.Warn().Msgf("Failed to parse broker address: %v", addr)
   127  			continue
   128  		}
   129  
   130  		cacheAddrs = append(cacheAddrs, ip)
   131  	}
   132  
   133  	resolver.CacheDNSRecord(host, cacheAddrs)
   134  }