1 files changed, 51 insertions, 9 deletions

diff --git a/src/studio.cpp b/src/studio.cpp
index 0c988dc..01f4830 100644
--- a/src/studio.cpp
+++ b/src/studio.cpp
@@ -25,10 +25,15 @@ namespace {
 using SetColorFn = void (*)(void *, const float *);
 using SetAlphaFn = void (*)(void *, float);
 using DrawStaticFn = void (*)(void *, void *, void *, int);
+using DrawModelFn = void (*)(void *, void *, void *, void *, void *, void *,
+                             int);
+using DrawArrayFn = void (*)(void *, void *, int, void *, int, int);
 
 SetColorFn g_orig_set_color = nullptr;
 SetAlphaFn g_orig_set_alpha = nullptr;
 DrawStaticFn g_orig_draw_static = nullptr;
+DrawModelFn g_orig_draw_model = nullptr;
+DrawArrayFn g_orig_draw_array = nullptr;
 
 bool g_install_tried = false;
 bool g_installed = false;
@@ -89,6 +94,44 @@ void hook_set_alpha(void *self, float alpha) {
   g_orig_set_alpha(self, alpha);
 }
 
+void force_prop_color(void *self) {
+  float c[4];
+  prop_color(c);
+  g_orig_set_color(self, c);
+  g_orig_set_alpha(self, 1.0f);
+}
+
+void hook_draw_model(void *self, void *results, void *info, void *bone,
+                     void *flex, void *flex_delayed, int flags) {
+  const char *model = info_model_name(info);
+  bool prop = prop_tint_on() && is_prop_model(model);
+  if (g_static_calls < 3 && prop)
+    logf("studio: DrawModel prop model=\"%s\"", model ? model : "?");
+  if (prop) {
+    force_prop_color(self);
+    g_prop_depth++;
+  }
+  g_orig_draw_model(self, results, info, bone, flex, flex_delayed, flags);
+  if (prop)
+    g_prop_depth--;
+}
+
+void hook_draw_array(void *self, void *info, int count, void *instances,
+                     int stride, int flags) {
+  const char *model = info_model_name(info);
+  bool prop = prop_tint_on() && is_prop_model(model);
+  if (g_static_calls < 3)
+    logf("studio: DrawModelArray count=%d model=\"%s\" prop=%d", count,
+         model ? model : "?", (int)prop);
+  if (prop) {
+    force_prop_color(self);
+    g_prop_depth++;
+  }
+  g_orig_draw_array(self, info, count, instances, stride, flags);
+  if (prop)
+    g_prop_depth--;
+}
+
 void hook_draw_static_prop(void *self, void *info, void *model_to_world,
                            int flags) {
   const char *model = info_model_name(info);
@@ -103,13 +146,7 @@ void hook_draw_static_prop(void *self, void *info, void *model_to_world,
     logf("studio: static draws=%d prop_draws=%d", g_static_calls,
          g_static_props);
   if (prop) {
-    // The engine may have set the modulation just before this call; override
-    // it here as well as from the SetColorModulation hook (the static-prop
-    // renderer can call it either side of DrawModelStaticProp).
-    float c[4];
-    prop_color(c);
-    g_orig_set_color(self, c);
-    g_orig_set_alpha(self, 1.0f);
+    force_prop_color(self);
     g_prop_depth++;
   }
   g_orig_draw_static(self, info, model_to_world, flags);
@@ -178,13 +215,18 @@ void studio_props_init() {
 
   g_orig_set_color = (SetColorFn)vt[27];
   g_orig_set_alpha = (SetAlphaFn)vt[28];
+  g_orig_draw_model = (DrawModelFn)vt[29];
   g_orig_draw_static = (DrawStaticFn)vt[30];
+  g_orig_draw_array = (DrawArrayFn)vt[46];
   vt[27] = (void *)&hook_set_color;
   vt[28] = (void *)&hook_set_alpha;
+  vt[29] = (void *)&hook_draw_model;
   vt[30] = (void *)&hook_draw_static_prop;
+  vt[46] = (void *)&hook_draw_array;
 
   *(void ***)sr = vt;
   g_installed = true;
-  logf("studio: IStudioRender %p hooked (setcolor=%p drawstatic=%p)", sr,
-       (void *)g_orig_set_color, (void *)g_orig_draw_static);
+  logf("studio: IStudioRender %p hooked (color=%p draw=%p static=%p "
+       "array=%p)", sr, (void *)g_orig_set_color, (void *)g_orig_draw_model,
+       (void *)g_orig_draw_static, (void *)g_orig_draw_array);
 }