wined3d: Validate the SM4 token count (AFL).

dlls/d3d11/shader.c

@@ -76,6 +76,7 @@ static HRESULT shdr_handler(const char *data, DWORD data_size, DWORD tag, void *
             if (desc->byte_code)
                 FIXME("Multiple shader code chunks.\n");
             desc->byte_code = (const DWORD *)data;
+            desc->byte_code_size = data_size;
             break;
 
         case TAG_AON9:
@@ -103,6 +104,7 @@ static HRESULT shdr_handler(const char *data, DWORD data_size, DWORD tag, void *
                 if (desc->byte_code)
                     FIXME("Multiple shader code chunks.\n");
                 desc->byte_code = (const DWORD *)byte_code;
+                desc->byte_code_size = data_size - header->byte_code_offset;
                 TRACE("Feature level 9 shader version 0%08x, 0%08x.\n", header->shader_version, *desc->byte_code);
             }
             else
@@ -126,6 +128,7 @@ static HRESULT shader_extract_from_dxbc(const void *dxbc, SIZE_T dxbc_length, st
     HRESULT hr;
 
     desc->byte_code = NULL;
+    desc->byte_code_size = 0;
     memset(&desc->input_signature, 0, sizeof(desc->input_signature));
     memset(&desc->output_signature, 0, sizeof(desc->output_signature));
 

dlls/d3d8/shader.c

@@ -117,6 +117,7 @@ HRESULT d3d8_vertex_shader_init(struct d3d8_vertex_shader *shader, struct d3d8_d
             FIXME("Usage %#x not implemented.\n", usage);
 
         desc.byte_code = byte_code;
+        desc.byte_code_size = ~(size_t)0;
         desc.input_signature.element_count = 0;
         desc.output_signature.element_count = 0;
         desc.max_version = 1;
@@ -166,6 +167,7 @@ HRESULT d3d8_pixel_shader_init(struct d3d8_pixel_shader *shader, struct d3d8_dev
     shader->handle = shader_handle;
 
     desc.byte_code = byte_code;
+    desc.byte_code_size = ~(size_t)0;
     desc.input_signature.element_count = 0;
     desc.output_signature.element_count = 0;
     desc.max_version = 1;

dlls/d3d9/shader.c

@@ -143,6 +143,7 @@ HRESULT vertexshader_init(struct d3d9_vertexshader *shader, struct d3d9_device *
     shader->IDirect3DVertexShader9_iface.lpVtbl = &d3d9_vertexshader_vtbl;
 
     desc.byte_code = byte_code;
+    desc.byte_code_size = ~(size_t)0;
     desc.input_signature.element_count = 0;
     desc.output_signature.element_count = 0;
     desc.max_version = 3;
@@ -294,6 +295,7 @@ HRESULT pixelshader_init(struct d3d9_pixelshader *shader, struct d3d9_device *de
     shader->IDirect3DPixelShader9_iface.lpVtbl = &d3d9_pixelshader_vtbl;
 
     desc.byte_code = byte_code;
+    desc.byte_code_size = ~(size_t)0;
     desc.input_signature.element_count = 0;
     desc.output_signature.element_count = 0;
     desc.max_version = 3;

dlls/wined3d/shader.c

@@ -2854,7 +2854,7 @@ const struct wined3d_shader_backend_ops none_shader_backend =
 };
 
 static HRESULT shader_set_function(struct wined3d_shader *shader, const DWORD *byte_code,
-        DWORD float_const_count, enum wined3d_shader_type type, unsigned int max_version)
+        size_t byte_code_size, DWORD float_const_count, enum wined3d_shader_type type, unsigned int max_version)
 {
     struct wined3d_shader_reg_maps *reg_maps = &shader->reg_maps;
     const struct wined3d_shader_frontend *fe;
@@ -2862,8 +2862,8 @@ static HRESULT shader_set_function(struct wined3d_shader *shader, const DWORD *b
     unsigned int backend_version;
     const struct wined3d_d3d_info *d3d_info = &shader->device->adapter->d3d_info;
 
-    TRACE("shader %p, byte_code %p, float_const_count %u, type %#x, max_version %u.\n",
-            shader, byte_code, float_const_count, type, max_version);
+    TRACE("shader %p, byte_code %p, byte_code_size %#lx, float_const_count %u, type %#x, max_version %u.\n",
+            shader, byte_code, (long)byte_code_size, float_const_count, type, max_version);
 
     list_init(&shader->constantsF);
     list_init(&shader->constantsB);
@@ -2878,7 +2878,7 @@ static HRESULT shader_set_function(struct wined3d_shader *shader, const DWORD *b
         return WINED3DERR_INVALIDCALL;
     }
     shader->frontend = fe;
-    shader->frontend_data = fe->shader_init(byte_code, &shader->output_signature);
+    shader->frontend_data = fe->shader_init(byte_code, byte_code_size, &shader->output_signature);
     if (!shader->frontend_data)
     {
         FIXME("Failed to initialize frontend.\n");
@@ -3193,7 +3193,7 @@ static HRESULT shader_init(struct wined3d_shader *shader, struct wined3d_device
     list_add_head(&device->shaders, &shader->shader_list_entry);
 
     if (FAILED(hr = shader_set_function(shader, desc->byte_code,
-            float_const_count, type, desc->max_version)))
+            desc->byte_code_size, float_const_count, type, desc->max_version)))
     {
         WARN("Failed to set function, hr %#x.\n", hr);
         shader_cleanup(shader);

dlls/wined3d/shader_sm1.c

@@ -533,7 +533,8 @@ static unsigned int shader_skip_unrecognized(const struct wined3d_sm1_data *priv
     return tokens_read;
 }
 
-static void *shader_sm1_init(const DWORD *byte_code, const struct wined3d_shader_signature *output_signature)
+static void *shader_sm1_init(const DWORD *byte_code, size_t byte_code_size,
+        const struct wined3d_shader_signature *output_signature)
 {
     struct wined3d_sm1_data *priv;
     BYTE major, minor;

dlls/wined3d/shader_sm4.c

@@ -1121,17 +1121,30 @@ static enum wined3d_data_type map_data_type(char t)
     }
 }
 
-static void *shader_sm4_init(const DWORD *byte_code, const struct wined3d_shader_signature *output_signature)
+static void *shader_sm4_init(const DWORD *byte_code, size_t byte_code_size,
+        const struct wined3d_shader_signature *output_signature)
 {
     DWORD version_token, token_count;
     struct wined3d_sm4_data *priv;
     unsigned int i;
 
+    if (byte_code_size / sizeof(*byte_code) < 2)
+    {
+        WARN("Invalid byte code size %lu.\n", (long)byte_code_size);
+        return NULL;
+    }
+
     version_token = byte_code[0];
     TRACE("Version: 0x%08x.\n", version_token);
     token_count = byte_code[1];
     TRACE("Token count: %u.\n", token_count);
 
+    if (token_count < 2 || byte_code_size / sizeof(*byte_code) < token_count)
+    {
+        WARN("Invalid token count %u.\n", token_count);
+        return NULL;
+    }
+
     if (!(priv = HeapAlloc(GetProcessHeap(), 0, sizeof(*priv))))
     {
         ERR("Failed to allocate private data\n");

dlls/wined3d/wined3d_private.h

@@ -1083,7 +1083,8 @@ struct wined3d_shader_loop_control
 
 struct wined3d_shader_frontend
 {
-    void *(*shader_init)(const DWORD *ptr, const struct wined3d_shader_signature *output_signature);
+    void *(*shader_init)(const DWORD *byte_code, size_t byte_code_size,
+            const struct wined3d_shader_signature *output_signature);
     void (*shader_free)(void *data);
     void (*shader_read_header)(void *data, const DWORD **ptr, struct wined3d_shader_version *shader_version);
     void (*shader_read_instruction)(void *data, const DWORD **ptr, struct wined3d_shader_instruction *ins);

include/wine/wined3d.h

@@ -1959,6 +1959,7 @@ struct wined3d_shader_signature
 struct wined3d_shader_desc
 {
     const DWORD *byte_code;
+    size_t byte_code_size;
     struct wined3d_shader_signature input_signature;
     struct wined3d_shader_signature output_signature;
     unsigned int max_version;