From 9d97f416750eacf1f15d005677980c64c6b0cc81 Mon Sep 17 00:00:00 2001
From: Sofian Hnaide <sofian.hnaide@gmail.com>
Date: Fri, 18 Jun 2021 10:45:48 -0700
Subject: [PATCH] fix race condition in setContext link that can leak
 subscription

---
 src/link/context/__tests__/index.ts | 33 +++++++++++++++++++++++++++++
 src/link/context/index.ts           |  4 ++++
 2 files changed, 37 insertions(+)

diff --git a/src/link/context/__tests__/index.ts b/src/link/context/__tests__/index.ts
index f677bd664f5..af2c4294021 100644
--- a/src/link/context/__tests__/index.ts
+++ b/src/link/context/__tests__/index.ts
@@ -205,3 +205,36 @@ it('unsubscribes without throwing before data', done => {
     done();
   }, 10);
 });
+
+it('does not start the next link subscription if the upstream subscription is already closed', done => {
+  let promiseResolved = false;
+  const withContext = setContext(() =>
+    sleep(5).then(() => {
+      promiseResolved = true;
+      return { dynamicallySet: true }
+    }),
+  );
+
+  let mockLinkCalled = false;
+  const mockLink = new ApolloLink(operation => {
+    mockLinkCalled = true;
+    fail('link should not be called');
+  });
+
+  const link = withContext.concat(mockLink);
+
+  let subscriptionReturnedData = false;
+  let handle = execute(link, { query }).subscribe(result => {
+    subscriptionReturnedData = true;
+    fail('subscription should not return data');
+  });
+
+  handle.unsubscribe();
+
+  setTimeout(() => {
+    expect(promiseResolved).toBe(true);
+    expect(mockLinkCalled).toBe(false);
+    expect(subscriptionReturnedData).toBe(false);
+    done();
+  }, 10);
+});
diff --git a/src/link/context/index.ts b/src/link/context/index.ts
index c427a769fa3..67b94f13af6 100644
--- a/src/link/context/index.ts
+++ b/src/link/context/index.ts
@@ -12,10 +12,13 @@ export function setContext(setter: ContextSetter): ApolloLink {
 
     return new Observable(observer => {
       let handle: ZenObservable.Subscription;
+      let closed = false;
       Promise.resolve(request)
         .then(req => setter(req, operation.getContext()))
         .then(operation.setContext)
         .then(() => {
+          // if the observer is already closed, no need to subscribe.
+          if (closed) return;
           handle = forward(operation).subscribe({
             next: observer.next.bind(observer),
             error: observer.error.bind(observer),
@@ -25,6 +28,7 @@ export function setContext(setter: ContextSetter): ApolloLink {
         .catch(observer.error.bind(observer));
 
       return () => {
+        closed = true;
         if (handle) handle.unsubscribe();
       };
     });