Fix kcov enabling and disabling procedures.

- When the KCOV_ENABLE_TRACE ioctl is called with the trace kind KCOV_TRACE_PC,
  the kcov mode should be set to KCOV_*MODE*_TRACE_PC.
- When the owning task of kcov exits, the memory mapping should not be cleared
  so it can be used by other tasks.
- Add more tests (also tested on native Linux kcov).

PiperOrigin-RevId: 335202585

5 files changed, 151 insertions, 27 deletions

diff --git a/pkg/sentry/fsimpl/sys/kcov.go b/pkg/sentry/fsimpl/sys/kcov.go
index b75d70ae6..1a6749e53 100644
--- a/pkg/sentry/fsimpl/sys/kcov.go
+++ b/pkg/sentry/fsimpl/sys/kcov.go
@@ -104,7 +104,7 @@ func (fd *kcovFD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) erro
 func (fd *kcovFD) Release(ctx context.Context) {
 	// kcov instances have reference counts in Linux, but this seems sufficient
 	// for our purposes.
-	fd.kcov.Reset()
+	fd.kcov.Clear()
 }
 
 // SetStat implements vfs.FileDescriptionImpl.SetStat.
diff --git a/pkg/sentry/kernel/kcov.go b/pkg/sentry/kernel/kcov.go
index d3e76ca7b..90ceb5ef5 100644
--- a/pkg/sentry/kernel/kcov.go
+++ b/pkg/sentry/kernel/kcov.go
@@ -89,7 +89,7 @@ func (kcov *Kcov) TaskWork(t *Task) {
 	kcov.mu.Lock()
 	defer kcov.mu.Unlock()
 
-	if kcov.mode != linux.KCOV_TRACE_PC {
+	if kcov.mode != linux.KCOV_MODE_TRACE_PC {
 		return
 	}
 
@@ -146,7 +146,7 @@ func (kcov *Kcov) InitTrace(size uint64) error {
 }
 
 // EnableTrace performs the KCOV_ENABLE_TRACE ioctl.
-func (kcov *Kcov) EnableTrace(ctx context.Context, traceMode uint8) error {
+func (kcov *Kcov) EnableTrace(ctx context.Context, traceKind uint8) error {
 	t := TaskFromContext(ctx)
 	if t == nil {
 		panic("kcovInode.EnableTrace() cannot be used outside of a task goroutine")
@@ -160,9 +160,9 @@ func (kcov *Kcov) EnableTrace(ctx context.Context, traceMode uint8) error {
 		return syserror.EINVAL
 	}
 
-	switch traceMode {
+	switch traceKind {
 	case linux.KCOV_TRACE_PC:
-		kcov.mode = traceMode
+		kcov.mode = linux.KCOV_MODE_TRACE_PC
 	case linux.KCOV_TRACE_CMP:
 		// We do not support KCOV_MODE_TRACE_CMP.
 		return syserror.ENOTSUP
@@ -175,6 +175,7 @@ func (kcov *Kcov) EnableTrace(ctx context.Context, traceMode uint8) error {
 	}
 
 	kcov.owningTask = t
+	t.SetKcov(kcov)
 	t.RegisterWork(kcov)
 
 	// Clear existing coverage data; the task expects to read only coverage data
@@ -196,28 +197,39 @@ func (kcov *Kcov) DisableTrace(ctx context.Context) error {
 	if t != kcov.owningTask {
 		return syserror.EINVAL
 	}
-	kcov.owningTask = nil
 	kcov.mode = linux.KCOV_MODE_INIT
-	kcov.resetLocked()
+	kcov.owningTask = nil
+	kcov.mappable = nil
 	return nil
 }
 
-// Reset is called when the owning task exits.
-func (kcov *Kcov) Reset() {
+// Clear resets the mode and clears the owning task and memory mapping for kcov.
+// It is called when the fd corresponding to kcov is closed. Note that the mode
+// needs to be set so that the next call to kcov.TaskWork() will exit early.
+func (kcov *Kcov) Clear() {
 	kcov.mu.Lock()
-	kcov.resetLocked()
+	kcov.clearLocked()
 	kcov.mu.Unlock()
 }
 
-// The kcov instance is reset when the owning task exits or when tracing is
-// disabled.
-func (kcov *Kcov) resetLocked() {
+func (kcov *Kcov) clearLocked() {
+	kcov.mode = linux.KCOV_MODE_INIT
 	kcov.owningTask = nil
 	if kcov.mappable != nil {
 		kcov.mappable = nil
 	}
 }
 
+// OnTaskExit is called when the owning task exits. It is similar to
+// kcov.Clear(), except the memory mapping is not cleared, so that the same
+// mapping can be used in the future if kcov is enabled again by another task.
+func (kcov *Kcov) OnTaskExit() {
+	kcov.mu.Lock()
+	kcov.mode = linux.KCOV_MODE_INIT
+	kcov.owningTask = nil
+	kcov.mu.Unlock()
+}
+
 // ConfigureMMap is called by the vfs.FileDescription for this kcov instance to
 // implement vfs.FileDescription.ConfigureMMap.
 func (kcov *Kcov) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
diff --git a/pkg/sentry/kernel/task.go b/pkg/sentry/kernel/task.go
index a436610c9..f796e0fa3 100644
--- a/pkg/sentry/kernel/task.go
+++ b/pkg/sentry/kernel/task.go
@@ -917,7 +917,7 @@ func (t *Task) SetKcov(k *Kcov) {
 // ResetKcov clears the kcov instance associated with t.
 func (t *Task) ResetKcov() {
 	if t.kcov != nil {
-		t.kcov.Reset()
+		t.kcov.OnTaskExit()
 		t.kcov = nil
 	}
 }
diff --git a/test/syscalls/linux/BUILD b/test/syscalls/linux/BUILD
index d9dbe2267..6a2ec9787 100644
--- a/test/syscalls/linux/BUILD
+++ b/test/syscalls/linux/BUILD
@@ -1079,6 +1079,7 @@ cc_binary(
         gtest,
         "//test/util:test_main",
         "//test/util:test_util",
+        "//test/util:thread_util",
     ],
 )
 
diff --git a/test/syscalls/linux/kcov.cc b/test/syscalls/linux/kcov.cc
index 6afcb4e75..6816c1fd0 100644
--- a/test/syscalls/linux/kcov.cc
+++ b/test/syscalls/linux/kcov.cc
@@ -16,39 +16,47 @@
 #include <sys/ioctl.h>
 #include <sys/mman.h>
 
+#include <atomic>
+
 #include "gtest/gtest.h"
 #include "test/util/capability_util.h"
 #include "test/util/file_descriptor.h"
 #include "test/util/test_util.h"
+#include "test/util/thread_util.h"
 
 namespace gvisor {
 namespace testing {
 
 namespace {
 
-// For this test to work properly, it must be run with coverage enabled. On
+// For this set of tests to run, they must be run with coverage enabled. On
 // native Linux, this involves compiling the kernel with kcov enabled. For
-// gVisor, we need to enable the Go coverage tool, e.g.
-// bazel test --collect_coverage_data --instrumentation_filter=//pkg/... <test>.
+// gVisor, we need to enable the Go coverage tool, e.g. bazel test --
+// collect_coverage_data --instrumentation_filter=//pkg/... <test>.
+
+constexpr char kcovPath[] = "/sys/kernel/debug/kcov";
+constexpr int kSize = 4096;
+constexpr int KCOV_INIT_TRACE = 0x80086301;
+constexpr int KCOV_ENABLE = 0x6364;
+constexpr int KCOV_DISABLE = 0x6365;
+
+uint64_t* KcovMmap(int fd) {
+  return (uint64_t*)mmap(nullptr, kSize * sizeof(uint64_t),
+                         PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+}
+
 TEST(KcovTest, Kcov) {
   SKIP_IF(!ASSERT_NO_ERRNO_AND_VALUE(HaveCapability((CAP_DAC_OVERRIDE))));
 
-  constexpr int kSize = 4096;
-  constexpr int KCOV_INIT_TRACE = 0x80086301;
-  constexpr int KCOV_ENABLE = 0x6364;
-  constexpr int KCOV_DISABLE = 0x6365;
-
   int fd;
-  ASSERT_THAT(fd = open("/sys/kernel/debug/kcov", O_RDWR),
+  ASSERT_THAT(fd = open(kcovPath, O_RDWR),
               AnyOf(SyscallSucceeds(), SyscallFailsWithErrno(ENOENT)));
-
   // Kcov not available.
   SKIP_IF(errno == ENOENT);
+  auto fd_closer = Cleanup([fd]() { close(fd); });
 
   ASSERT_THAT(ioctl(fd, KCOV_INIT_TRACE, kSize), SyscallSucceeds());
-  uint64_t* area = (uint64_t*)mmap(nullptr, kSize * sizeof(uint64_t),
-                                   PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
-
+  uint64_t* area = KcovMmap(fd);
   ASSERT_TRUE(area != MAP_FAILED);
   ASSERT_THAT(ioctl(fd, KCOV_ENABLE, 0), SyscallSucceeds());
 
@@ -67,6 +75,109 @@ TEST(KcovTest, Kcov) {
   ASSERT_THAT(ioctl(fd, KCOV_DISABLE, 0), SyscallSucceeds());
 }
 
+TEST(KcovTest, PrematureMmap) {
+  SKIP_IF(!ASSERT_NO_ERRNO_AND_VALUE(HaveCapability((CAP_DAC_OVERRIDE))));
+
+  int fd;
+  ASSERT_THAT(fd = open(kcovPath, O_RDWR),
+              AnyOf(SyscallSucceeds(), SyscallFailsWithErrno(ENOENT)));
+  // Kcov not available.
+  SKIP_IF(errno == ENOENT);
+  auto fd_closer = Cleanup([fd]() { close(fd); });
+
+  // Cannot mmap before KCOV_INIT_TRACE.
+  uint64_t* area = KcovMmap(fd);
+  ASSERT_TRUE(area == MAP_FAILED);
+}
+
+// Tests that multiple kcov fds can be used simultaneously.
+TEST(KcovTest, MultipleFds) {
+  SKIP_IF(!ASSERT_NO_ERRNO_AND_VALUE(HaveCapability((CAP_DAC_OVERRIDE))));
+
+  int fd1;
+  ASSERT_THAT(fd1 = open(kcovPath, O_RDWR),
+              AnyOf(SyscallSucceeds(), SyscallFailsWithErrno(ENOENT)));
+  // Kcov not available.
+  SKIP_IF(errno == ENOENT);
+
+  int fd2;
+  ASSERT_THAT(fd2 = open(kcovPath, O_RDWR), SyscallSucceeds());
+  auto fd_closer = Cleanup([fd1, fd2]() {
+    close(fd1);
+    close(fd2);
+  });
+
+  auto t1 = ScopedThread([&] {
+    ASSERT_THAT(ioctl(fd1, KCOV_INIT_TRACE, kSize), SyscallSucceeds());
+    uint64_t* area = KcovMmap(fd1);
+    ASSERT_TRUE(area != MAP_FAILED);
+    ASSERT_THAT(ioctl(fd1, KCOV_ENABLE, 0), SyscallSucceeds());
+  });
+
+  ASSERT_THAT(ioctl(fd2, KCOV_INIT_TRACE, kSize), SyscallSucceeds());
+  uint64_t* area = KcovMmap(fd2);
+  ASSERT_TRUE(area != MAP_FAILED);
+  ASSERT_THAT(ioctl(fd2, KCOV_ENABLE, 0), SyscallSucceeds());
+}
+
+// Tests behavior for two threads trying to use the same kcov fd.
+TEST(KcovTest, MultipleThreads) {
+  SKIP_IF(!ASSERT_NO_ERRNO_AND_VALUE(HaveCapability((CAP_DAC_OVERRIDE))));
+
+  int fd;
+  ASSERT_THAT(fd = open(kcovPath, O_RDWR),
+              AnyOf(SyscallSucceeds(), SyscallFailsWithErrno(ENOENT)));
+  // Kcov not available.
+  SKIP_IF(errno == ENOENT);
+  auto fd_closer = Cleanup([fd]() { close(fd); });
+
+  // Test the behavior of multiple threads trying to use the same kcov fd
+  // simultaneously.
+  std::atomic<bool> t1_enabled(false), t1_disabled(false), t2_failed(false),
+      t2_exited(false);
+  auto t1 = ScopedThread([&] {
+    ASSERT_THAT(ioctl(fd, KCOV_INIT_TRACE, kSize), SyscallSucceeds());
+    uint64_t* area = KcovMmap(fd);
+    ASSERT_TRUE(area != MAP_FAILED);
+    ASSERT_THAT(ioctl(fd, KCOV_ENABLE, 0), SyscallSucceeds());
+    t1_enabled = true;
+
+    // After t2 has made sure that enabling kcov again fails, disable it.
+    while (!t2_failed) {
+      sched_yield();
+    }
+    ASSERT_THAT(ioctl(fd, KCOV_DISABLE, 0), SyscallSucceeds());
+    t1_disabled = true;
+
+    // Wait for t2 to enable kcov and then exit, after which we should be able
+    // to enable kcov again, without needing to set up a new memory mapping.
+    while (!t2_exited) {
+      sched_yield();
+    }
+    ASSERT_THAT(ioctl(fd, KCOV_ENABLE, 0), SyscallSucceeds());
+  });
+
+  auto t2 = ScopedThread([&] {
+    // Wait for t1 to enable kcov, and make sure that enabling kcov again fails.
+    while (!t1_enabled) {
+      sched_yield();
+    }
+    ASSERT_THAT(ioctl(fd, KCOV_ENABLE, 0), SyscallFailsWithErrno(EINVAL));
+    t2_failed = true;
+
+    // Wait for t1 to disable kcov, after which using fd should now succeed.
+    while (!t1_disabled) {
+      sched_yield();
+    }
+    uint64_t* area = KcovMmap(fd);
+    ASSERT_TRUE(area != MAP_FAILED);
+    ASSERT_THAT(ioctl(fd, KCOV_ENABLE, 0), SyscallSucceeds());
+  });
+
+  t2.Join();
+  t2_exited = true;
+}
+
 }  // namespace
 
 }  // namespace testing