3 files changed, 15 insertions, 8 deletions
diff --git a/src/gallium/auxiliary/draw/draw_vs_exec.c b/src/gallium/auxiliary/draw/draw_vs_exec.c
index 54a2b2ab04..7a02f6334b 100644
--- a/src/gallium/auxiliary/draw/draw_vs_exec.c
+++ b/src/gallium/auxiliary/draw/draw_vs_exec.c
@@ -39,6 +39,7 @@
 #include "draw_vs.h"
 
 #include "tgsi/util/tgsi_parse.h"
+#include "tgsi/util/tgsi_scan.h"
 
 
 struct exec_vertex_shader {
@@ -165,21 +166,23 @@ draw_create_vs_exec(struct draw_context *draw,
 		    const struct pipe_shader_state *state)
 {
    struct exec_vertex_shader *vs = CALLOC_STRUCT( exec_vertex_shader );
-   uint nt = tgsi_num_tokens(state->tokens);
 
    if (vs == NULL) 
       return NULL;
 
    /* we make a private copy of the tokens */
-   vs->base.state.tokens = mem_dup(state->tokens, nt * sizeof(state->tokens[0]));
-   tgsi_scan_shader(state->tokens, &vs->base.info);
+   vs->base.state.tokens = tgsi_dup_tokens(state->tokens);
+   if (!vs->base.state.tokens) {
+      FREE(vs);
+      return NULL;
+   }
 
+   tgsi_scan_shader(state->tokens, &vs->base.info);
 
    vs->base.prepare = vs_exec_prepare;
    vs->base.run_linear = vs_exec_run_linear;
    vs->base.delete = vs_exec_delete;
    vs->machine = &draw->machine;
 
-
    return &vs->base;
 }
diff --git a/src/gallium/auxiliary/draw/draw_vs_llvm.c b/src/gallium/auxiliary/draw/draw_vs_llvm.c
index dcada66514..be6907ed04 100644
--- a/src/gallium/auxiliary/draw/draw_vs_llvm.c
+++ b/src/gallium/auxiliary/draw/draw_vs_llvm.c
@@ -138,14 +138,17 @@ draw_create_vs_llvm(struct draw_context *draw,
 		    const struct pipe_shader_state *templ)
 {
    struct draw_llvm_vertex_shader *vs;
-   uint nt = tgsi_num_tokens(templ->tokens);
 
    vs = CALLOC_STRUCT( draw_llvm_vertex_shader );
    if (vs == NULL) 
       return NULL;
 
    /* we make a private copy of the tokens */
-   vs->base.state.tokens = mem_dup(templ->tokens, nt * sizeof(templ->tokens[0]));
+   vs->base.state.tokens = tgsi_dup_tokens(templ->tokens);
+   if (!vs->base.state.tokens) {
+      FREE(vs);
+      return NULL;
+   }
 
    tgsi_scan_shader(vs->base.state.tokens, &vs->base.info);
 
diff --git a/src/gallium/auxiliary/draw/draw_vs_sse.c b/src/gallium/auxiliary/draw/draw_vs_sse.c
index a57c938fbf..5929ea76b2 100644
--- a/src/gallium/auxiliary/draw/draw_vs_sse.c
+++ b/src/gallium/auxiliary/draw/draw_vs_sse.c
@@ -188,7 +188,6 @@ draw_create_vs_sse(struct draw_context *draw,
                           const struct pipe_shader_state *templ)
 {
    struct draw_sse_vertex_shader *vs;
-   uint nt = tgsi_num_tokens(templ->tokens);
 
    if (!rtasm_cpu_has_sse2())
       return NULL;
@@ -198,7 +197,9 @@ draw_create_vs_sse(struct draw_context *draw,
       return NULL;
 
    /* we make a private copy of the tokens */
-   vs->base.state.tokens = mem_dup(templ->tokens, nt * sizeof(templ->tokens[0]));
+   vs->base.state.tokens = tgsi_dup_tokens(templ->tokens);
+   if (!vs->base.state.tokens)
+      goto fail;
 
    tgsi_scan_shader(templ->tokens, &vs->base.info);