diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBase.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBase.scala index 9b59ce6d925..613a77aab39 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBase.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/HuggingFaceCodegenBase.scala @@ -72,9 +72,11 @@ object HuggingFaceCodegenBase { | |# Defensive format check for MODEL_ID before it is interpolated into |# HF URL paths. The base host is hardcoded so the worst case isn't - |# SSRF, but rejecting `..` segments / query strings / fragments / - |# control chars keeps the operator's request shape predictable. - |_HF_MODEL_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*(/[A-Za-z0-9._-]+)+$$") + |# SSRF, but rejecting `..` traversal / query strings / fragments / + |# control chars keeps the operator's request shape predictable. The + |# leading (?!.*\.\.) rejects any `..`; the trailing /segment group is + |# optional so single-segment legacy IDs like `gpt2` are also accepted. + |_HF_MODEL_ID_PATTERN = re.compile(r"^(?!.*\.\.)[A-Za-z0-9][A-Za-z0-9._-]*(/[A-Za-z0-9._-]+)*$$") | |class ProcessTableOperator(UDFTableOperator): | @@ -488,7 +490,7 @@ object HuggingFaceCodegenBase { | if not _HF_MODEL_ID_PATTERN.match(self.MODEL_ID or ""): | raise ValueError( | f"Invalid Hugging Face model ID '{self.MODEL_ID}'. " - | f"Expected format like 'org/model-name' or 'org/model-name/revision'." + | f"Expected a model ID like 'gpt2', 'org/model-name', or 'org/model-name/revision'." | ) | | # --- resolve API token --- diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala index bf6549fb802..439b111fae9 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala @@ -141,6 +141,10 @@ class HuggingFaceInferenceOpDescSpec extends AnyFlatSpec with Matchers { code should include("if not _HF_MODEL_ID_PATTERN.match(") code should include("raise ValueError(") code should include("Invalid Hugging Face model ID") + // #7196: reject `..` path traversal (leading negative lookahead) and accept + // single-segment legacy IDs like `gpt2` (trailing /segment group optional). + code should include("(?!.*\\.\\.)") + code should include("(/[A-Za-z0-9._-]+)*$") } it should "not leak raw user-input strings into the generated Python source" in {