1 files changed, 11 insertions, 4 deletions

diff --git a/source/blender/editors/space_node/node_templates.c b/source/blender/editors/space_node/node_templates.c
index 7d4095326f8..da0058585cb 100644
--- a/source/blender/editors/space_node/node_templates.c
+++ b/source/blender/editors/space_node/node_templates.c
@@ -45,6 +45,7 @@
 #include "UI_interface.h"
 
 #include "ED_node.h" /* own include */
+#include "node_intern.h"
 
 #include "ED_undo.h"
 
@@ -662,17 +663,23 @@ static void ui_template_node_link_menu(bContext *C, uiLayout *layout, void *but_
   ui_node_menu_column(arg, NODE_CLASS_GROUP, N_("Group"));
 }
 
-void uiTemplateNodeLink(uiLayout *layout, bNodeTree *ntree, bNode *node, bNodeSocket *sock)
+void uiTemplateNodeLink(
+    uiLayout *layout, bContext *C, bNodeTree *ntree, bNode *node, bNodeSocket *sock)
 {
   uiBlock *block = uiLayoutGetBlock(layout);
   NodeLinkArg *arg;
   uiBut *but;
+  float socket_col[4];
 
   arg = MEM_callocN(sizeof(NodeLinkArg), "NodeLinkArg");
   arg->ntree = ntree;
   arg->node = node;
   arg->sock = sock;
 
+  PointerRNA node_ptr;
+  RNA_pointer_create((ID *)ntree, &RNA_Node, node, &node_ptr);
+  node_socket_color_get(C, ntree, &node_ptr, sock, socket_col);
+
   UI_block_layout_set_current(block, layout);
 
   if (sock->link || sock->type == SOCK_SHADER || (sock->flag & SOCK_HIDE_VALUE)) {
@@ -687,8 +694,8 @@ void uiTemplateNodeLink(uiLayout *layout, bNodeTree *ntree, bNode *node, bNodeSo
   }
 
   UI_but_type_set_menu_from_pulldown(but);
+  UI_but_node_link_set(but, sock, socket_col);
 
-  but->flag |= UI_BUT_NODE_LINK;
   but->poin = (char *)but;
   but->func_argN = arg;
 
@@ -798,7 +805,7 @@ static void ui_node_draw_input(
   }
   else if (lnode) {
     /* input linked to a node */
-    uiTemplateNodeLink(split, ntree, node, input);
+    uiTemplateNodeLink(split, C, ntree, node, input);
 
     if (depth == 0 || !(input->flag & SOCK_COLLAPSED)) {
       if (depth == 0) {
@@ -835,7 +842,7 @@ static void ui_node_draw_input(
       row = uiLayoutRow(split, false);
     }
 
-    uiTemplateNodeLink(row, ntree, node, input);
+    uiTemplateNodeLink(row, C, ntree, node, input);
   }
 
   /* clear */