From d0e4023d88d3ba253be5fb8080aca47bce0a7e0c Mon Sep 17 00:00:00 2001
From: Hajime Hoshi <hajimehoshi@gmail.com>
Date: Sat, 29 Jul 2023 18:57:24 +0900
Subject: [PATCH] internal/ui: add a length check for uniform variables

---
 image.go              |   4 ++
 internal/ui/shader.go |  12 +++++
 shader_test.go        | 106 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 122 insertions(+)

diff --git a/image.go b/image.go
index f2cbb653d..30cd54fce 100644
--- a/image.go
+++ b/image.go
@@ -577,6 +577,8 @@ var _ [len(DrawTrianglesShaderOptions{}.Images) - graphics.ShaderImageCount]stru
 //
 // When a specified image is non-nil and is disposed, DrawTrianglesShader panics.
 //
+// If a specified uniform variable's length or type doesn't match with an expected one, DrawTrianglesShader panics.
+//
 // When the image i is disposed, DrawTrianglesShader does nothing.
 func (i *Image) DrawTrianglesShader(vertices []Vertex, indices []uint16, shader *Shader, options *DrawTrianglesShaderOptions) {
 	i.copyCheck()
@@ -724,6 +726,8 @@ var _ [len(DrawRectShaderOptions{}.Images)]struct{} = [graphics.ShaderImageCount
 //
 // When one of the specified image is non-nil and is disposed, DrawRectShader panics.
 //
+// If a specified uniform variable's length or type doesn't match with an expected one, DrawRectShader panics.
+//
 // When the image i is disposed, DrawRectShader does nothing.
 func (i *Image) DrawRectShader(width, height int, shader *Shader, options *DrawRectShaderOptions) {
 	i.copyCheck()
diff --git a/internal/ui/shader.go b/internal/ui/shader.go
index ded2e4af6..d5d7cb473 100644
--- a/internal/ui/shader.go
+++ b/internal/ui/shader.go
@@ -69,13 +69,25 @@ func (s *Shader) AppendUniforms(dst []uint32, uniforms map[string]any) []uint32
 			t := v.Type()
 			switch t.Kind() {
 			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+				if typ.Uint32Count() != 1 {
+					panic(fmt.Sprintf("ui: unexpected uniform value for %s (%s)", name, typ.String()))
+				}
 				dst[idx] = uint32(v.Int())
 			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
+				if typ.Uint32Count() != 1 {
+					panic(fmt.Sprintf("ui: unexpected uniform value for %s (%s)", name, typ.String()))
+				}
 				dst[idx] = uint32(v.Uint())
 			case reflect.Float32, reflect.Float64:
+				if typ.Uint32Count() != 1 {
+					panic(fmt.Sprintf("ui: unexpected uniform value for %s (%s)", name, typ.String()))
+				}
 				dst[idx] = math.Float32bits(float32(v.Float()))
 			case reflect.Slice, reflect.Array:
 				l := v.Len()
+				if typ.Uint32Count() != l {
+					panic(fmt.Sprintf("ui: unexpected uniform value for %s (%s)", name, typ.String()))
+				}
 				switch t.Elem().Kind() {
 				case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
 					for i := 0; i < l; i++ {
diff --git a/shader_test.go b/shader_test.go
index 95c886471..7aa9fe30e 100644
--- a/shader_test.go
+++ b/shader_test.go
@@ -1854,3 +1854,109 @@ func Fragment(position vec4, texCoord vec2, color vec4) vec4 {
 		t.Errorf("got: %v, want: %v", got, want)
 	}
 }
+
+func TestShaderUniformSizes(t *testing.T) {
+	const w, h = 16, 16
+
+	dst := ebiten.NewImage(w, h)
+	s, err := ebiten.NewShader([]byte(`//kage:unit pixel
+
+package main
+
+var U vec4
+var V [3]float
+
+func Fragment(position vec4, texCoord vec2, color vec4) vec4 {
+	return vec4(0)
+}
+`))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	tests := []struct {
+		uniforms map[string]any
+		err      bool
+	}{
+		{
+			uniforms: nil,
+			err:      false,
+		},
+		{
+			uniforms: map[string]any{
+				"U": 1,
+			},
+			err: true,
+		},
+		{
+			uniforms: map[string]any{
+				"U": "1",
+			},
+			err: true,
+		},
+		{
+			uniforms: map[string]any{
+				"U": []int32{1, 2, 3},
+			},
+			err: true,
+		},
+		{
+			uniforms: map[string]any{
+				"U": []int32{1, 2, 3, 4},
+			},
+			err: false,
+		},
+		{
+			uniforms: map[string]any{
+				"U": []int32{1, 2, 3, 4, 5},
+			},
+			err: true,
+		},
+		{
+			uniforms: map[string]any{
+				"V": 1,
+			},
+			err: true,
+		},
+		{
+			uniforms: map[string]any{
+				"V": "1",
+			},
+			err: true,
+		},
+		{
+			uniforms: map[string]any{
+				"V": []int32{1, 2},
+			},
+			err: true,
+		},
+		{
+			uniforms: map[string]any{
+				"V": []int32{1, 2, 3},
+			},
+			err: false,
+		},
+		{
+			uniforms: map[string]any{
+				"V": []int32{1, 2, 3, 4},
+			},
+			err: true,
+		},
+	}
+	for _, tc := range tests {
+		tc := tc
+		t.Run(fmt.Sprintf("%v", tc.uniforms), func(t *testing.T) {
+			defer func() {
+				r := recover()
+				if r != nil && !tc.err {
+					t.Errorf("DrawRectShader must not panic but did")
+				} else if r == nil && tc.err {
+					t.Errorf("DrawRectShader must panic but does not")
+				}
+			}()
+			op := &ebiten.DrawRectShaderOptions{}
+			op.Uniforms = tc.uniforms
+			dst.DrawRectShader(w, h, s, op)
+		})
+	}
+}