crypto/subtle: add DIT closure

Add a new function, WithDataIndependentTiming, which takes a function as
an argument, and encloses it with calls to set/unset the DIT PSTATE bit
on Arm64.

Since DIT is OS thread-local, for the duration of the execution of
WithDataIndependentTiming, we lock the goroutine to the OS thread, using
LockOSThread. For long running operations, this is likely to not be
performant, but we expect this to be tightly scoped around cryptographic
operations that have bounded execution times.

If locking to the OS thread turns out to be too slow, another option is
to add a bit to the g state indicating if a goroutine has DIT enabled,
and then have the scheduler enable/disable DIT when scheduling a g.

Additionally, we add a new GODEBUG, dataindependenttiming, which allows
setting DIT for an entire program. Running a program with
dataindependenttiming=1 enables DIT for the program during
initialization. In an ideal world PSTATE.DIT would be inherited from
the parent thread, so we'd only need to set it in the main thread and
then all subsequent threads would inherit the value. While this does
happen in the Linux kernel [0], it is not the case for darwin [1].
Rather than add complex logic to only set it on darwin for each new
thread, we just unconditionally set it in mstart1 and cgocallbackg1
regardless of the OS. DIT will already impose some overhead, and the
cost of setting the bit is only ~two instructions (CALL, MSR), so it
should be cheap enough.

Fixes #66450
Updates #49702

[0] https://github.com/torvalds/linux/blob/e8bdb3c8be08c9a3edc0a373c0aa8729355a0705/arch/arm64/kernel/process.c#L373
[1] https://github.com/apple-oss-distributions/xnu/blob/8d741a5de7ff4191bf97d57b9f54c2f6d4a15585/osfmk/arm64/status.c#L1666

Change-Id: I78eda691ff9254b0415f2b54770e5850a0179749
Reviewed-on: https://go-review.googlesource.com/c/go/+/598336
Reviewed-by: Michael Knyszek <mknyszek@google.com>
Reviewed-by: Filippo Valsorda <filippo@golang.org>
Reviewed-by: Michael Pratt <mpratt@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

9 files changed, 178 insertions, 0 deletions

diff --git a/src/crypto/subtle/dit.go b/src/crypto/subtle/dit.go
new file mode 100644
index 0000000000..c23df971f0
--- /dev/null
+++ b/src/crypto/subtle/dit.go
@@ -0,0 +1,50 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package subtle
+
+import (
+	"internal/runtime/sys"
+	"runtime"
+)
+
+// WithDataIndependentTiming enables architecture specific features which ensure
+// that the timing of specific instructions is independent of their inputs
+// before executing f. On f returning it disables these features.
+//
+// WithDataIndependentTiming should only be used when f is written to make use
+// of constant-time operations. WithDataIndependentTiming does not make
+// variable-time code constant-time.
+//
+// WithDataIndependentTiming may lock the current goroutine to the OS thread for
+// the duration of f. Calls to WithDataIndependentTiming may be nested.
+//
+// On Arm64 processors with FEAT_DIT, WithDataIndependentTiming enables
+// PSTATE.DIT. See https://developer.arm.com/documentation/ka005181/1-0/?lang=en.
+//
+// Currently, on all other architectures WithDataIndependentTiming executes f immediately
+// with no other side-effects.
+//
+//go:noinline
+func WithDataIndependentTiming(f func()) {
+	if !sys.DITSupported {
+		f()
+		return
+	}
+
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+
+	alreadyEnabled := sys.EnableDIT()
+
+	// disableDIT is called in a deferred function so that if f panics we will
+	// still disable DIT, in case the panic is recovered further up the stack.
+	defer func() {
+		if !alreadyEnabled {
+			sys.DisableDIT()
+		}
+	}()
+
+	f()
+}
diff --git a/src/crypto/subtle/dit_test.go b/src/crypto/subtle/dit_test.go
new file mode 100644
index 0000000000..8753ed623f
--- /dev/null
+++ b/src/crypto/subtle/dit_test.go
@@ -0,0 +1,61 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package subtle
+
+import (
+	"internal/cpu"
+	"internal/runtime/sys"
+	"testing"
+)
+
+func TestWithDataIndependentTiming(t *testing.T) {
+	if !cpu.ARM64.HasDIT {
+		t.Skip("CPU does not support DIT")
+	}
+
+	WithDataIndependentTiming(func() {
+		if !sys.DITEnabled() {
+			t.Fatal("dit not enabled within WithDataIndependentTiming closure")
+		}
+
+		WithDataIndependentTiming(func() {
+			if !sys.DITEnabled() {
+				t.Fatal("dit not enabled within nested WithDataIndependentTiming closure")
+			}
+		})
+
+		if !sys.DITEnabled() {
+			t.Fatal("dit not enabled after return from nested WithDataIndependentTiming closure")
+		}
+	})
+
+	if sys.DITEnabled() {
+		t.Fatal("dit not unset after returning from WithDataIndependentTiming closure")
+	}
+}
+
+func TestDITPanic(t *testing.T) {
+	if !cpu.ARM64.HasDIT {
+		t.Skip("CPU does not support DIT")
+	}
+
+	defer func() {
+		e := recover()
+		if e == nil {
+			t.Fatal("didn't panic")
+		}
+		if sys.DITEnabled() {
+			t.Error("DIT still enabled after panic inside of WithDataIndependentTiming closure")
+		}
+	}()
+
+	WithDataIndependentTiming(func() {
+		if !sys.DITEnabled() {
+			t.Fatal("dit not enabled within WithDataIndependentTiming closure")
+		}
+
+		panic("bad")
+	})
+}
diff --git a/src/internal/godebugs/table.go b/src/internal/godebugs/table.go
index d00014eaae..da6ca78773 100644
--- a/src/internal/godebugs/table.go
+++ b/src/internal/godebugs/table.go
@@ -26,6 +26,7 @@ type Info struct {
 // (Otherwise the test in this package will fail.)
 var All = []Info{
 	{Name: "asynctimerchan", Package: "time", Changed: 23, Old: "1"},
+	{Name: "dataindependenttiming", Package: "crypto/subtle", Opaque: true},
 	{Name: "execerrdot", Package: "os/exec"},
 	{Name: "gocachehash", Package: "cmd/go"},
 	{Name: "gocachetest", Package: "cmd/go"},
diff --git a/src/internal/runtime/sys/dit_arm64.go b/src/internal/runtime/sys/dit_arm64.go
new file mode 100644
index 0000000000..643fd770d5
--- /dev/null
+++ b/src/internal/runtime/sys/dit_arm64.go
@@ -0,0 +1,17 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build arm64
+
+package sys
+
+import (
+	"internal/cpu"
+)
+
+var DITSupported = cpu.ARM64.HasDIT
+
+func EnableDIT() bool
+func DITEnabled() bool
+func DisableDIT()
diff --git a/src/internal/runtime/sys/dit_arm64.s b/src/internal/runtime/sys/dit_arm64.s
new file mode 100644
index 0000000000..fcb44d6f22
--- /dev/null
+++ b/src/internal/runtime/sys/dit_arm64.s
@@ -0,0 +1,18 @@
+#include "textflag.h"
+
+TEXT ·EnableDIT(SB),$0-1
+    MRS DIT, R0
+    UBFX $24, R0, $1, R1
+    MOVB R1, ret+0(FP)
+    MSR $1, DIT
+    RET
+
+TEXT ·DITEnabled(SB),$0-1
+    MRS DIT, R0
+    UBFX $24, R0, $1, R1
+    MOVB R1, ret+0(FP)
+    RET
+
+TEXT ·DisableDIT(SB),$0
+    MSR $0, DIT
+    RET
diff --git a/src/internal/runtime/sys/no_dit.go b/src/internal/runtime/sys/no_dit.go
new file mode 100644
index 0000000000..0589d0ca14
--- /dev/null
+++ b/src/internal/runtime/sys/no_dit.go
@@ -0,0 +1,13 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build !arm64
+
+package sys
+
+var DITSupported = false
+
+func EnableDIT() bool  { return false }
+func DITEnabled() bool { return false }
+func DisableDIT()      {}
diff --git a/src/runtime/cgocall.go b/src/runtime/cgocall.go
index 0effcb8053..326674cd2e 100644
--- a/src/runtime/cgocall.go
+++ b/src/runtime/cgocall.go
@@ -425,6 +425,13 @@ func cgocallbackg1(fn, frame unsafe.Pointer, ctxt uintptr) {
 	restore := true
 	defer unwindm(&restore)
 
+	var ditAlreadySet bool
+	if debug.dataindependenttiming == 1 && gp.m.isextra {
+		// We only need to enable DIT for threads that were created by C, as it
+		// should already by enabled on threads that were created by Go.
+		ditAlreadySet = sys.EnableDIT()
+	}
+
 	if raceenabled {
 		raceacquire(unsafe.Pointer(&racecgosync))
 	}
@@ -440,6 +447,11 @@ func cgocallbackg1(fn, frame unsafe.Pointer, ctxt uintptr) {
 		racereleasemerge(unsafe.Pointer(&racecgosync))
 	}
 
+	if debug.dataindependenttiming == 1 && !ditAlreadySet {
+		// Only unset DIT if it wasn't already enabled when cgocallback was called.
+		sys.DisableDIT()
+	}
+
 	// Do not unwind m->g0->sched.sp.
 	// Our caller, cgocallback, will do that.
 	restore = false
diff --git a/src/runtime/proc.go b/src/runtime/proc.go
index 3f360ef129..17c375de1a 100644
--- a/src/runtime/proc.go
+++ b/src/runtime/proc.go
@@ -1848,6 +1848,10 @@ func mstart1() {
 		mstartm0()
 	}
 
+	if debug.dataindependenttiming == 1 {
+		sys.EnableDIT()
+	}
+
 	if fn := gp.m.mstartfn; fn != nil {
 		fn()
 	}
diff --git a/src/runtime/runtime1.go b/src/runtime/runtime1.go
index 56886ea571..7a092e8039 100644
--- a/src/runtime/runtime1.go
+++ b/src/runtime/runtime1.go
@@ -331,6 +331,7 @@ var debug struct {
 	traceadvanceperiod       int32
 	traceCheckStackOwnership int32
 	profstackdepth           int32
+	dataindependenttiming    int32
 
 	// debug.malloc is used as a combined debug check
 	// in the malloc function and should be set
@@ -367,6 +368,7 @@ var dbgvars = []*dbgVar{
 	{name: "asynctimerchan", atomic: &debug.asynctimerchan},
 	{name: "cgocheck", value: &debug.cgocheck},
 	{name: "clobberfree", value: &debug.clobberfree},
+	{name: "dataindependenttiming", value: &debug.dataindependenttiming},
 	{name: "disablethp", value: &debug.disablethp},
 	{name: "dontfreezetheworld", value: &debug.dontfreezetheworld},
 	{name: "efence", value: &debug.efence},