6 files changed, 221 insertions, 0 deletions
diff --git a/pkg/procid/BUILD b/pkg/procid/BUILD
new file mode 100644
index 000000000..aa3e3ac0b
--- /dev/null
+++ b/pkg/procid/BUILD
@@ -0,0 +1,34 @@
+load("//tools:defs.bzl", "go_library", "go_test")
+
+package(licenses = ["notice"])
+
+go_library(
+    name = "procid",
+    srcs = [
+        "procid.go",
+        "procid_amd64.s",
+        "procid_arm64.s",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+go_test(
+    name = "procid_test",
+    size = "small",
+    srcs = [
+        "procid_test.go",
+    ],
+    library = ":procid",
+    deps = ["//pkg/sync"],
+)
+
+go_test(
+    name = "procid_net_test",
+    size = "small",
+    srcs = [
+        "procid_net_test.go",
+        "procid_test.go",
+    ],
+    library = ":procid",
+    deps = ["//pkg/sync"],
+)
diff --git a/pkg/procid/procid.go b/pkg/procid/procid.go
new file mode 100644
index 000000000..78b92422c
--- /dev/null
+++ b/pkg/procid/procid.go
@@ -0,0 +1,21 @@
+// Copyright 2018 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package procid provides a way to get the current system thread identifier.
+package procid
+
+// Current returns the current system thread identifier.
+//
+// Precondition: This should only be called with the runtime OS thread locked.
+func Current() uint64
diff --git a/pkg/procid/procid_amd64.s b/pkg/procid/procid_amd64.s
new file mode 100644
index 000000000..7c622e5d7
--- /dev/null
+++ b/pkg/procid/procid_amd64.s
@@ -0,0 +1,30 @@
+// Copyright 2018 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// +build amd64
+// +build go1.8
+// +build !go1.16
+
+#include "textflag.h"
+
+TEXT ·Current(SB),NOSPLIT,$0-8
+	// The offset specified here is the m_procid offset for Go1.8+.
+	// Changes to this offset should be caught by the tests, and major
+	// version changes require an explicit tag change above.
+	MOVQ TLS, AX
+	MOVQ 0(AX)(TLS*1), AX
+	MOVQ 48(AX), AX // g_m (may change in future versions)
+	MOVQ 72(AX), AX // m_procid (may change in future versions)
+	MOVQ AX, ret+0(FP)
+	RET
diff --git a/pkg/procid/procid_arm64.s b/pkg/procid/procid_arm64.s
new file mode 100644
index 000000000..48ebb5fd1
--- /dev/null
+++ b/pkg/procid/procid_arm64.s
@@ -0,0 +1,29 @@
+// Copyright 2018 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// +build arm64
+// +build go1.8
+// +build !go1.16
+
+#include "textflag.h"
+
+TEXT ·Current(SB),NOSPLIT,$0-8
+	// The offset specified here is the m_procid offset for Go1.8+.
+	// Changes to this offset should be caught by the tests, and major
+	// version changes require an explicit tag change above.
+	MOVD g, R0      // g
+	MOVD 48(R0), R0 // g_m (may change in future versions)
+	MOVD 72(R0), R0 // m_procid (may change in future versions)
+	MOVD R0, ret+0(FP)
+	RET
diff --git a/pkg/procid/procid_net_test.go b/pkg/procid/procid_net_test.go
new file mode 100644
index 000000000..b628e2285
--- /dev/null
+++ b/pkg/procid/procid_net_test.go
@@ -0,0 +1,21 @@
+// Copyright 2018 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package procid
+
+// This file is just to force the inclusion of the "net" package, which will
+// make the test binary a cgo one.
+import (
+	_ "net"
+)
diff --git a/pkg/procid/procid_test.go b/pkg/procid/procid_test.go
new file mode 100644
index 000000000..9ec08c3d6
--- /dev/null
+++ b/pkg/procid/procid_test.go
@@ -0,0 +1,86 @@
+// Copyright 2018 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package procid
+
+import (
+	"os"
+	"runtime"
+	"syscall"
+	"testing"
+
+	"gvisor.dev/gvisor/pkg/sync"
+)
+
+// runOnMain is used to send functions to run on the main (initial) thread.
+var runOnMain = make(chan func(), 10)
+
+func checkProcid(t *testing.T, start *sync.WaitGroup, done *sync.WaitGroup) {
+	defer done.Done()
+
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+
+	start.Done()
+	start.Wait()
+
+	procID := Current()
+	tid := syscall.Gettid()
+
+	if procID != uint64(tid) {
+		t.Logf("Bad procid: expected %v, got %v", tid, procID)
+		t.Fail()
+	}
+}
+
+func TestProcidInitialized(t *testing.T) {
+	var start sync.WaitGroup
+	var done sync.WaitGroup
+
+	count := 100
+	start.Add(count + 1)
+	done.Add(count + 1)
+
+	// Run the check on the main thread.
+	//
+	// When cgo is not included, the only case when procid isn't initialized
+	// is in the main (initial) thread, so we have to test this case
+	// specifically.
+	runOnMain <- func() {
+		checkProcid(t, &start, &done)
+	}
+
+	// Run the check on a number of different threads.
+	for i := 0; i < count; i++ {
+		go checkProcid(t, &start, &done)
+	}
+
+	done.Wait()
+}
+
+func TestMain(m *testing.M) {
+	// Make sure we remain at the main (initial) thread.
+	runtime.LockOSThread()
+
+	// Start running tests in a different goroutine.
+	go func() {
+		os.Exit(m.Run())
+	}()
+
+	// Execute any functions that have been sent for execution in the main
+	// thread.
+	for f := range runOnMain {
+		f()
+	}
+}