From c8a50763f27f7f7f60964a9db5cecad16d6f2193 Mon Sep 17 00:00:00 2001
From: Jing Chen <chjing@google.com>
Date: Tue, 12 Mar 2024 00:23:42 -0700
Subject: [PATCH] Add unimplemented methods for tpuFd from
 vfs.FileDescriptionImpl.

PiperOrigin-RevId: 614936899
---
 pkg/sentry/devices/tpuproxy/device.go |  4 ++++
 pkg/sentry/devices/tpuproxy/tpu.go    | 19 ++++++++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/pkg/sentry/devices/tpuproxy/device.go b/pkg/sentry/devices/tpuproxy/device.go
index 4ba5b575d..59c819a42 100644
--- a/pkg/sentry/devices/tpuproxy/device.go
+++ b/pkg/sentry/devices/tpuproxy/device.go
@@ -74,6 +74,10 @@ func (dev *tpuDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, o
 		unix.Close(hostFD)
 		return nil, err
 	}
+	if err := fdnotifier.AddFD(int32(hostFD), &fd.queue); err != nil {
+		unix.Close(hostFD)
+		return nil, err
+	}
 	return &fd.vfsfd, nil
 }
 
diff --git a/pkg/sentry/devices/tpuproxy/tpu.go b/pkg/sentry/devices/tpuproxy/tpu.go
index e806d0ac9..51c4b6734 100644
--- a/pkg/sentry/devices/tpuproxy/tpu.go
+++ b/pkg/sentry/devices/tpuproxy/tpu.go
@@ -16,8 +16,12 @@
 package tpuproxy
 
 import (
+	"fmt"
+
+	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/errors/linuxerr"
+	"gvisor.dev/gvisor/pkg/fdnotifier"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 	"gvisor.dev/gvisor/pkg/usermem"
@@ -35,24 +39,37 @@ type tpuFD struct {
 
 	hostFD int32
 	device *tpuDevice
+	queue  waiter.Queue
 }
 
 // Release implements vfs.FileDescriptionImpl.Release.
 func (fd *tpuFD) Release(context.Context) {
+	fdnotifier.RemoveFD(fd.hostFD)
+	fd.queue.Notify(waiter.EventHUp)
+	unix.Close(int(fd.hostFD))
 }
 
 // EventRegister implements waiter.Waitable.EventRegister.
 func (fd *tpuFD) EventRegister(e *waiter.Entry) error {
+	fd.queue.EventRegister(e)
+	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
+		fd.queue.EventUnregister(e)
+		return err
+	}
 	return nil
 }
 
 // EventUnregister implements waiter.Waitable.EventUnregister.
 func (fd *tpuFD) EventUnregister(e *waiter.Entry) {
+	fd.queue.EventUnregister(e)
+	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
+		panic(fmt.Sprint("UpdateFD:", err))
+	}
 }
 
 // Readiness implements waiter.Waitable.Readiness.
 func (fd *tpuFD) Readiness(mask waiter.EventMask) waiter.EventMask {
-	return waiter.EventErr
+	return fdnotifier.NonBlockingPoll(fd.hostFD, mask)
 }
 
 // Epollable implements vfs.FileDescriptionImpl.Epollable.