Skip to content

Commit

Permalink
Add methods to manage NAT ports
Browse files Browse the repository at this point in the history
  • Loading branch information
SRIKKANTH committed Jan 27, 2025
1 parent 922622b commit 07858ea
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 9 deletions.
6 changes: 1 addition & 5 deletions lisa/sut_orchestrator/hyperv/platform_.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@ def _initialize(self, *args: Any, **kwargs: Any) -> None:
self._source_vhd: Optional[PurePath] = None
self._source_factory = Factory[Source](Source)
self._source_files: Optional[List[PurePath]] = None
self._external_forwarding_port = 50000

self.device_pool = HyperVDevicePool(
node=self._server,
Expand Down Expand Up @@ -315,14 +314,11 @@ def _deploy_environment(self, environment: Environment, log: Logger) -> None:
# If the switch type is internal, we need to add a NAT mapping to access the
# VM from the outside of HyperV host.
if default_switch.type == HypervSwitchType.INTERNAL:
hv.add_nat_mapping(
port = hv.add_nat_mapping(
nat_name=default_switch.name,
internal_ip=ip_addr,
external_port=self._external_forwarding_port,
)
ip_addr = node_context.host.public_address
port = self._external_forwarding_port
self._external_forwarding_port += 1
username = self.runbook.admin_username
password = self.runbook.admin_password
node.set_connection_info(
Expand Down
28 changes: 24 additions & 4 deletions lisa/tools/hyperv.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ class HyperV(Tool):
# 192.168.5.12
IP_REGEX = r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}"
_default_switch: Optional[VMSwitch] = None
_external_forwarding_port_start = 50000
_assigned_nat_ports: set[int] = set()

@property
def command(self) -> str:
Expand Down Expand Up @@ -293,9 +295,24 @@ def create_nat(self, name: str, ip_range: str) -> None:
force_run=True,
)

def add_nat_mapping(
self, nat_name: str, internal_ip: str, external_port: int
) -> None:
def _get_next_available_nat_port(self) -> int:
# Start checking from the start_port and find the first available one
port = self._external_forwarding_port_start
while port in self._assigned_nat_ports:
port += 1
self._assigned_nat_ports.add(port) # Assign this port
return port

def _release_nat_port(self, port: int) -> None:
# Release a port if needed
if port in self._assigned_nat_ports:
self._assigned_nat_ports.remove(port)
else:
print(f"Port {port} is not assigned.")

def add_nat_mapping(self, nat_name: str, internal_ip: str) -> int:
external_port = self._get_next_available_nat_port()

# delete existing NAT mapping
self.delete_nat_mapping(external_port)

Expand All @@ -306,21 +323,24 @@ def add_nat_mapping(
f"-InternalPort 22 -ExternalPort {external_port}",
force_run=True,
)
return external_port

def delete_nat_mapping(self, external_port: int) -> None:
# create a new NAT
# get the NAT mapping id for the port
mapping_id = self.node.tools[PowerShell].run_cmdlet(
f"Get-NetNatStaticMapping | "
f"Where-Object {{$_.ExternalPort -eq {external_port}}}"
f" | Select-Object -ExpandProperty StaticMappingID",
force_run=True,
)
if mapping_id:
# delete the NAT mapping if it exists
self.node.tools[PowerShell].run_cmdlet(
f"Remove-NetNatStaticMapping -StaticMappingID {mapping_id} "
"-Confirm:$false",
force_run=True,
)
self._release_nat_port(external_port)
else:
self._log.debug(f"Mapping for port {external_port} does not exist")

Expand Down

0 comments on commit 07858ea

Please sign in to comment.