From 68cbd241e2172b8018328d77ea087dd5974a580f Mon Sep 17 00:00:00 2001
From: killa <zhubin.gzb@antgroup.com>
Date: Fri, 12 Jul 2024 00:32:47 +0800
Subject: [PATCH] feat: impl httpclient.safeCurl (#5339)

Co-authored-by: fengmk2 <suqian.yf@antgroup.com>
---
 lib/core/httpclient.js           | 11 +++++++++
 lib/core/httpclient_next.js      | 16 +++++++++++++
 test/lib/core/httpclient.test.js | 40 ++++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+)

diff --git a/lib/core/httpclient.js b/lib/core/httpclient.js
index a270a8c026..97b167b0e0 100644
--- a/lib/core/httpclient.js
+++ b/lib/core/httpclient.js
@@ -44,6 +44,17 @@ class HttpClient extends urllib.HttpClient2 {
   async curl(...args) {
     return await this.request(...args);
   }
+
+  async safeCurl(url, options = {}) {
+    const ssrfConfig = this.app.config.security.ssrf;
+    if (ssrfConfig?.checkAddress) {
+      options.checkAddress = ssrfConfig.checkAddress;
+    } else {
+      this.app.logger.warn('[egg-security] please configure `config.security.ssrf` first');
+    }
+
+    return await this.curl(url, options);
+  }
 }
 
 function normalizeConfig(app) {
diff --git a/lib/core/httpclient_next.js b/lib/core/httpclient_next.js
index 1a18e6ce23..7d0fc4afea 100644
--- a/lib/core/httpclient_next.js
+++ b/lib/core/httpclient_next.js
@@ -1,5 +1,6 @@
 const { HttpClient } = require('urllib-next');
 const ms = require('humanize-ms');
+const SSRF_HTTPCLIENT = Symbol('SSRF_HTTPCLIENT');
 
 class HttpClientNext extends HttpClient {
   constructor(app, options) {
@@ -33,6 +34,21 @@ class HttpClientNext extends HttpClient {
   async curl(...args) {
     return await this.request(...args);
   }
+
+  async safeCurl(url, options = {}) {
+    if (!this[SSRF_HTTPCLIENT]) {
+      const ssrfConfig = this.app.config.security.ssrf;
+      if (ssrfConfig?.checkAddress) {
+        options.checkAddress = ssrfConfig.checkAddress;
+      } else {
+        this.app.logger.warn('[egg-security] please configure `config.security.ssrf` first');
+      }
+      this[SSRF_HTTPCLIENT] = new HttpClientNext(this.app, {
+        checkAddress: ssrfConfig.checkAddress,
+      });
+    }
+    return await this[SSRF_HTTPCLIENT].request(url, options);
+  }
 }
 
 function normalizeConfig(app) {
diff --git a/test/lib/core/httpclient.test.js b/test/lib/core/httpclient.test.js
index a98daf7390..e0b26b0da3 100644
--- a/test/lib/core/httpclient.test.js
+++ b/test/lib/core/httpclient.test.js
@@ -98,6 +98,26 @@ describe('test/lib/core/httpclient.test.js', () => {
     });
   });
 
+  it('should support safeCurl', async () => {
+    let ip;
+    let family;
+    let host;
+    mm(client.app.config, 'security', {
+      ssrf: {
+        checkAddress(aIp, aFamilay, aHost) {
+          ip = aIp;
+          family = aFamilay;
+          host = aHost;
+          return true;
+        },
+      },
+    });
+    await client.safeCurl(url);
+    assert(ip);
+    assert(family);
+    assert(host);
+  });
+
   describe('HttpClientNext', () => {
     it('should request ok with log', async () => {
       const args = {
@@ -145,6 +165,26 @@ describe('test/lib/core/httpclient.test.js', () => {
         return true;
       });
     });
+
+    it('should support safeCurl', async () => {
+      let ip;
+      let family;
+      let host;
+      mm(clientNext.app.config, 'security', {
+        ssrf: {
+          checkAddress(aIp, aFamilay, aHost) {
+            ip = aIp;
+            family = aFamilay;
+            host = aHost;
+            return true;
+          },
+        },
+      });
+      await clientNext.safeCurl(url);
+      assert(ip);
+      assert(family);
+      assert(host);
+    });
   });
 
   describe('httpclient.httpAgent.timeout < 30000', () => {