diff --git a/netbox/extras/models/search.py b/netbox/extras/models/search.py index e77122f0f..50b0965c9 100644 --- a/netbox/extras/models/search.py +++ b/netbox/extras/models/search.py @@ -2,6 +2,8 @@ from django.contrib.contenttypes.fields import GenericForeignKey from django.contrib.contenttypes.models import ContentType from django.db import models +from utilities.fields import RestrictedGenericForeignKey + __all__ = ( 'CachedValue', ) @@ -18,7 +20,7 @@ class CachedValue(models.Model): related_name='+' ) object_id = models.PositiveBigIntegerField() - object = GenericForeignKey( + object = RestrictedGenericForeignKey( ct_field='object_type', fk_field='object_id' ) diff --git a/netbox/netbox/search/backends.py b/netbox/netbox/search/backends.py index 2fd3e40b1..82c1fa452 100644 --- a/netbox/netbox/search/backends.py +++ b/netbox/netbox/search/backends.py @@ -11,6 +11,7 @@ from django.db.models.signals import post_delete, post_save from extras.models import CachedValue from extras.registry import registry from netbox.constants import SEARCH_MAX_RESULTS +from utilities.querysets import RestrictedPrefetch from . import FieldTypes, LookupTypes, SearchResult, get_registry # The cache for the initialized backend. @@ -151,6 +152,7 @@ class CachedValueSearchBackend(SearchBackend): def search(self, request, value, lookup=DEFAULT_LOOKUP_TYPE): # Define the search parameters + # TODO: Filter object types to only those which the use has permission to view params = { f'value__{lookup}': value } @@ -168,14 +170,24 @@ class CachedValueSearchBackend(SearchBackend): ) ) + # Construct a Prefetch to pre-fetch only those related objects for which the + # user has permission to view. + prefetch = RestrictedPrefetch('object', request.user, 'view') + # Wrap the base query to return only the lowest-weight result for each object # Hat-tip to https://blog.oyam.dev/django-filter-by-window-function/ for the solution sql, params = queryset.query.sql_with_params() - return CachedValue.objects.prefetch_related('object').raw( + results = CachedValue.objects.prefetch_related(prefetch).raw( f"SELECT * FROM ({sql}) t WHERE row_number = 1", params ) + # Omit any results pertaining to an object the user does not have permission to view + # TODO: We'll have to figure out how to handle pagination + return [ + r for r in results if r.object is not None + ] + @classmethod def cache(cls, instance, data): ct = ContentType.objects.get_for_model(instance) diff --git a/netbox/utilities/fields.py b/netbox/utilities/fields.py index a9b851def..87eb8f2fe 100644 --- a/netbox/utilities/fields.py +++ b/netbox/utilities/fields.py @@ -1,3 +1,6 @@ +from collections import defaultdict + +from django.contrib.contenttypes.fields import GenericForeignKey from django.core.validators import RegexValidator from django.db import models @@ -71,3 +74,62 @@ class NaturalOrderingField(models.CharField): [self.target_field], kwargs, ) + + +class RestrictedGenericForeignKey(GenericForeignKey): + + # Replicated from GenericForeignKey + def get_prefetch_queryset(self, instances, queryset=None): + # Compensate for the hack in RestrictedPrefetch + restrict_params = queryset if type(queryset) is dict else {} + + # For efficiency, group the instances by content type and then do one + # query per model + fk_dict = defaultdict(set) + # We need one instance for each group in order to get the right db: + instance_dict = {} + ct_attname = self.model._meta.get_field(self.ct_field).get_attname() + for instance in instances: + # We avoid looking for values if either ct_id or fkey value is None + ct_id = getattr(instance, ct_attname) + if ct_id is not None: + fk_val = getattr(instance, self.fk_field) + if fk_val is not None: + fk_dict[ct_id].add(fk_val) + instance_dict[ct_id] = instance + + ret_val = [] + for ct_id, fkeys in fk_dict.items(): + instance = instance_dict[ct_id] + ct = self.get_content_type(id=ct_id, using=instance._state.db) + if restrict_params: + # Override the default behavior to call restrict() on each model's queryset + qs = ct.model_class().objects.filter(pk__in=fkeys).restrict(**restrict_params) + ret_val.extend(qs) + else: + # Default behavior + ret_val.extend(ct.get_all_objects_for_this_type(pk__in=fkeys)) + + # For doing the join in Python, we have to match both the FK val and the + # content type, so we use a callable that returns a (fk, class) pair. + def gfk_key(obj): + ct_id = getattr(obj, ct_attname) + if ct_id is None: + return None + else: + model = self.get_content_type( + id=ct_id, using=obj._state.db + ).model_class() + return ( + model._meta.pk.get_prep_value(getattr(obj, self.fk_field)), + model, + ) + + return ( + ret_val, + lambda obj: (obj.pk, obj.__class__), + gfk_key, + True, + self.name, + False, + ) diff --git a/netbox/utilities/querysets.py b/netbox/utilities/querysets.py index 955a10d64..a7f65d46c 100644 --- a/netbox/utilities/querysets.py +++ b/netbox/utilities/querysets.py @@ -1,9 +1,36 @@ -from django.db.models import QuerySet +from django.db.models import Prefetch, QuerySet from users.constants import CONSTRAINT_TOKEN_USER from utilities.permissions import permission_is_exempt, qs_filter_from_constraints +class RestrictedPrefetch(Prefetch): + """ + Extend Django's Prefetch to accept a user and action to be passed to the + `restrict()` method of the related object's queryset. + """ + def __init__(self, lookup, user, action='view', queryset=None, to_attr=None): + self.restrict_user = user + self.restrict_action = action + + super().__init__(lookup, queryset=queryset, to_attr=to_attr) + + def get_current_queryset(self, level): + params = { + 'user': self.restrict_user, + 'action': self.restrict_action, + } + + qs = super().get_current_queryset(level) + if qs: + return qs.filter(**params) + + # Bit of a hack. If no queryset is defined, pass through the dict of restrict() + # kwargs to be called later. This is necessary e.g. for GenericForeignKey fields, + # which do not permit setting a queryset on a Prefetch object. + return params + + class RestrictedQuerySet(QuerySet): def restrict(self, user, action='view'):