From 72dde793d88801aa4706d31d5b01028e92eddf60 Mon Sep 17 00:00:00 2001
From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
Date: Mon, 18 Nov 2024 22:55:24 +0530
Subject: [PATCH] Fix workspace extensions filer accidentally reading notebooks
 (#1891)

## Changes
The workspace extensions filer should not read or stat a notebook called
`foo` if the user calls `.Stat(ctx, "foo")`.

Instead, the filer should return a file not found error. This is because
the contract for the workspace extensions filer is to only work for
notebooks when the file path / name includes the extension (example:
`foo.ipynb` or `foo.sql` instead of just `foo`)

## Tests
Integration tests.
---
 internal/filer_test.go                        | 57 +++++++++++++++++++
 .../workspace_files_extensions_client.go      | 43 +++++++++++++-
 2 files changed, 98 insertions(+), 2 deletions(-)

diff --git a/internal/filer_test.go b/internal/filer_test.go
index 20207d34..a2760d91 100644
--- a/internal/filer_test.go
+++ b/internal/filer_test.go
@@ -723,6 +723,63 @@ func TestAccWorkspaceFilesExtensionsDirectoriesAreNotNotebooks(t *testing.T) {
 	assert.ErrorIs(t, err, fs.ErrNotExist)
 }
 
+func TestAccWorkspaceFilesExtensionsNotebooksAreNotReadAsFiles(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	wf, _ := setupWsfsExtensionsFiler(t)
+
+	// Create a notebook
+	err := wf.Write(ctx, "foo.ipynb", strings.NewReader(readFile(t, "testdata/notebooks/py1.ipynb")))
+	require.NoError(t, err)
+
+	// Reading foo should fail. Even though the WSFS name for the notebook is foo
+	// reading the notebook should only work with the .ipynb extension.
+	_, err = wf.Read(ctx, "foo")
+	assert.ErrorIs(t, err, fs.ErrNotExist)
+
+	_, err = wf.Read(ctx, "foo.ipynb")
+	assert.NoError(t, err)
+}
+
+func TestAccWorkspaceFilesExtensionsNotebooksAreNotStatAsFiles(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	wf, _ := setupWsfsExtensionsFiler(t)
+
+	// Create a notebook
+	err := wf.Write(ctx, "foo.ipynb", strings.NewReader(readFile(t, "testdata/notebooks/py1.ipynb")))
+	require.NoError(t, err)
+
+	// Stating foo should fail. Even though the WSFS name for the notebook is foo
+	// stating the notebook should only work with the .ipynb extension.
+	_, err = wf.Stat(ctx, "foo")
+	assert.ErrorIs(t, err, fs.ErrNotExist)
+
+	_, err = wf.Stat(ctx, "foo.ipynb")
+	assert.NoError(t, err)
+}
+
+func TestAccWorkspaceFilesExtensionsNotebooksAreNotDeletedAsFiles(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	wf, _ := setupWsfsExtensionsFiler(t)
+
+	// Create a notebook
+	err := wf.Write(ctx, "foo.ipynb", strings.NewReader(readFile(t, "testdata/notebooks/py1.ipynb")))
+	require.NoError(t, err)
+
+	// Deleting foo should fail. Even though the WSFS name for the notebook is foo
+	// deleting the notebook should only work with the .ipynb extension.
+	err = wf.Delete(ctx, "foo")
+	assert.ErrorIs(t, err, fs.ErrNotExist)
+
+	err = wf.Delete(ctx, "foo.ipynb")
+	assert.NoError(t, err)
+}
+
 func TestAccWorkspaceFilesExtensions_ExportFormatIsPreserved(t *testing.T) {
 	t.Parallel()
 
diff --git a/libs/filer/workspace_files_extensions_client.go b/libs/filer/workspace_files_extensions_client.go
index 53b77dd5..2a605209 100644
--- a/libs/filer/workspace_files_extensions_client.go
+++ b/libs/filer/workspace_files_extensions_client.go
@@ -244,6 +244,17 @@ func (w *workspaceFilesExtensionsClient) Write(ctx context.Context, name string,
 
 // Try to read the file as a regular file. If the file is not found, try to read it as a notebook.
 func (w *workspaceFilesExtensionsClient) Read(ctx context.Context, name string) (io.ReadCloser, error) {
+	// Ensure that the file / notebook exists. We do this check here to avoid reading
+	// the content of a notebook called `foo` when the user actually wanted
+	// to read the content of a file called `foo`.
+	//
+	// To read the content of a notebook called `foo` in the workspace the user
+	// should use the name with the extension included like `foo.ipynb` or `foo.sql`.
+	_, err := w.Stat(ctx, name)
+	if err != nil {
+		return nil, err
+	}
+
 	r, err := w.wsfs.Read(ctx, name)
 
 	// If the file is not found, it might be a notebook.
@@ -276,7 +287,18 @@ func (w *workspaceFilesExtensionsClient) Delete(ctx context.Context, name string
 		return ReadOnlyError{"delete"}
 	}
 
-	err := w.wsfs.Delete(ctx, name, mode...)
+	// Ensure that the file / notebook exists. We do this check here to avoid
+	// deleting the a notebook called `foo` when the user actually wanted to
+	// delete a file called `foo`.
+	//
+	// To delete a notebook called `foo` in the workspace the user should use the
+	// name with the extension included like `foo.ipynb` or `foo.sql`.
+	_, err := w.Stat(ctx, name)
+	if err != nil {
+		return err
+	}
+
+	err = w.wsfs.Delete(ctx, name, mode...)
 
 	// If the file is not found, it might be a notebook.
 	if errors.As(err, &FileDoesNotExistError{}) {
@@ -315,7 +337,24 @@ func (w *workspaceFilesExtensionsClient) Stat(ctx context.Context, name string)
 		return wsfsFileInfo{ObjectInfo: stat.ObjectInfo}, nil
 	}
 
-	return info, err
+	if err != nil {
+		return nil, err
+	}
+
+	// If an object is found and it is a notebook, return a FileDoesNotExistError.
+	// If a notebook is found by the workspace files client, without having stripped
+	// the extension, this implies that no file with the same name exists.
+	//
+	// This check is done to avoid returning the stat for a notebook called `foo`
+	// when the user actually wanted to stat a file called `foo`.
+	//
+	// To stat the metadata of a notebook called `foo` in the workspace the user
+	// should use the name with the extension included like `foo.ipynb` or `foo.sql`.
+	if info.Sys().(workspace.ObjectInfo).ObjectType == workspace.ObjectTypeNotebook {
+		return nil, FileDoesNotExistError{name}
+	}
+
+	return info, nil
 }
 
 // Note: The import API returns opaque internal errors for namespace clashes