3 files changed, 30 insertions, 2 deletions
diff --git a/pkg/sentry/arch/arch.go b/pkg/sentry/arch/arch.go
index 575b7ba66..4cd7a9af5 100644
--- a/pkg/sentry/arch/arch.go
+++ b/pkg/sentry/arch/arch.go
@@ -115,6 +115,12 @@ type Context interface {
 	// SetStack sets the current stack pointer.
 	SetStack(value uintptr)
 
+	// TLS returns the current TLS pointer.
+	TLS() uintptr
+
+	// SetTLS sets the current TLS pointer. Returns false if value is invalid.
+	SetTLS(value uintptr) bool
+
 	// SetRSEQInterruptedIP sets the register that contains the old IP when a
 	// restartable sequence is interrupted.
 	SetRSEQInterruptedIP(value uintptr)
diff --git a/pkg/sentry/arch/arch_amd64.go b/pkg/sentry/arch/arch_amd64.go
index bb80a7bed..2507774f7 100644
--- a/pkg/sentry/arch/arch_amd64.go
+++ b/pkg/sentry/arch/arch_amd64.go
@@ -158,6 +158,22 @@ func (c *context64) SetStack(value uintptr) {
 	c.Regs.Rsp = uint64(value)
 }
 
+// TLS returns the current TLS pointer.
+func (c *context64) TLS() uintptr {
+	return uintptr(c.Regs.Fs_base)
+}
+
+// SetTLS sets the current TLS pointer. Returns false if value is invalid.
+func (c *context64) SetTLS(value uintptr) bool {
+	if !isValidSegmentBase(uint64(value)) {
+		return false
+	}
+
+	c.Regs.Fs = 0
+	c.Regs.Fs_base = uint64(value)
+	return true
+}
+
 // SetRSEQInterruptedIP implements Context.SetRSEQInterruptedIP.
 func (c *context64) SetRSEQInterruptedIP(value uintptr) {
 	c.Regs.R10 = uint64(value)
diff --git a/pkg/sentry/arch/arch_x86.go b/pkg/sentry/arch/arch_x86.go
index 59bf89d99..e50a76083 100644
--- a/pkg/sentry/arch/arch_x86.go
+++ b/pkg/sentry/arch/arch_x86.go
@@ -353,10 +353,10 @@ func (s *State) PtraceSetRegs(src io.Reader) (int, error) {
 	if !isUserSegmentSelector(regs.Ss) {
 		return 0, syscall.EIO
 	}
-	if regs.Fs_base >= uint64(maxAddr64) {
+	if !isValidSegmentBase(regs.Fs_base) {
 		return 0, syscall.EIO
 	}
-	if regs.Gs_base >= uint64(maxAddr64) {
+	if !isValidSegmentBase(regs.Gs_base) {
 		return 0, syscall.EIO
 	}
 	// CS and SS are validated, but changes to them are otherwise silently
@@ -389,6 +389,12 @@ func isUserSegmentSelector(reg uint64) bool {
 	return reg&3 == 3
 }
 
+// isValidSegmentBase returns true if the given segment base specifies a
+// canonical user address.
+func isValidSegmentBase(reg uint64) bool {
+	return reg < uint64(maxAddr64)
+}
+
 // ptraceFPRegsSize is the size in bytes of Linux's user_i387_struct, the type
 // manipulated by PTRACE_GETFPREGS and PTRACE_SETFPREGS on x86. Equivalently,
 // ptraceFPRegsSize is the size in bytes of the x86 FXSAVE area.