Adding new InstanceConfig value for VMSS & VM tags (#1560)

* Adding new instanceconfig value for tags.

* Removing bad import.

* Updating where tags are generated.

* Updating tag generation for scalesets.

* Updating tag generation in vm.

* Updating vm tag generation.

* Updating vm tag generation.

* Fixing extension.

* Fixing import.

* Fixing typing.

* Fixing get_vm calls.

* Fixing calls to get_vm.

* Fixing optional tag.
This commit is contained in:
Noah McGregor Harper 2022-01-05 13:16:03 -08:00 committed by GitHub
parent c5421894c1
commit 83e48e7e7b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 76 additions and 16 deletions

View File

@ -813,6 +813,20 @@ Each event will be submitted via HTTP POST to the user provided URL.
"default": "Standard_B2s",
"title": "Proxy Vm Sku",
"type": "string"
},
"vm_tags": {
"additionalProperties": {
"type": "string"
},
"title": "Vm Tags",
"type": "object"
},
"vmss_tags": {
"additionalProperties": {
"type": "string"
},
"title": "Vmss Tags",
"type": "object"
}
},
"required": [
@ -5955,6 +5969,20 @@ Each event will be submitted via HTTP POST to the user provided URL.
"default": "Standard_B2s",
"title": "Proxy Vm Sku",
"type": "string"
},
"vm_tags": {
"additionalProperties": {
"type": "string"
},
"title": "Vm Tags",
"type": "object"
},
"vmss_tags": {
"additionalProperties": {
"type": "string"
},
"title": "Vmss Tags",
"type": "object"
}
},
"required": [

View File

@ -49,7 +49,9 @@ def create_vm(
password: str,
ssh_public_key: str,
nsg: Optional[NSG],
tags: Optional[Dict[str, str]],
) -> Union[None, Error]:
resource_group = get_base_resource_group()
logging.info("creating vm %s:%s:%s", resource_group, location, name)
@ -114,8 +116,12 @@ def create_vm(
},
}
if "ONEFUZZ_OWNER" in os.environ:
params["tags"] = {"OWNER": os.environ["ONEFUZZ_OWNER"]}
if tags:
params["tags"] = tags.copy()
owner = os.environ.get("ONEFUZZ_OWNER")
if owner:
params["tags"]["OWNER"] = owner
try:
compute_client.virtual_machines.begin_create_or_update(
@ -230,6 +236,7 @@ class VM(BaseModel):
image: str
auth: Authentication
nsg: Optional[NSG]
tags: Optional[Dict[str, str]]
@validator("name", allow_reuse=True)
def check_name(cls, value: Union[UUID, str]) -> Union[UUID, str]:
@ -264,6 +271,7 @@ class VM(BaseModel):
self.auth.password,
self.auth.public_key,
self.nsg,
self.tags,
)
def delete(self) -> bool:

View File

@ -356,6 +356,7 @@ def create_vmss(
)
params["tags"] = tags.copy()
owner = os.environ.get("ONEFUZZ_OWNER")
if owner:
params["tags"]["OWNER"] = owner

View File

@ -148,8 +148,8 @@ def azmon_extension(
"publisher": "Microsoft.Azure.Monitor",
"location": region,
"type": "AzureMonitorLinuxAgent",
"typeHandlerVersion": "1.9",
"autoUpgradeMinorVersion": False,
"typeHandlerVersion": "1.0",
"autoUpgradeMinorVersion": True,
"settings": {"GCS_AUTO_CONFIG": True},
"protectedsettings": {
"configVersion": config_version,

View File

@ -69,19 +69,24 @@ class Proxy(ORMMixin):
def key_fields(cls) -> Tuple[str, Optional[str]]:
return ("region", "proxy_id")
def get_vm(self) -> VM:
sku = InstanceConfig.fetch().proxy_vm_sku
def get_vm(self, config: InstanceConfig) -> VM:
sku = config.proxy_vm_sku
tags = None
if config.vm_tags:
tags = config.vm_tags
vm = VM(
name="proxy-%s" % base58.b58encode(self.proxy_id.bytes).decode(),
region=self.region,
sku=sku,
image=PROXY_IMAGE,
auth=self.auth,
tags=tags,
)
return vm
def init(self) -> None:
vm = self.get_vm()
config = InstanceConfig.fetch()
vm = self.get_vm(config)
vm_data = vm.get()
if vm_data:
if vm_data.provisioning_state == "Failed":
@ -101,7 +106,6 @@ class Proxy(ORMMixin):
self.set_failed(result)
return
config = InstanceConfig.fetch()
nsg_config = config.proxy_nsg_config
result = nsg.set_allowed_sources(nsg_config)
if isinstance(result, Error):
@ -145,7 +149,8 @@ class Proxy(ORMMixin):
self.set_state(VmState.stopping)
def extensions_launch(self) -> None:
vm = self.get_vm()
config = InstanceConfig.fetch()
vm = self.get_vm(config)
vm_data = vm.get()
if not vm_data:
self.set_failed(
@ -177,7 +182,8 @@ class Proxy(ORMMixin):
self.save()
def stopping(self) -> None:
vm = self.get_vm()
config = InstanceConfig.fetch()
vm = self.get_vm(config)
if not vm.is_deleted():
logging.info(PROXY_LOG_PREFIX + "stopping proxy: %s", self.region)
vm.delete()

View File

@ -47,7 +47,11 @@ class Repro(BASE_REPRO, ORMMixin):
self.state = VmState.stopping
self.save()
def get_vm(self) -> VM:
def get_vm(self, config: InstanceConfig) -> VM:
tags = None
if config.vm_tags:
tags = config.vm_tags
task = Task.get_by_task_id(self.task_id)
if isinstance(task, Error):
raise Exception("previously existing task missing: %s" % self.task_id)
@ -71,10 +75,12 @@ class Repro(BASE_REPRO, ORMMixin):
sku=vm_config.sku,
image=vm_config.image,
auth=self.auth,
tags=tags,
)
def init(self) -> None:
vm = self.get_vm()
config = InstanceConfig.fetch()
vm = self.get_vm(config)
vm_data = vm.get()
if vm_data:
if vm_data.provisioning_state == "Failed":
@ -96,7 +102,6 @@ class Repro(BASE_REPRO, ORMMixin):
self.set_failed(result)
return
config = InstanceConfig.fetch()
nsg_config = config.proxy_nsg_config
result = nsg.set_allowed_sources(nsg_config)
if isinstance(result, Error):
@ -128,7 +133,8 @@ class Repro(BASE_REPRO, ORMMixin):
return None
def extensions_launch(self) -> None:
vm = self.get_vm()
config = InstanceConfig.fetch()
vm = self.get_vm(config)
vm_data = vm.get()
if not vm_data:
self.set_error(
@ -159,7 +165,8 @@ class Repro(BASE_REPRO, ORMMixin):
self.save()
def stopping(self) -> None:
vm = self.get_vm()
config = InstanceConfig.fetch()
vm = self.get_vm(config)
if not vm.is_deleted():
logging.info("vm stopping: %s", self.vm_id)
vm.delete()

View File

@ -221,7 +221,9 @@ class Scaleset(BASE_SCALESET, ORMMixin):
SCALESET_LOG_PREFIX + "creating scaleset. scaleset_id:%s",
self.scaleset_id,
)
extensions = fuzz_extensions(pool, self)
result = create_vmss(
self.region,
self.scaleset_id,

View File

@ -16,6 +16,7 @@ from onefuzztypes.responses import BoolResult
from ..onefuzzlib.azure.creds import get_base_region, get_regions
from ..onefuzzlib.azure.vmss import list_available_skus
from ..onefuzzlib.config import InstanceConfig
from ..onefuzzlib.endpoint_authorization import call_if_user, check_can_manage_pools
from ..onefuzzlib.request import not_ok, ok, parse_request
from ..onefuzzlib.workers.pools import Pool
@ -44,6 +45,7 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
def post(req: func.HttpRequest) -> func.HttpResponse:
request = parse_request(ScalesetCreate, req)
instance_config = InstanceConfig.fetch()
if isinstance(request, Error):
return not_ok(request, context="ScalesetCreate")
@ -88,6 +90,10 @@ def post(req: func.HttpRequest) -> func.HttpResponse:
context="scalesetcreate",
)
tags = request.tags
if instance_config.vmss_tags:
tags.update(instance_config.vmss_tags)
scaleset = Scaleset.create(
pool_name=request.pool_name,
vm_sku=request.vm_sku,
@ -96,7 +102,7 @@ def post(req: func.HttpRequest) -> func.HttpResponse:
size=request.size,
spot_instances=request.spot_instances,
ephemeral_os_disks=request.ephemeral_os_disks,
tags=request.tags,
tags=tags,
)
# don't return auths during create, only 'get' with include_auth
scaleset.auth = None

View File

@ -864,6 +864,8 @@ class InstanceConfig(BaseModel):
proxy_vm_sku: str = Field(default="Standard_B2s")
api_access_rules: Optional[Dict[Endpoint, ApiAccessRule]] = None
group_membership: Optional[Dict[PrincipalID, List[GroupId]]] = None
vm_tags: Optional[Dict[str, str]] = None
vmss_tags: Optional[Dict[str, str]] = None
def update(self, config: "InstanceConfig") -> None:
for field in config.__fields__: