Confirm Location's protocol support on edit

core/controllers/profile_state/data_filters.py

@@ -27,16 +27,10 @@ Inputs:
 def get_supported_protocols_for_location(location_codes: str, profile_type: str = None) -> Result:
 
     country, city = location_codes.split("_")
-    print(f"country is {country}")
-    print(f"city is {city}")
 
     # this IS doing sql:
     location = get_profile_location_data(country_code=country, city_code=city)
 
-    print(f"location is {location}")
-    print(f"location type is {type(location)}")
-    # return location
-
     if not location:
         error_msg = f"This is a developer error.  No location data found for {location_codes}"
         logger.error(error_msg)
@@ -58,10 +52,7 @@ def _get_supported_protocols_for_location(location: Location, as_strings: bool =
 
     for each_protocol in ALL_PROTOCOLS:
         supported = location.supports_protocol(enum_value=each_protocol)
-        print(f"supported for {each_protocol} is {supported}")
-        print(f"type of supported is {type(supported)}")
         if supported:
-            print(f"THIS IS TRUE for {each_protocol}")
             SUPPORTED_PROTOCOLS.append(each_protocol)
 
     return SUPPORTED_PROTOCOLS

core/controllers/profile_state/update_profile.py

@@ -14,7 +14,7 @@ from core.models.BaseProfile import get_application_version
 # from core.models.Subscription import Subscription
 # from core.models.session.SessionProfile import SessionProfile
 # from core.models.system.SystemProfile import SystemProfile
-# from core.models.orm_models.Location import Location
+from core.models.orm_models.Location import Location
 # from core.models.orm_models.Operator import Operator
 
 SYSTEMWIDE_CHOICES = ['wireguard', 'hysteria2', 'vless']
@@ -93,12 +93,15 @@ def update_profile(
     # ============= LOCATION =============
     elif key == "location":
         country_code, city_code = new_value.split("_")
-        location = get_profile_location_data(   # SQLAlchemy
+        new_location = get_profile_location_data(   # SQLAlchemy
             country_code=country_code,
             city_code=city_code
         )
-        # SQLAlchemy foreign key assignment — fills in id, timezone, operator, etc.
+        pre_existing_protocol = profile.connection.code
-        profile.location = location
+        if new_location.supports_protocol(pre_existing_protocol):
+            profile.location = new_location
+        else:
+            return Result(valid=False, message=f"The {new_location.name} location doesn't currently support the {str(pre_existing_protocol)} protocol", error_type=ResultError.WRONG_LOCATION)
 
         # Subscription gets wiped on a location change, and is outside ORM
         profile.subscription = None
@@ -117,12 +120,23 @@ def update_profile(
 
 def change_protocol(profile: Profile, new_value: str) -> Result:
     logger.info(f"changing the protcol to {new_value}")
+
     # ============= SAME VALUE =============
     if profile.connection.code == new_value:
         error_msg = f"This is NOT changing the protocol, it already was {new_value}"
         logger.error(error_msg)
         return Result(valid=False, message=error_msg, error_type=ResultError.INVALID_INPUT)
 
+    # ============= LOCATION SUPPORTS PROTOCOL =============
+    logger.info("checking protocol support for this location..")
+    pre_existing_location = profile.location
+
+    if not pre_existing_location.supports_protocol(new_value):
+        logger.error(f"User is trying to update the protocol to {new_value}, but selected an invalid location to do so.")
+        error_msg = f"This location does not yet support the {str(new_value)} protocol."
+        return Result(valid=False, message=error_msg, error_type=ResultError.WRONG_LOCATION)
+    logger.info("Passed location check")
+
     # ============= SYSTEMWIDE =============
     logger.info(f"Evaluating the profile's connection type of {profile.type}")
     if profile.type == ProfileType.SYSTEM:

core/models/orm_models/Location.py

@@ -56,9 +56,9 @@ class Location(BaseModel):
 
 
 
-    def supports_protocol(self, enum_value: ProtocolChoice) -> bool:
+    def supports_protocol(self, enum_value: ProtocolChoice|str) -> bool:
-        # print(f"inside {self.country_name} {self.name} location model's supports_protocol, the submitted protocol is {enum_value}")
+        """Check if a given location supports a given protocol"""
-        print(f"enum_value type is {type(enum_value)}")
+        
         if enum_value is None:
             return False
 
@@ -72,7 +72,6 @@ class Location(BaseModel):
             string_version = enum_value
 
         field_name = f'is_{string_version}_capable'
-        print(f"field_name is {field_name}")
         return bool(getattr(self, field_name))