From 51ba5ffadf6d17a7f71a314e1be830c1ae04e67b Mon Sep 17 00:00:00 2001
From: Daniel Pope <mauve@mauveweb.co.uk>
Date: Fri, 24 Jul 2020 09:55:57 +0100
Subject: [PATCH] Fix OOB access in swizzle; add '0' and '1' swizzles

---
 src_c/math.c      | 27 +++++++++++++++++++--------
 test/math_test.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/src_c/math.c b/src_c/math.c
index 5ae4dab493..5a5ed72f84 100644
--- a/src_c/math.c
+++ b/src_c/math.c
@@ -1584,6 +1584,7 @@ vector_str(pgVector *self)
 static PyObject *
 vector_getAttr_swizzle(pgVector *self, PyObject *attr_name)
 {
+    double value;
     double *coords;
     Py_ssize_t i, idx, len;
     PyObject *attr_unicode = NULL;
@@ -1634,23 +1635,33 @@ vector_getAttr_swizzle(pgVector *self, PyObject *attr_name)
             case 'y':
             case 'z':
                 idx = attr[i] - 'x';
-                break;
+                goto swizzle_idx;
             case 'w':
                 idx = 3;
+
+        swizzle_idx:
+                if (idx >= self->dim) {
+                    goto swizzle_failed;
+                };
+                value = coords[idx];
+                break;
+
+            case '0':
+                value = 0.0f;
                 break;
+            case '1':
+                value = 1.0f;
+                break;
+
             default:
                 goto swizzle_failed;
         }
         if (len == 2 || len == 3) {
-            ((pgVector *)res)->coords[i] = coords[idx];
-        }
-        else if (idx < self->dim) {
-            if (PyTuple_SetItem(res, i, PyFloat_FromDouble(coords[idx])) != 0)
+            ((pgVector *)res)->coords[i] = value;
+        } else {
+            if (PyTuple_SetItem(res, i, PyFloat_FromDouble(value)) != 0)
                 goto internal_error;
         }
-        else {
-            goto swizzle_failed;
-        }
     }
     /* swizzling succeeded! */
     Py_DECREF(attr_unicode);
diff --git a/test/math_test.py b/test/math_test.py
index ef22246f36..41fd502489 100644
--- a/test/math_test.py
+++ b/test/math_test.py
@@ -2111,6 +2111,34 @@ class Vector(pygame.math.Vector3):
         v *= 2
         self.assertEqual(v, (4.0, 4.0, 4.0))
 
+    def test_swizzle_constants(self):
+        """We can get constant values from a swizzle."""
+        v = Vector2(7, 6)
+        self.assertEqual(
+            v.xy1,
+            (7.0, 6.0, 1.0),
+        )
+
+    def test_swizzle_four_constants(self):
+        """We can get 4 constant values from a swizzle."""
+        v = Vector2(7, 6)
+        self.assertEqual(
+            v.xy01,
+            (7.0, 6.0, 0.0, 1.0),
+        )
+
+    def test_swizzle_oob(self):
+        """An out-of-bounds swizzle raises an AttributeError."""
+        v = Vector2(7, 6)
+        with self.assertRaises(AttributeError):
+            v.xyz
+
+    def test_swizzle_set_oob(self):
+        """An out-of-bounds swizzle set raises an AttributeError."""
+        v = Vector2(7, 6)
+        with self.assertRaises(AttributeError):
+            v.xz = (1, 1)
+
 
 if __name__ == "__main__":
     unittest.main()