Moved wireguard NAT configuration to service manager.

mysteriumnetwork · Dec 6, 2018 · b73e31c · b73e31c
1 parent 51dd516
commit b73e31c
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 28 deletions.
diff --git a/services/wireguard/endpoint/endpoint.go b/services/wireguard/endpoint/endpoint.go
@@ -48,6 +48,13 @@ type connectionEndpoint struct {
 func (ce *connectionEndpoint) Start(config *wg.ServiceConfig) error {
 	ce.iface = ce.resourceAllocator.AllocateInterface()
 	ce.endpoint.Port = ce.resourceAllocator.AllocatePort()
+	if ce.ipResolver != nil {
+		publicIP, err := ce.ipResolver.GetPublicIP()
+		if err != nil {
+			return err
+		}
+		ce.endpoint.IP = net.ParseIP(publicIP)
+	}
 
 	if config == nil {
 		privateKey, err := GeneratePrivateKey()
@@ -63,29 +70,6 @@ func (ce *connectionEndpoint) Start(config *wg.ServiceConfig) error {
 		ce.privateKey = config.Consumer.PrivateKey
 	}
 
-	if ce.ipResolver != nil {
-		publicIP, err := ce.ipResolver.GetPublicIP()
-		if err != nil {
-			return err
-		}
-		ce.endpoint.IP = net.ParseIP(publicIP)
-
-		outboundIP, err := ce.ipResolver.GetOutboundIP()
-		if err != nil {
-			return err
-		}
-
-		ce.natService = nat.NewService()
-		ce.natService.Add(nat.RuleForwarding{
-			SourceAddress: ce.subnet.String(),
-			TargetIP:      outboundIP,
-		})
-
-		if err := ce.natService.Start(); err != nil {
-			return err
-		}
-	}
-
 	var deviceConfig deviceConfig
 	deviceConfig.listenPort = ce.endpoint.Port
 	deviceConfig.privateKey = ce.privateKey
@@ -117,10 +101,6 @@ func (ce *connectionEndpoint) Stop() error {
 		return err
 	}
 
-	if ce.natService != nil {
-		ce.natService.Stop()
-	}
-
 	if err := ce.resourceAllocator.ReleasePort(ce.endpoint.Port); err != nil {
 		return err
 	}

diff --git a/services/wireguard/service/service.go b/services/wireguard/service/service.go
@@ -25,6 +25,7 @@ import (
 	"github.com/mysteriumnetwork/node/core/location"
 	"github.com/mysteriumnetwork/node/identity"
 	"github.com/mysteriumnetwork/node/money"
+	"github.com/mysteriumnetwork/node/nat"
 	dto_discovery "github.com/mysteriumnetwork/node/service_discovery/dto"
 	wg "github.com/mysteriumnetwork/node/services/wireguard"
 	"github.com/mysteriumnetwork/node/services/wireguard/endpoint"
@@ -39,6 +40,7 @@ func NewManager(locationResolver location.Resolver, ipResolver ip.Resolver, conn
 		locationResolver:   locationResolver,
 		ipResolver:         ipResolver,
 		connectionEndpoint: connectionEndpoint,
+		natService:         nat.NewService(),
 	}
 }
 
@@ -48,13 +50,30 @@ type Manager struct {
 	ipResolver         ip.Resolver
 	connectionEndpoint wg.ConnectionEndpoint
 	wg                 sync.WaitGroup
+	natService         nat.NATService
 }
 
 // Start starts service - does not block
 func (manager *Manager) Start(providerID identity.Identity) (dto_discovery.ServiceProposal, session.ConfigProvider, error) {
 	if err := manager.connectionEndpoint.Start(nil); err != nil {
 		return dto_discovery.ServiceProposal{}, nil, err
 	}
+	config, err := manager.connectionEndpoint.Config()
+	if err != nil {
+		return dto_discovery.ServiceProposal{}, nil, err
+	}
+
+	outboundIP, err := manager.ipResolver.GetOutboundIP()
+	if err != nil {
+		return dto_discovery.ServiceProposal{}, nil, err
+	}
+	manager.natService.Add(nat.RuleForwarding{
+		SourceAddress: config.Subnet.String(),
+		TargetIP:      outboundIP,
+	})
+	if err := manager.natService.Start(); err != nil {
+		return dto_discovery.ServiceProposal{}, nil, err
+	}
 
 	sessionConfigProvider := func() (session.ServiceConfiguration, error) {
 		privateKey, err := endpoint.GeneratePrivateKey()
@@ -75,7 +94,6 @@ func (manager *Manager) Start(providerID identity.Identity) (dto_discovery.Servi
 		if err != nil {
 			return wg.ServiceConfig{}, err
 		}
-
 		config.Consumer.PrivateKey = privateKey
 		return config, nil
 	}
@@ -116,6 +134,7 @@ func (manager *Manager) Wait() error {
 // Stop stops service.
 func (manager *Manager) Stop() error {
 	manager.wg.Done()
+	manager.natService.Stop()
 	if err := manager.connectionEndpoint.Stop(); err != nil {
 		return err
 	}