internal/sockstate: Restore connection checking on Linux. (#2722)

GMS server handler is supposed to cancel running queries if the connection
which issued them goes away. It does this by checking the connection state
out-of-band anytime the query is running and canceling theh query if the
connection goes away. The connection checking code is platform-specific and
currently only works on Linux.

In commit 538696b I introduced a bug where the
connection checking code tries to inspect the socket state of an already closed
file descriptor. This change fixes the behavior so that the file descriptor is
left open until the necessary socket state is extracted.

Author: reltuk

internal/sockstate/netstat_darwin.go

@@ -32,9 +32,10 @@ func tcpSocks(accept AcceptFn) ([]sockTabEntry, error) {
 }
 
 func GetConnInode(c *net.TCPConn) (uint64, error) {
-	_, err := getConnFd(c)
+	_, finalize, err := getConnFd(c)
 	if err != nil {
 		return 0, err
 	}
+	defer finalize()
 	return 0, ErrSocketCheckNotImplemented.New()
 }

internal/sockstate/netstat_freebsd.go

@@ -32,9 +32,10 @@ func tcpSocks(accept AcceptFn) ([]sockTabEntry, error) {
 }
 
 func GetConnInode(c *net.TCPConn) (uint64, error) {
-	_, err := getConnFd(c)
+	_, finalize, err := getConnFd(c)
 	if err != nil {
 		return 0, err
 	}
+	defer finalize()
 	return 0, ErrSocketCheckNotImplemented.New()
 }

internal/sockstate/netstat_linux.go

@@ -279,10 +279,11 @@ func tcpSocks(accept AcceptFn) ([]sockTabEntry, error) {
 
 // GetConnInode returns the inode number of an fd.
 func GetConnInode(conn *net.TCPConn) (n uint64, err error) {
-	fd, err := getConnFd(conn)
+	fd, finalize, err := getConnFd(conn)
 	if err != nil {
 		return 0, err
 	}
+	defer finalize()
 
 	if isWSL || isProcBlocked {
 		return 0, ErrSocketCheckNotImplemented.New()

internal/sockstate/netstat_unix.go

@@ -22,10 +22,10 @@ import (
 	"syscall"
 )
 
-func getConnFd(c *net.TCPConn) (fd uintptr, err error) {
+func getConnFd(c *net.TCPConn) (fd uintptr, finalize func() error, err error) {
 	f, err := c.File()
 	if err != nil {
-		return 0, err
+		return 0, nil, err
 	}
 
 	fd = f.Fd()
@@ -34,7 +34,5 @@ func getConnFd(c *net.TCPConn) (fd uintptr, err error) {
 	// blocking Close() in some cases.
 	syscall.SetNonblock(int(fd), true)
 
-	f.Close()
-
-	return fd, nil
+	return fd, f.Close, nil
 }