From 5de5284c70b44acc73eaaad651fd2edc72d54f28 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=94=BF?= <389428742@qq.com>
Date: Mon, 16 Aug 2021 22:06:12 +0800
Subject: [PATCH] fix: correct aspect chain bug (#1204)

Co-authored-by: chenzhaozheng <chenzhaozheng@LAPTOP-GK4GS2RH.localdomain>
---
 packages/core/src/service/aspectService.ts      |  7 ++++++-
 packages/core/test/baseFramework.test.ts        |  3 +++
 .../fixtures/base-app-aspect/src/aspect/c.ts    |  5 +++++
 .../test/fixtures/base-app-aspect/src/home.ts   | 17 +++++++++++++----
 packages/ws/src/framework.ts                    |  4 ++--
 5 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/packages/core/src/service/aspectService.ts b/packages/core/src/service/aspectService.ts
index b9055f0b325b..7f7c586adb6b 100644
--- a/packages/core/src/service/aspectService.ts
+++ b/packages/core/src/service/aspectService.ts
@@ -6,6 +6,7 @@ import {
   isAsyncFunction,
   isProxy,
   listModule,
+  isClass,
 } from '@midwayjs/decorator';
 import * as pm from 'picomatch';
 import { IAspectService, IMidwayContainer } from '../interface';
@@ -181,7 +182,11 @@ export class MidwayAspectService implements IAspectService {
               }
               const resultTemp = aspectIns.afterReturn?.(joinPoint, result);
               result = typeof resultTemp === 'undefined' ? result : resultTemp;
-              return result;
+              if (result && isClass(result.constructor)) {
+                return this.wrapperAspectToInstance(ins);
+              } else {
+                return result;
+              }
             } catch (err) {
               error = err;
               if (aspectIns.afterThrow) {
diff --git a/packages/core/test/baseFramework.test.ts b/packages/core/test/baseFramework.test.ts
index 312149e20795..d890d0eb487a 100644
--- a/packages/core/test/baseFramework.test.ts
+++ b/packages/core/test/baseFramework.test.ts
@@ -474,6 +474,9 @@ describe('/test/baseFramework.test.ts', () => {
     } catch (err) {
       expect(err.message).toMatch('ccc');
     }
+    // aspect chain
+    const result = await userController1.test1().test2().getUser1();
+    expect(result).toEqual('before test user');
   });
 
   it('should inject global value in component', async () => {
diff --git a/packages/core/test/fixtures/base-app-aspect/src/aspect/c.ts b/packages/core/test/fixtures/base-app-aspect/src/aspect/c.ts
index 554aa3e6a82f..1e6ddb0e2290 100644
--- a/packages/core/test/fixtures/base-app-aspect/src/aspect/c.ts
+++ b/packages/core/test/fixtures/base-app-aspect/src/aspect/c.ts
@@ -11,4 +11,9 @@ export class MyAspect3 implements IMethodAspect {
       throw new Error('ccc');
     }
   }
+  async before(point: JoinPoint) {
+    console.log('before around in aspect3');
+    console.log('before methodName:', point.methodName);
+    point.args = ['before test user'];
+  }
 }
diff --git a/packages/core/test/fixtures/base-app-aspect/src/home.ts b/packages/core/test/fixtures/base-app-aspect/src/home.ts
index f8c4ba889f7c..df8c4d3d8b49 100644
--- a/packages/core/test/fixtures/base-app-aspect/src/home.ts
+++ b/packages/core/test/fixtures/base-app-aspect/src/home.ts
@@ -8,13 +8,12 @@ class Parent {
   }
 
   async hello1() {
-    return 'hello world 1'
+    return 'hello world 1';
   }
 }
 
 @Provide()
 export class Home extends Parent {
-
   bbb = 'aaa';
 
   ccc: string;
@@ -28,14 +27,24 @@ export class Home extends Parent {
   }
 }
 
-
 @Provide()
 export class UserController {
-
   @Inject()
   ctx;
 
   async getUser() {
     throw new Error('bbb');
   }
+
+  async getUser1(uid = 'user1') {
+    return uid;
+  }
+
+  test1() {
+    return this;
+  }
+
+  test2() {
+    return this;
+  }
 }
diff --git a/packages/ws/src/framework.ts b/packages/ws/src/framework.ts
index 9f1da674d42b..86a88393af79 100644
--- a/packages/ws/src/framework.ts
+++ b/packages/ws/src/framework.ts
@@ -53,8 +53,8 @@ export class MidwayWSFramework extends BaseFramework<
       server = this.configurationOptions.server ?? http.createServer();
     }
 
-    server.on('upgrade', (request, socket, head) => {
-      this.app.handleUpgrade(request, socket as any, head, ws => {
+    server.on('upgrade', (request, socket: any, head) => {
+      this.app.handleUpgrade(request, socket, head, ws => {
         this.app.emit('connection', ws, request);
       });
     });