9094 want writable uregs[] in dtrace

Reviewed by: Jerry Jelinek <jerry.jelinek@joyent.com>
Reviewed by: Alex Wilson <alex.wilson@joyent.com>
Reviewed by: Bryan Cantrill <bryan@joyent.com>
Reviewed by: Adam Leventhal <adam.leventhal@gmail.com>
Approved by: Richard Lowe <richlowe@richlowe.net>

15 files changed, 484 insertions, 50 deletions

diff --git a/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite1.c b/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite1.c
new file mode 100644
index 0000000000..fba70fb709
--- /dev/null
+++ b/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite1.c
@@ -0,0 +1,42 @@
+/*
+ * This file and its contents are supplied under the terms of the
+ * Common Development and Distribution License ("CDDL"), version 1.0.
+ * You may only use this file in accordance with the terms of version
+ * 1.0 of the CDDL.
+ *
+ * A full copy of the text of the CDDL should have accompanied this
+ * source.  A copy of the CDDL is also available via the Internet at
+ * http://www.illumos.org/license/CDDL.
+ */
+
+/*
+ * Copyright 2017 Joyent, Inc.
+ */
+
+#include <unistd.h>
+
+int
+startup_wait(int *val)
+{
+	return (*val);
+}
+
+int
+baz(void)
+{
+	return (0);
+}
+
+int
+main(int argc, char **argv)
+{
+	int wait = 1;
+
+	while (startup_wait(&wait)) {
+		usleep(1000);
+	}
+
+	(void) baz();
+
+	return (0);
+}
diff --git a/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite1.d b/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite1.d
new file mode 100644
index 0000000000..238fefcdc2
--- /dev/null
+++ b/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite1.d
@@ -0,0 +1,38 @@
+/*
+ * This file and its contents are supplied under the terms of the
+ * Common Development and Distribution License ("CDDL"), version 1.0.
+ * You may only use this file in accordance with the terms of version
+ * 1.0 of the CDDL.
+ *
+ * A full copy of the text of the CDDL should have accompanied this
+ * source.  A copy of the CDDL is also available via the Internet at
+ * http://www.illumos.org/license/CDDL.
+ */
+
+/*
+ * Copyright 2017 Joyent, Inc.
+ */
+
+
+#pragma D option quiet
+#pragma D option destructive
+
+pid$1:a.out:startup_wait:entry
+{
+	this->val = (int *)alloca(sizeof (int));
+	*this->val = 0;
+	copyout(this->val, arg0, sizeof (int));
+}
+
+pid$1:a.out:baz:return
+{
+	uregs[R_PC] += 1;
+	/* should not be reached due to error */
+	exit(0);
+}
+
+syscall::rexit:entry
+/pid == $1/
+{
+	exit(1);
+}
diff --git a/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite2.c b/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite2.c
new file mode 100644
index 0000000000..24d3b61966
--- /dev/null
+++ b/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite2.c
@@ -0,0 +1,22 @@
+/*
+ * This file and its contents are supplied under the terms of the
+ * Common Development and Distribution License ("CDDL"), version 1.0.
+ * You may only use this file in accordance with the terms of version
+ * 1.0 of the CDDL.
+ *
+ * A full copy of the text of the CDDL should have accompanied this
+ * source.  A copy of the CDDL is also available via the Internet at
+ * http://www.illumos.org/license/CDDL.
+ */
+
+/*
+ * Copyright 2017 Joyent, Inc.
+ */
+
+#include <unistd.h>
+
+int
+main(int argc, char **argv)
+{
+	return (0);
+}
diff --git a/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite2.d b/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite2.d
new file mode 100644
index 0000000000..9be786b4ba
--- /dev/null
+++ b/usr/src/cmd/dtrace/test/tst/i386/pid/err.uregswrite2.d
@@ -0,0 +1,22 @@
+/*
+ * This file and its contents are supplied under the terms of the
+ * Common Development and Distribution License ("CDDL"), version 1.0.
+ * You may only use this file in accordance with the terms of version
+ * 1.0 of the CDDL.
+ *
+ * A full copy of the text of the CDDL should have accompanied this
+ * source.  A copy of the CDDL is also available via the Internet at
+ * http://www.illumos.org/license/CDDL.
+ */
+
+/*
+ * Copyright 2017 Joyent, Inc.
+ */
+
+#pragma D option quiet
+
+pid$1:a.out:main:entry
+{
+	/* should cause an error due to lack of destructive option */
+	uregs[R_EAX] = 0;
+}
diff --git a/usr/src/cmd/dtrace/test/tst/i386/pid/tst.uregswrite.c b/usr/src/cmd/dtrace/test/tst/i386/pid/tst.uregswrite.c
new file mode 100644
index 0000000000..642cd40f8d
--- /dev/null
+++ b/usr/src/cmd/dtrace/test/tst/i386/pid/tst.uregswrite.c
@@ -0,0 +1,60 @@
+/*
+ * This file and its contents are supplied under the terms of the
+ * Common Development and Distribution License ("CDDL"), version 1.0.
+ * You may only use this file in accordance with the terms of version
+ * 1.0 of the CDDL.
+ *
+ * A full copy of the text of the CDDL should have accompanied this
+ * source.  A copy of the CDDL is also available via the Internet at
+ * http://www.illumos.org/license/CDDL.
+ */
+
+/*
+ * Copyright 2017 Joyent, Inc.
+ */
+
+#include <unistd.h>
+
+int
+startup_wait(int *val)
+{
+	return (*val);
+}
+
+int
+baz(void)
+{
+	return (0);
+}
+
+int bar_val = 0;
+
+int
+bar(void)
+{
+	return (bar_val++);
+}
+
+int
+main(int argc, char **argv)
+{
+	int wait = 1;
+
+	while (startup_wait(&wait)) {
+		usleep(1000);
+	}
+
+	if (baz() != 1)
+		return (1);
+
+	if (bar() != 0)
+		return (1);
+
+	if (bar() != 2)
+		return (1);
+
+	if (bar() != 2)
+		return (1);
+
+	return (0);
+}
diff --git a/usr/src/cmd/dtrace/test/tst/i386/pid/tst.uregswrite.d b/usr/src/cmd/dtrace/test/tst/i386/pid/tst.uregswrite.d
new file mode 100644
index 0000000000..0185614f59
--- /dev/null
+++ b/usr/src/cmd/dtrace/test/tst/i386/pid/tst.uregswrite.d
@@ -0,0 +1,42 @@
+/*
+ * This file and its contents are supplied under the terms of the
+ * Common Development and Distribution License ("CDDL"), version 1.0.
+ * You may only use this file in accordance with the terms of version
+ * 1.0 of the CDDL.
+ *
+ * A full copy of the text of the CDDL should have accompanied this
+ * source.  A copy of the CDDL is also available via the Internet at
+ * http://www.illumos.org/license/CDDL.
+ */
+
+/*
+ * Copyright 2017 Joyent, Inc.
+ */
+
+
+#pragma D option quiet
+#pragma D option destructive
+
+pid$1:a.out:startup_wait:entry
+{
+	this->val = (int *)alloca(sizeof (int));
+	*this->val = 0;
+	copyout(this->val, arg0, sizeof (int));
+}
+
+pid$1:a.out:baz:return
+{
+	uregs[R_EAX] = 1;
+}
+
+pid$1:a.out:bar:return
+/(uregs[R_EAX] % 2) != 0/
+{
+	uregs[R_EAX] += 1;
+}
+
+syscall::rexit:entry
+/pid == $1/
+{
+	exit(arg0);
+}
diff --git a/usr/src/lib/libdtrace/common/dt_cg.c b/usr/src/lib/libdtrace/common/dt_cg.c
index 28db9b2262..9f3625e6ee 100644
--- a/usr/src/lib/libdtrace/common/dt_cg.c
+++ b/usr/src/lib/libdtrace/common/dt_cg.c
@@ -27,6 +27,7 @@
 
 /*
  * Copyright (c) 2012 by Delphix. All rights reserved.
+ * Copyright 2017 Joyent, Inc.
  */
 
 #include <sys/types.h>
@@ -1115,23 +1116,14 @@ dt_cg_asgn_op(dt_node_t *dnp, dt_irlist_t *dlp, dt_regset_t *drp)
 	}
 
 	/*
-	 * If we are storing to a variable, generate an stv instruction from
-	 * the variable specified by the identifier.  If we are storing to a
-	 * memory address, generate code again for the left-hand side using
-	 * DT_NF_REF to get the address, and then generate a store to it.
-	 * In both paths, we assume dnp->dn_reg already has the new value.
+	 * If we are storing to a memory address, generate code again for the
+	 * left-hand side using DT_NF_REF to get the address, and then generate
+	 * a store to it.
+	 *
+	 * Both here and the other variable-store paths, we assume dnp->dn_reg
+	 * already has the new value.
 	 */
-	if (dnp->dn_left->dn_kind == DT_NODE_VAR) {
-		idp = dt_ident_resolve(dnp->dn_left->dn_ident);
-
-		if (idp->di_kind == DT_IDENT_ARRAY)
-			dt_cg_arglist(idp, dnp->dn_left->dn_args, dlp, drp);
-
-		idp->di_flags |= DT_IDFLG_DIFW;
-		instr = DIF_INSTR_STV(dt_cg_stvar(idp),
-		    idp->di_id, dnp->dn_reg);
-		dt_irlist_append(dlp, dt_cg_node_alloc(DT_LBL_NONE, instr));
-	} else {
+	if (dnp->dn_left->dn_kind != DT_NODE_VAR) {
 		uint_t rbit = dnp->dn_left->dn_flags & DT_NF_REF;
 
 		assert(dnp->dn_left->dn_flags & DT_NF_WRITABLE);
@@ -1145,7 +1137,33 @@ dt_cg_asgn_op(dt_node_t *dnp, dt_irlist_t *dlp, dt_regset_t *drp)
 
 		dnp->dn_left->dn_flags &= ~DT_NF_REF;
 		dnp->dn_left->dn_flags |= rbit;
+		return;
 	}
+
+	idp = dt_ident_resolve(dnp->dn_left->dn_ident);
+	idp->di_flags |= DT_IDFLG_DIFW;
+
+	/*
+	 * Storing to an array variable is a special case.
+	 * Only 'uregs[]' supports this for the time being.
+	 */
+	if (idp->di_kind == DT_IDENT_ARRAY &&
+	    idp->di_id <= DIF_VAR_ARRAY_MAX) {
+		dt_node_t *idx = dnp->dn_left->dn_args;
+
+		dt_cg_node(idx, dlp, drp);
+		instr = DIF_INSTR_FMT(DIF_OP_STGA, idp->di_id, idx->dn_reg,
+		    dnp->dn_reg);
+		dt_irlist_append(dlp, dt_cg_node_alloc(DT_LBL_NONE, instr));
+		dt_regset_free(drp, idx->dn_reg);
+		return;
+	}
+
+	if (idp->di_kind == DT_IDENT_ARRAY)
+		dt_cg_arglist(idp, dnp->dn_left->dn_args, dlp, drp);
+
+	instr = DIF_INSTR_STV(dt_cg_stvar(idp), idp->di_id, dnp->dn_reg);
+	dt_irlist_append(dlp, dt_cg_node_alloc(DT_LBL_NONE, instr));
 }
 
 static void
diff --git a/usr/src/lib/libdtrace/common/dt_dis.c b/usr/src/lib/libdtrace/common/dt_dis.c
index c0af36420e..60195f3970 100644
--- a/usr/src/lib/libdtrace/common/dt_dis.c
+++ b/usr/src/lib/libdtrace/common/dt_dis.c
@@ -27,7 +27,7 @@
 
 /*
  * Copyright (c) 2013 by Delphix. All rights reserved.
- * Copyright (c) 2013 Joyent, Inc. All rights reserved.
+ * Copyright (c) 2017 Joyent, Inc.
  */
 
 #include <strings.h>
@@ -47,7 +47,7 @@ dt_dis_log(const dtrace_difo_t *dp, const char *name, dif_instr_t in, FILE *fp)
 /*ARGSUSED*/
 static void
 dt_dis_branch(const dtrace_difo_t *dp, const char *name,
-	dif_instr_t in, FILE *fp)
+    dif_instr_t in, FILE *fp)
 {
 	(void) fprintf(fp, "%-4s %u", name, DIF_INSTR_LABEL(in));
 }
@@ -63,7 +63,7 @@ dt_dis_load(const dtrace_difo_t *dp, const char *name, dif_instr_t in, FILE *fp)
 /*ARGSUSED*/
 static void
 dt_dis_store(const dtrace_difo_t *dp, const char *name,
-	dif_instr_t in, FILE *fp)
+    dif_instr_t in, FILE *fp)
 {
 	(void) fprintf(fp, "%-4s %%r%u, [%%r%u]", name,
 	    DIF_INSTR_R1(in), DIF_INSTR_RD(in));
@@ -167,6 +167,19 @@ dt_dis_stv(const dtrace_difo_t *dp, const char *name, dif_instr_t in, FILE *fp)
 }
 
 static void
+dt_dis_sta(const dtrace_difo_t *dp, const char *name, dif_instr_t in, FILE *fp)
+{
+	uint_t var = DIF_INSTR_VAR(in);
+	const char *vname;
+
+	(void) fprintf(fp, "%-4s DT_VAR(%u), %%r%u, %%r%u",
+	    name, var, DIF_INSTR_R2(in), DIF_INSTR_RD(in));
+
+	if ((vname = dt_dis_varname(dp, var, dt_dis_scope(name))) != NULL)
+		(void) fprintf(fp, "\t\t! DT_VAR(%u) = \"%s\"", var, vname);
+}
+
+static void
 dt_dis_setx(const dtrace_difo_t *dp, const char *name, dif_instr_t in, FILE *fp)
 {
 	uint_t intptr = DIF_INSTR_INTEGER(in);
@@ -428,6 +441,7 @@ dt_dis(const dtrace_difo_t *dp, FILE *fp)
 		{ "rldx", dt_dis_load },	/* DIF_OP_RLDX */
 		{ "xlate", dt_dis_xlate },	/* DIF_OP_XLATE */
 		{ "xlarg", dt_dis_xlate },	/* DIF_OP_XLARG */
+		{ "stga", dt_dis_sta },		/* DIF_OP_XLARG */
 	};
 
 	const struct opent *op;
diff --git a/usr/src/lib/libdtrace/common/dt_open.c b/usr/src/lib/libdtrace/common/dt_open.c
index 38c8146039..d6e90bd953 100644
--- a/usr/src/lib/libdtrace/common/dt_open.c
+++ b/usr/src/lib/libdtrace/common/dt_open.c
@@ -21,7 +21,7 @@
 
 /*
  * Copyright (c) 2003, 2010, Oracle and/or its affiliates. All rights reserved.
- * Copyright (c) 2013, Joyent, Inc. All rights reserved.
+ * Copyright (c) 2017, Joyent, Inc.
  * Copyright (c) 2012, 2016 by Delphix. All rights reserved.
  */
 
@@ -427,8 +427,8 @@ static const dt_ident_t _dtrace_globals[] = {
 	&dt_idops_type, "uid_t" },
 { "umod", DT_IDENT_ACTFUNC, 0, DT_ACT_UMOD, DT_ATTR_STABCMN,
 	DT_VERS_1_2, &dt_idops_func, "_usymaddr(uintptr_t)" },
-{ "uregs", DT_IDENT_ARRAY, 0, DIF_VAR_UREGS, DT_ATTR_STABCMN, DT_VERS_1_0,
-	&dt_idops_regs, NULL },
+{ "uregs", DT_IDENT_ARRAY, DT_IDFLG_WRITE, DIF_VAR_UREGS, DT_ATTR_STABCMN,
+	DT_VERS_1_0, &dt_idops_regs, NULL },
 { "ustack", DT_IDENT_ACTFUNC, 0, DT_ACT_USTACK, DT_ATTR_STABCMN, DT_VERS_1_0,
 	&dt_idops_func, "stack(...)" },
 { "ustackdepth", DT_IDENT_SCALAR, 0, DIF_VAR_USTACKDEPTH,
diff --git a/usr/src/pkg/manifests/system-dtrace-tests.mf b/usr/src/pkg/manifests/system-dtrace-tests.mf
index 3a5d1e8e05..2fc1affe39 100644
--- a/usr/src/pkg/manifests/system-dtrace-tests.mf
+++ b/usr/src/pkg/manifests/system-dtrace-tests.mf
@@ -156,6 +156,12 @@ $(sparc_ONLY)file \
 $(sparc_ONLY)file \
     path=opt/SUNWdtrt/tst/$(ARCH)/pid/err.D_PROC_ALIGN.misaligned.exe \
     mode=0555
+$(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/err.uregswrite1.d mode=0444
+$(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/err.uregswrite1.exe \
+    mode=0555
+$(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/err.uregswrite2.d mode=0444
+$(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/err.uregswrite2.exe \
+    mode=0555
 $(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/tst.badinstr.d mode=0444
 $(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/tst.badinstr.exe mode=0555
 $(sparc_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/tst.br.d mode=0444
@@ -169,6 +175,9 @@ $(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/tst.ret.d mode=0444
 $(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/tst.ret.exe mode=0555
 $(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/tst.retlist.exe mode=0555
 $(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/tst.retlist.ksh mode=0444
+$(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/tst.uregswrite.d mode=0444
+$(i386_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/pid/tst.uregswrite.exe \
+    mode=0555
 $(sparc_ONLY)file path=opt/SUNWdtrt/tst/$(ARCH)/usdt/tst.tailcall.ksh \
     mode=0444
 file path=opt/SUNWdtrt/tst/$(ARCH)/ustack/tst.annotated.d mode=0444
diff --git a/usr/src/uts/common/dtrace/dtrace.c b/usr/src/uts/common/dtrace/dtrace.c
index ced7249084..b1685f42fc 100644
--- a/usr/src/uts/common/dtrace/dtrace.c
+++ b/usr/src/uts/common/dtrace/dtrace.c
@@ -21,7 +21,7 @@
 
 /*
  * Copyright (c) 2003, 2010, Oracle and/or its affiliates. All rights reserved.
- * Copyright (c) 2016, Joyent, Inc. All rights reserved.
+ * Copyright (c) 2017, Joyent, Inc.
  * Copyright (c) 2012, 2014 by Delphix. All rights reserved.
  */
 
@@ -3545,6 +3545,34 @@ dtrace_dif_variable(dtrace_mstate_t *mstate, dtrace_state_t *state, uint64_t v,
 	}
 }
 
+static void
+dtrace_dif_variable_write(dtrace_mstate_t *mstate, dtrace_state_t *state,
+    uint64_t v, uint64_t ndx, uint64_t data)
+{
+	switch (v) {
+	case DIF_VAR_UREGS: {
+		klwp_t *lwp;
+
+		if (dtrace_destructive_disallow ||
+		    !dtrace_priv_proc_control(state, mstate)) {
+			return;
+		}
+
+		if ((lwp = curthread->t_lwp) == NULL) {
+			DTRACE_CPUFLAG_SET(CPU_DTRACE_BADADDR);
+			cpu_core[CPU->cpu_id].cpuc_dtrace_illval = NULL;
+			return;
+		}
+
+		dtrace_setreg(lwp->lwp_regs, ndx, data);
+		return;
+	}
+
+	default:
+		DTRACE_CPUFLAG_SET(CPU_DTRACE_ILLOP);
+		return;
+	}
+}
 
 typedef enum dtrace_json_state {
 	DTRACE_JSON_REST = 1,
@@ -6024,6 +6052,11 @@ dtrace_dif_emulate(dtrace_difo_t *difo, dtrace_mstate_t *mstate,
 			regs[rd] = dtrace_dif_variable(mstate, state, id, 0);
 			break;
 
+		case DIF_OP_STGA:
+			dtrace_dif_variable_write(mstate, state, r1, regs[r2],
+			    regs[rd]);
+			break;
+
 		case DIF_OP_STGS:
 			id = DIF_INSTR_VAR(instr);
 
@@ -9392,6 +9425,15 @@ dtrace_difo_validate(dtrace_difo_t *dp, dtrace_vstate_t *vstate, uint_t nregs,
 			if (rd == 0)
 				err += efunc(pc, "cannot write to %r0\n");
 			break;
+		case DIF_OP_STGA:
+			if (r1 > DIF_VAR_ARRAY_MAX)
+				err += efunc(pc, "invalid array %u\n", r1);
+			if (r2 >= nregs)
+				err += efunc(pc, "invalid register %u\n", r2);
+			if (rd >= nregs)
+				err += efunc(pc, "invalid register %u\n", rd);
+			dp->dtdo_destructive = 1;
+			break;
 		case DIF_OP_LDGS:
 		case DIF_OP_LDTS:
 		case DIF_OP_LDLS:
@@ -9737,12 +9779,23 @@ dtrace_difo_validate_helper(dtrace_difo_t *dp)
 			break;
 
 		case DIF_OP_LDTA:
+			if (v < DIF_VAR_OTHER_UBASE) {
+				err += efunc(pc, "illegal variable load\n");
+				break;
+			}
+			/* FALLTHROUGH */
 		case DIF_OP_LDTS:
 		case DIF_OP_LDGAA:
 		case DIF_OP_LDTAA:
 			err += efunc(pc, "illegal dynamic variable load\n");
 			break;
 
+		case DIF_OP_STGA:
+			if (v < DIF_VAR_OTHER_UBASE) {
+				err += efunc(pc, "illegal variable store\n");
+				break;
+			}
+			/* FALLTHROUGH */
 		case DIF_OP_STTS:
 		case DIF_OP_STGAA:
 		case DIF_OP_STTAA:
diff --git a/usr/src/uts/common/sys/dtrace.h b/usr/src/uts/common/sys/dtrace.h
index 4ee1a94510..44ca7d8ae7 100644
--- a/usr/src/uts/common/sys/dtrace.h
+++ b/usr/src/uts/common/sys/dtrace.h
@@ -25,7 +25,7 @@
  */
 
 /*
- * Copyright 2016 Joyent, Inc.
+ * Copyright 2017 Joyent, Inc.
  * Copyright (c) 2013 by Delphix. All rights reserved.
  */
 
@@ -189,6 +189,7 @@ typedef enum dtrace_probespec {
 #define	DIF_OP_RLDX	77		/* rldx  [r1], rd */
 #define	DIF_OP_XLATE	78		/* xlate xlrindex, rd */
 #define	DIF_OP_XLARG	79		/* xlarg xlrindex, rd */
+#define	DIF_OP_STGA	80		/* stga var, ri, rd */
 
 #define	DIF_INTOFF_MAX		0xffff	/* highest integer table offset */
 #define	DIF_STROFF_MAX		0xffff	/* highest string table offset */
diff --git a/usr/src/uts/common/sys/dtrace_impl.h b/usr/src/uts/common/sys/dtrace_impl.h
index de7d0d25fb..a55d0d7e1f 100644
--- a/usr/src/uts/common/sys/dtrace_impl.h
+++ b/usr/src/uts/common/sys/dtrace_impl.h
@@ -25,7 +25,7 @@
  */
 
 /*
- * Copyright 2016 Joyent, Inc.
+ * Copyright 2017 Joyent, Inc.
  * Copyright (c) 2012 by Delphix. All rights reserved.
  */
 
@@ -1262,6 +1262,7 @@ extern void dtrace_copyoutstr(uintptr_t, uintptr_t, size_t,
     volatile uint16_t *);
 extern void dtrace_getpcstack(pc_t *, int, int, uint32_t *);
 extern ulong_t dtrace_getreg(struct regs *, uint_t);
+extern void dtrace_setreg(struct regs *, uint_t, ulong_t);
 extern uint64_t dtrace_getvmreg(uint_t, volatile uint16_t *);
 extern int dtrace_getstackdepth(int);
 extern void dtrace_getupcstack(uint64_t *, int);
diff --git a/usr/src/uts/intel/dtrace/dtrace_isa.c b/usr/src/uts/intel/dtrace/dtrace_isa.c
index f649473dfb..f52180b0c0 100644
--- a/usr/src/uts/intel/dtrace/dtrace_isa.c
+++ b/usr/src/uts/intel/dtrace/dtrace_isa.c
@@ -26,6 +26,7 @@
 
 /*
  * Copyright (c) 2013, 2014 by Delphix. All rights reserved.
+ * Copyright (c) 2017 Joyent, Inc.
  */
 
 #include <sys/dtrace_impl.h>
@@ -546,39 +547,42 @@ dtrace_getstackdepth(int aframes)
 	return (depth - aframes);
 }
 
+#if defined(__amd64)
+static const int dtrace_regmap[] = {
+	REG_GS,		/* GS */
+	REG_FS,		/* FS */
+	REG_ES,		/* ES */
+	REG_DS,		/* DS */
+	REG_RDI,	/* EDI */
+	REG_RSI,	/* ESI */
+	REG_RBP,	/* EBP */
+	REG_RSP,	/* ESP */
+	REG_RBX,	/* EBX */
+	REG_RDX,	/* EDX */
+	REG_RCX,	/* ECX */
+	REG_RAX,	/* EAX */
+	REG_TRAPNO,	/* TRAPNO */
+	REG_ERR,	/* ERR */
+	REG_RIP,	/* EIP */
+	REG_CS,		/* CS */
+	REG_RFL,	/* EFL */
+	REG_RSP,	/* UESP */
+	REG_SS		/* SS */
+};
+#endif
+
+
 ulong_t
 dtrace_getreg(struct regs *rp, uint_t reg)
 {
 #if defined(__amd64)
-	int regmap[] = {
-		REG_GS,		/* GS */
-		REG_FS,		/* FS */
-		REG_ES,		/* ES */
-		REG_DS,		/* DS */
-		REG_RDI,	/* EDI */
-		REG_RSI,	/* ESI */
-		REG_RBP,	/* EBP */
-		REG_RSP,	/* ESP */
-		REG_RBX,	/* EBX */
-		REG_RDX,	/* EDX */
-		REG_RCX,	/* ECX */
-		REG_RAX,	/* EAX */
-		REG_TRAPNO,	/* TRAPNO */
-		REG_ERR,	/* ERR */
-		REG_RIP,	/* EIP */
-		REG_CS,		/* CS */
-		REG_RFL,	/* EFL */
-		REG_RSP,	/* UESP */
-		REG_SS		/* SS */
-	};
-
 	if (reg <= SS) {
-		if (reg >= sizeof (regmap) / sizeof (int)) {
+		if (reg >= sizeof (dtrace_regmap) / sizeof (int)) {
 			DTRACE_CPUFLAG_SET(CPU_DTRACE_ILLOP);
 			return (0);
 		}
 
-		reg = regmap[reg];
+		reg = dtrace_regmap[reg];
 	} else {
 		reg -= SS + 1;
 	}
@@ -651,6 +655,105 @@ dtrace_getreg(struct regs *rp, uint_t reg)
 #endif
 }
 
+void
+dtrace_setreg(struct regs *rp, uint_t reg, ulong_t val)
+{
+#if defined(__amd64)
+	if (reg <= SS) {
+		ASSERT(reg < (sizeof (dtrace_regmap) / sizeof (int)));
+
+		reg = dtrace_regmap[reg];
+	} else {
+		reg -= SS + 1;
+	}
+
+	switch (reg) {
+	case REG_RDI:
+		rp->r_rdi = val;
+		break;
+	case REG_RSI:
+		rp->r_rsi = val;
+		break;
+	case REG_RDX:
+		rp->r_rdx = val;
+		break;
+	case REG_RCX:
+		rp->r_rcx = val;
+		break;
+	case REG_R8:
+		rp->r_r8 = val;
+		break;
+	case REG_R9:
+		rp->r_r9 = val;
+		break;
+	case REG_RAX:
+		rp->r_rax = val;
+		break;
+	case REG_RBX:
+		rp->r_rbx = val;
+		break;
+	case REG_RBP:
+		rp->r_rbp = val;
+		break;
+	case REG_R10:
+		rp->r_r10 = val;
+		break;
+	case REG_R11:
+		rp->r_r11 = val;
+		break;
+	case REG_R12:
+		rp->r_r12 = val;
+		break;
+	case REG_R13:
+		rp->r_r13 = val;
+		break;
+	case REG_R14:
+		rp->r_r14 = val;
+		break;
+	case REG_R15:
+		rp->r_r15 = val;
+		break;
+	case REG_RSP:
+		rp->r_rsp = val;
+		break;
+	default:
+		DTRACE_CPUFLAG_SET(CPU_DTRACE_ILLOP);
+		return;
+	}
+
+#else /* defined(__amd64) */
+	switch (reg) {
+	case EAX:
+		rp->r_eax = val;
+		break;
+	case ECX:
+		rp->r_ecx = val;
+		break;
+	case EDX:
+		rp->r_edx = val;
+		break;
+	case EBX:
+		rp->r_ebx = val;
+		break;
+	case ESP:
+		rp->r_esp = val;
+		break;
+	case EBP:
+		rp->r_ebp = val;
+		break;
+	case ESI:
+		rp->r_esi = val;
+		break;
+	case EDI:
+		rp->r_edi = val;
+		break;
+	default:
+		DTRACE_CPUFLAG_SET(CPU_DTRACE_ILLOP);
+		return;
+	}
+#endif /* defined(__amd64) */
+}
+
 static int
 dtrace_copycheck(uintptr_t uaddr, uintptr_t kaddr, size_t size)
 {
diff --git a/usr/src/uts/sparc/dtrace/dtrace_isa.c b/usr/src/uts/sparc/dtrace/dtrace_isa.c
index 20d9bc420b..541f733b06 100644
--- a/usr/src/uts/sparc/dtrace/dtrace_isa.c
+++ b/usr/src/uts/sparc/dtrace/dtrace_isa.c
@@ -24,7 +24,7 @@
  * Use is subject to license terms.
  */
 /*
- * Copyright 2011 Joyent, Inc.  All rights reserved.
+ * Copyright 2017 Joyent, Inc.
  */
 
 #include <sys/dtrace_impl.h>
@@ -940,6 +940,15 @@ got_fp:
 	return (value);
 }
 
+/* ARGSUSED */
+int
+dtrace_setreg(struct regs *rp, uint_t reg, ulong_t val)
+{
+	/* Not supported at this time */
+	DTRACE_CPUFLAG_SET(CPU_DTRACE_ILLOP);
+	return (0);
+}
+
 /*ARGSUSED*/
 uint64_t
 dtrace_getvmreg(uint_t ndx, volatile uint16_t *flags)