[JIT] Do not allow creating generics with None types (#44958)

malfet · facebook-github-bot · commit 989d877c95a9 · 2020-09-22T21:50:40.000-07:00
Summary: Otherwise, invoking something like `python -c "import torch._C;print(torch._C.ListType(None))"` will result in SIGSEGV Discovered while trying to create a torch script for function with the following type annotation `Tuple[int, Ellipsis] -> None` Pull Request resolved: #44958 Reviewed By: suo Differential Revision: D23799906 Pulled By: malfet fbshipit-source-id: 916a243007d13ed3e7a5b282dd712da3d66e3bf7
diff --git a/aten/src/ATen/core/jit_type.h b/aten/src/ATen/core/jit_type.h
@@ -263,7 +263,12 @@ struct SingleElementType : public Type {
   }
 
  protected:
-  SingleElementType(TypePtr elem) : Type(Kind), elem(std::move(elem)) {}
+  SingleElementType(TypePtr elem) : Type(Kind), elem(std::move(elem)) {
+    if (!this->elem) {
+      throw std::runtime_error(c10::str(
+            "Can not create ", typeKindToString(Kind), " with None type"));
+    }
+  }
 
  private:
   TypePtr elem;
diff --git a/aten/src/ATen/core/type.cpp b/aten/src/ATen/core/type.cpp
@@ -716,6 +716,9 @@ TupleType::TupleType(
       schema_(std::move(schema)) {
   has_free_variables_ =
       std::any_of(elements_.begin(), elements_.end(), [](TypePtr v) {
+        if (!v) {
+          throw std::runtime_error("Can not create tuple with None type");
+        }
         return v->hasFreeVariables();
       });
   if (schema_) {
diff --git a/test/jit/test_list_dict.py b/test/jit/test_list_dict.py
@@ -1155,6 +1155,11 @@ def annotated_fn(x: torch.Tensor) -> List:
         with self.assertRaisesRegex(RuntimeError, r"Attempted to use List without a contained type"):
             torch.jit.script(annotated_fn)
 
+    def test_list_none(self):
+        with self.assertRaisesRegex(RuntimeError, "Can not create ListType with None type"):
+            x = torch._C.ListType(None)
+
+
 
 class TestDict(JitTestCase):
     def dict(self):
