From b1442b990769ba433a6e8cad8d7dfd7d87a940c8 Mon Sep 17 00:00:00 2001
From: Tianyu Yao <skyyao@fb.com>
Date: Tue, 17 Mar 2020 15:09:18 -0700
Subject: [PATCH] Disallow user defined object to be used as key in useFragment

Reviewed By: jstejada

Differential Revision: D20497252

fbshipit-source-id: 5d7cfeb181480710fb32c9cb3d8ab6db0c979131
---
 .../__flowtests__/useFragment-flowtest.js     | 10 +++++++
 .../relay-experimental/__flowtests__/utils.js |  5 ++++
 .../__tests__/useFragment-test.js             |  1 +
 packages/relay-experimental/useFragment.js    | 28 +++++++++++++++----
 4 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/packages/relay-experimental/__flowtests__/useFragment-flowtest.js b/packages/relay-experimental/__flowtests__/useFragment-flowtest.js
index cd4cc3c9280aa..c33b30f45b026 100644
--- a/packages/relay-experimental/__flowtests__/useFragment-flowtest.js
+++ b/packages/relay-experimental/__flowtests__/useFragment-flowtest.js
@@ -23,6 +23,7 @@ import {
   keyNonNullablePlural,
   keyNullable,
   keyNullablePlural,
+  fragmentData,
 } from './utils';
 import type {
   NonNullableData,
@@ -52,4 +53,13 @@ import type {
 // $FlowExpectedError: Key should be one of the generated types
 (useFragment(fragmentInput, 'INVALID_KEY'): NullableData);
 
+// $FlowExpectedError: Key should not be a user provided object
+(useFragment(fragmentInput, {a: 123}): NullableData);
+
+// $FlowExpectedError: Key should not be an empty object
+(useFragment(fragmentInput, {}): NullableData);
+
+// $FlowExpectedError: Key should be the `<name>$key` type from generated flow
+(useFragment(fragmentInput, fragmentData): NullableData);
+
 /* eslint-enable react-hooks/rules-of-hooks */
diff --git a/packages/relay-experimental/__flowtests__/utils.js b/packages/relay-experimental/__flowtests__/utils.js
index 4a061103befff..a4d5ef086abb5 100644
--- a/packages/relay-experimental/__flowtests__/utils.js
+++ b/packages/relay-experimental/__flowtests__/utils.js
@@ -72,6 +72,11 @@ declare export var keyAnotherNullable: ?{
   ...
 };
 
+declare export var fragmentData: {
+  +$refType: FragmentReference,
+  ...
+};
+
 export type QueryOperation = {|
   +variables: QueryVariables,
   +response: {...},
diff --git a/packages/relay-experimental/__tests__/useFragment-test.js b/packages/relay-experimental/__tests__/useFragment-test.js
index 5094f877dda55..e1283bb2388d1 100644
--- a/packages/relay-experimental/__tests__/useFragment-test.js
+++ b/packages/relay-experimental/__tests__/useFragment-test.js
@@ -46,6 +46,7 @@ describe('useFragment', () => {
   let ContextProvider;
 
   function useFragment(fragmentNode, fragmentRef) {
+    // $FlowFixMe non-generated fragmentRef is disallowd
     const data = useFragmentOriginal(fragmentNode, fragmentRef);
     renderSpy(data);
     return data;
diff --git a/packages/relay-experimental/useFragment.js b/packages/relay-experimental/useFragment.js
index 8f39f69f10889..b403024795bb0 100644
--- a/packages/relay-experimental/useFragment.js
+++ b/packages/relay-experimental/useFragment.js
@@ -19,7 +19,7 @@ const useStaticFragmentNodeWarning = require('./useStaticFragmentNodeWarning');
 const {useTrackLoadQueryInRender} = require('./loadQuery');
 const {getFragment} = require('relay-runtime');
 
-import type {GraphQLTaggedNode} from 'relay-runtime';
+import type {GraphQLTaggedNode, FragmentReference} from 'relay-runtime';
 
 // NOTE: These declares ensure that the type of the returned data is:
 //   - non-nullable if the provided ref type is non-nullable
@@ -28,12 +28,16 @@ import type {GraphQLTaggedNode} from 'relay-runtime';
 //     non-nullable refs
 //   - array of nullable if the privoided ref type is an array of nullable refs
 
-declare function useFragment<TKey: {+$data?: mixed, ...}>(
+declare function useFragment<
+  TKey: {+$data?: mixed, +$fragmentRefs: FragmentReference, ...},
+>(
   fragmentInput: GraphQLTaggedNode,
   fragmentRef: TKey,
 ): $Call<<TFragmentData>({+$data?: TFragmentData, ...}) => TFragmentData, TKey>;
 
-declare function useFragment<TKey: ?{+$data?: mixed, ...}>(
+declare function useFragment<
+  TKey: ?{+$data?: mixed, +$fragmentRefs: FragmentReference, ...},
+>(
   fragmentInput: GraphQLTaggedNode,
   fragmentRef: TKey,
 ): $Call<
@@ -41,7 +45,13 @@ declare function useFragment<TKey: ?{+$data?: mixed, ...}>(
   TKey,
 >;
 
-declare function useFragment<TKey: $ReadOnlyArray<{+$data?: mixed, ...}>>(
+declare function useFragment<
+  TKey: $ReadOnlyArray<{
+    +$data?: mixed,
+    +$fragmentRefs: FragmentReference,
+    ...
+  }>,
+>(
   fragmentInput: GraphQLTaggedNode,
   fragmentRef: TKey,
 ): $Call<
@@ -51,7 +61,13 @@ declare function useFragment<TKey: $ReadOnlyArray<{+$data?: mixed, ...}>>(
   TKey,
 >;
 
-declare function useFragment<TKey: ?$ReadOnlyArray<{+$data?: mixed, ...}>>(
+declare function useFragment<
+  TKey: ?$ReadOnlyArray<{
+    +$data?: mixed,
+    +$fragmentRefs: FragmentReference,
+    ...
+  }>,
+>(
   fragmentInput: GraphQLTaggedNode,
   fragmentRef: TKey,
 ): $Call<
@@ -63,7 +79,7 @@ declare function useFragment<TKey: ?$ReadOnlyArray<{+$data?: mixed, ...}>>(
 
 function useFragment(
   fragmentInput: GraphQLTaggedNode,
-  fragmentRef: ?$ReadOnlyArray<{+$data?: mixed, ...}> | ?{+$data?: mixed, ...},
+  fragmentRef: mixed,
 ): mixed {
   // We need to use this hook in order to be able to track if
   // loadQuery was called during render