From 4d15bd042d3d636f8dc70605f9f692ecdd6e1a71 Mon Sep 17 00:00:00 2001
From: edvardchen <pigeon73101@gmail.com>
Date: Tue, 13 Aug 2019 11:20:38 +0800
Subject: [PATCH] feat(ges): prevent interceptor from postprocess response on
 unary call

simplify ctx.response
---
 .../grpc-experimental-server.test.ts          | 15 ++++---
 .../grpc-experimental-server/src/index.ts     | 39 +++++++------------
 2 files changed, 24 insertions(+), 30 deletions(-)

diff --git a/packages/grpc-experimental-server/__tests__/grpc-experimental-server.test.ts b/packages/grpc-experimental-server/__tests__/grpc-experimental-server.test.ts
index 1b40f04..b98fe1f 100644
--- a/packages/grpc-experimental-server/__tests__/grpc-experimental-server.test.ts
+++ b/packages/grpc-experimental-server/__tests__/grpc-experimental-server.test.ts
@@ -72,7 +72,7 @@ describe('grpc-experimental-server', () => {
     });
     runTest({
       implementations: {
-        async getFeature(call, callback) {
+        getFeature(call: ServerUnaryCall<unknown>, callback: sendUnaryData<unknown>) {
           const error: ServiceError = new Error('unexpected');
           error.code = status.PERMISSION_DENIED;
           callback(error, new Feature());
@@ -83,10 +83,14 @@ describe('grpc-experimental-server', () => {
           let finished = false;
           server().use(async (ctx, next) => {
             ctx.onFinished(error => {
-              expect(error).toBeInstanceOf(Error);
+              try {
+                expect(error).toBeInstanceOf(Error);
+              } catch (e) {
+                done(e);
+              }
               finished = true;
             });
-            await next();
+            next().catch(() => undefined);
           });
           // @ts-ignore
           client().getFeature(new Point(), () => {
@@ -101,10 +105,11 @@ describe('grpc-experimental-server', () => {
   describe('postprocess', () => {
     runTest({
       testcase(server, client) {
-        it('should get Feature', () => {
+        it('should get Feature', done => {
           server().use(async (ctx, next) => {
             await next();
-            expect(ctx.response.value).toBeInstanceOf(Feature);
+            expect(ctx.response).toBeInstanceOf(Feature);
+            done();
           });
 
           // @ts-ignore
diff --git a/packages/grpc-experimental-server/src/index.ts b/packages/grpc-experimental-server/src/index.ts
index a2e0020..d487582 100644
--- a/packages/grpc-experimental-server/src/index.ts
+++ b/packages/grpc-experimental-server/src/index.ts
@@ -1,6 +1,5 @@
 import compose from 'koa-compose';
 import {
-  Metadata,
   Server,
   UntypedServiceImplementation,
   ServiceDefinition,
@@ -25,11 +24,7 @@ type ServerCall =
 type ServerNonStreamCall = ServerUnaryCall<unknown> | ServerReadableStream<unknown>;
 
 export class Context {
-  response: {
-    value?: unknown;
-    trailer?: Metadata;
-    flags?: number;
-  } = {};
+  response: unknown;
   constructor(public call: ServerCall, public definition: MethodDefinition<unknown, unknown>) {}
   onFinished(listener: (err: Error | null) => void): void {
     const emitter = this.call as EventEmitter;
@@ -85,15 +80,24 @@ export default class ExperimentalServer extends Server {
 
       const ctx = new Context(call, definition);
 
-      const pending = this.handleRequest(ctx, async () => {
+      this.handleRequest(ctx, async () => {
         // unary call
         if (grpcCallback) {
           const nonStreamCall = call as ServerNonStreamCall;
           return new Promise((resolve, reject) => {
-            const callback: sendUnaryData<unknown> = (error, value, trailer, flags) => {
-              if (error) return reject(error);
-              ctx.response = { ...ctx.response, value, trailer, flags };
+            const callback: sendUnaryData<unknown> = (error, value, ...rest) => {
+              if (error) {
+                grpcCallback(error, value);
+                reject(error);
+                (call as EventEmitter).emit('finish', error);
+                return;
+              }
+
+              // REAL SEND RESPONSE
+              grpcCallback(null, value, ...rest);
+              ctx.response = value;
               resolve();
+              (call as EventEmitter).emit('finish');
             };
             // @ts-ignore
             original(nonStreamCall, callback);
@@ -104,21 +108,6 @@ export default class ExperimentalServer extends Server {
         original(call);
         return;
       });
-      if (grpcCallback) {
-        pending
-          .then(() => {
-            const {
-              response: { value, trailer, flags },
-            } = ctx;
-            // real send response
-            grpcCallback(null, value, trailer, flags);
-            (call as EventEmitter).emit('finish');
-          })
-          .catch(error => {
-            grpcCallback(error, null);
-            (call as EventEmitter).emit('finish', error);
-          });
-      }
     };
   }