Refactoring: Move some logic to the AlephFirecrackerVM class.

vm_supervisor/ip.py

@@ -36,13 +36,6 @@ class TapInterface:
  ip_network: IPv4NetworkWithInterfaces
  used_by: Dict[int, IPv4Interface]
 
- def __init__(
- self, device_name: str, ip_network: IPv4NetworkWithInterfaces, network
- ):
- self.device_name: str = device_name
- self.ip_network: IPv4NetworkWithInterfaces = ip_network
- self.network = network
-
  @property
  def guest_ip(self) -> IPv4Interface:
  return self.ip_network[2]
@@ -51,8 +44,15 @@ def guest_ip(self) -> IPv4Interface:
  def host_ip(self) -> IPv4Interface:
  return self.ip_network[1]
 
- async def create(self):
- logger.debug("Create network interface")
+ def __init__(
+ self, device_name: str, ip_network: IPv4NetworkWithInterfaces
+ ):
+ self.device_name: str = device_name
+ self.ip_network: IPv4NetworkWithInterfaces = ip_network
+
+ def create(self):
+ """Create """
+ logger.debug(f"Creating network interface {self.device_name} on {self.ip_network}")
 
  run(["/usr/bin/ip", "tuntap", "add", self.device_name, "mode", "tap"])
  run(
@@ -68,10 +68,9 @@ async def create(self):
  run(["/usr/bin/ip", "link", "set", self.device_name, "up"])
  logger.debug(f"Network interface created: {self.device_name}")
 
- async def delete(self, vm_id: int) -> None:
+ async def delete(self) -> None:
  """Asks the firewall to teardown any rules for the VM with id provided.
  Then removes the interface from the host."""
- self.network.firewall.teardown_nftables_for_vm(vm_id)
  logger.debug(f"Removing interface {self.device_name}")
  await asyncio.sleep(0.1) # Avoids Device/Resource busy bug
  run(["ip", "tuntap", "del", self.device_name, "mode", "tap"])
@@ -473,15 +472,17 @@ def teardown_nftables_for_vm(self, vm_id: int) -> None:
 
 class Network:
  firewall: Firewall
- def __init__(self):
- self.firewall = Firewall()
+ address_pool: IPv4NetworkWithInterfaces
  ipv4_forward_state_before_setup = None
- address_pool: IPv4NetworkWithInterfaces = IPv4NetworkWithInterfaces("172.16.0.0/12")
  network_size = 24
  network_initialized = False
  external_interface = "eth0"
  vm_info: Dict = {}
 
+ def __init__(self):
+ self.firewall = Firewall()
+ self.address_pool = IPv4NetworkWithInterfaces("172.16.0.0/12")
+
  def get_network_for_tap(self, vm_id: int) -> IPv4NetworkWithInterfaces:
  subnets = list(self.address_pool.subnets(new_prefix=self.network_size))
  return subnets[vm_id]
@@ -526,7 +527,6 @@ async def create_tap(self, vm_id: int) -> TapInterface:
  Currently this will create a new tap interface for each vm, but the structure is made this way
  to facilitate future tap interface sharing.
  """
- interface = TapInterface(f"vmtap{vm_id}", self.get_network_for_tap(vm_id), self)
- await interface.create()
- self.firewall.setup_nftables_for_vm(vm_id, interface)
+ interface = TapInterface(f"vmtap{vm_id}", self.get_network_for_tap(vm_id))
+ interface.create()
  return interface

vm_supervisor/models.py

@@ -15,7 +15,7 @@
 from .utils import dumps_for_json, create_task_log_exceptions
 from .vm import AlephFirecrackerVM
 from .vm.firecracker_microvm import AlephFirecrackerResources
-from .ip import TapInterface
+from .ip import TapInterface, Firewall
 
 logger = logging.getLogger(__name__)
 
@@ -102,7 +102,10 @@ async def prepare(self):
  self.resources = resources
 
  async def create(
- self, vm_id: int, tap_interface: TapInterface
+ self,
+ vm_id: int,
+ tap_interface: Optional[TapInterface] = None,
+ firewall: Optional[Firewall] = None,
  ) -> AlephFirecrackerVM:
  if not self.resources:
  raise ValueError("Execution resources must be configured first")
@@ -114,6 +117,7 @@ async def create(
  enable_networking=self.program.environment.internet,
  hardware_resources=self.program.resources,
  tap_interface=tap_interface,
+ firewall=firewall,
  )
  try:
  await vm.setup()

vm_supervisor/pool.py

@@ -40,7 +40,9 @@ async def create_a_vm(
  vm_id = self.get_unique_vm_id()
 
  tap_interface = await self.network.create_tap(vm_id)
- await execution.create(vm_id=vm_id, tap_interface=tap_interface)
+ await execution.create(
+ vm_id=vm_id, tap_interface=tap_interface, firewall=self.network.firewall
+ )
  return execution
 
  def get_unique_vm_id(self) -> int:

vm_supervisor/vm/firecracker_microvm.py

@@ -32,7 +32,7 @@
 from guest_api.__main__ import run_guest_api
 from ..conf import settings
 from ..storage import get_code_path, get_runtime_path, get_data_path, get_volume_path
-from ..ip import TapInterface
+from ..ip import TapInterface, Firewall
 
 logger = logging.getLogger(__name__)
 set_start_method("spawn")
@@ -213,6 +213,7 @@ class AlephFirecrackerVM:
  fvm: Optional[MicroVM] = None
  guest_api_process: Optional[Process] = None
  tap_interface: Optional[TapInterface] = None
+ firewall: Optional[Firewall] = None
 
  def __init__(
  self,
@@ -223,6 +224,7 @@ def __init__(
  enable_console: Optional[bool] = None,
  hardware_resources: MachineResources = MachineResources(),
  tap_interface: Optional[TapInterface] = None,
+ firewall: Optional[Firewall] = None,
  ):
  self.vm_id = vm_id
  self.vm_hash = vm_hash
@@ -233,6 +235,7 @@ def __init__(
  self.enable_console = enable_console
  self.hardware_resources = hardware_resources
  self.tap_interface = tap_interface
+ self.firewall = firewall
 
  def to_dict(self):
  if self.fvm.proc and psutil:
@@ -314,13 +317,15 @@ async def setup(self):
 
  logger.debug(config.json(by_alias=True, exclude_none=True, indent=4))
 
+ self.configure_firewall()
+
  try:
  await fvm.start(config)
  logger.debug("setup done")
  self.fvm = fvm
  except Exception:
  await fvm.teardown()
- await self.tap_interface.delete(self.vm_id)
+ await self.teardown_networking()
  raise
 
  async def start(self):
@@ -433,7 +438,7 @@ async def stop_guest_api(self):
  async def teardown(self):
  if self.fvm:
  await self.fvm.teardown()
- await self.tap_interface.delete(self.vm_id)
+ await self.teardown_networking()
  await self.stop_guest_api()
 
  async def run_code(
@@ -474,3 +479,24 @@ async def communicate(reader, writer, scope):
  logger.debug("Cleaning VM socket resources")
  writer.close()
  await writer.wait_closed()
+
+ async def teardown_networking(self):
+ """Remove VM specific networking configuration, including firewall rules and tap interface."""
+ if self.firewall:
+ self.cleanup_firewall()
+ if self.tap_interface:
+ await self.tap_interface.delete()
+
+ def configure_firewall(self):
+ """Configure the firewall to forward packets to the VM."""
+ if self.firewall:
+ self.firewall.setup_nftables_for_vm(
+ vm_id=self.vm_id, interface=self.tap_interface
+ )
+
+ def cleanup_firewall(self):
+ """Remove firewall rules specific to the VM."""
+ if self.firewall:
+ self.firewall.teardown_nftables_for_vm(
+ vm_id=self.vm_id
+ )