datasource escape with special check (sofastack#411)

Huoxi-any · Dec 28, 2020 · cd1dca3 · cd1dca3
1 parent ea6f55d
commit cd1dca3
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 3 deletions.
diff --git a/...source-plugin/src/main/java/com/alipay/sofa/tracer/plugins/datasource/utils/SqlUtils.java b/...source-plugin/src/main/java/com/alipay/sofa/tracer/plugins/datasource/utils/SqlUtils.java
@@ -36,9 +36,7 @@ public static String getSqlEscaped(String sql) {
         if (limitSql != null && limitSql.length() > DIGEST_LOG_SQL_LIMIT) {
             limitSql = limitSql.substring(0, DIGEST_LOG_SQL_LIMIT) + " ...";
         }
-        return escape(
-            escape(escape(limitSql, DEFAULT_SEPARATOR, DEFAULT_SEPARATOR_ESCAPE), DEFAULT_NEW_LINE,
-                EMPTY_STRING), DEFAULT_RETURN, EMPTY_STRING);
+        return escape(escapeWithSpecialCheck(limitSql), DEFAULT_SEPARATOR, DEFAULT_SEPARATOR_ESCAPE);
     }
 
     private static String escape(String str, String oldStr, String newStr) {
@@ -47,4 +45,24 @@ private static String escape(String str, String oldStr, String newStr) {
         }
         return StringUtils.replace(str, oldStr, newStr);
     }
+
+    private static String escapeWithSpecialCheck(String str) {
+        if (str == null) {
+            return "";
+        }
+        StringBuilder appender = new StringBuilder();
+        int len = str.length();
+        appender.ensureCapacity(appender.length() + len);
+        for (int i = 0; i < len; i++) {
+            char c = str.charAt(i);
+            if (c == '\n' || c == '\r' || c == '|') {
+                c = ' ';
+            }
+            if (c != ' ' || c == ' ' && appender.length() > 0
+                && appender.charAt(appender.length() - 1) != ' ') {
+                appender.append(c);
+            }
+        }
+        return appender.toString().trim();
+    }
 }
diff --git a/...atasource-plugin/src/test/java/com/sofa/tracer/plugins/datasource/DataSourceUtilTest.java b/...atasource-plugin/src/test/java/com/sofa/tracer/plugins/datasource/DataSourceUtilTest.java
@@ -17,6 +17,7 @@
 package com.sofa.tracer.plugins.datasource;
 
 import com.alipay.common.tracer.core.utils.ReflectionUtils;
+import com.alipay.sofa.tracer.plugins.datasource.utils.SqlUtils;
 import com.sofa.tracer.plugins.datasource.bean.ConcreteClassService;
 import org.junit.Assert;
 import org.junit.Test;
@@ -38,4 +39,21 @@ public void testReflectionUtils() throws Throwable {
         method = ReflectionUtils.findMethod(concreteClassService.getClass(), "serviceB");
         Assert.assertEquals("serviceB", method.invoke(concreteClassService));
     }
+
+    @Test
+    public void testGetSqlEscaped() {
+        // start with blank and special character
+        String str1 = "  select app1\n" + ",app2\r\n" + " from table where          " + " id = a;";
+        String result = SqlUtils.getSqlEscaped(str1);
+        Assert.assertTrue(!result.contains("\n"));
+        Assert.assertTrue(!result.contains("\r"));
+        Assert.assertTrue(result.equals("select app1 %2Capp2 from table where id = a;"));
+        Assert.assertTrue(!result.startsWith(" "));
+        // normal string
+        String str2 = "select app1 ,app2 from table where id = a;";
+        result = SqlUtils.getSqlEscaped(str2);
+        Assert.assertTrue(!result.contains("\n"));
+        Assert.assertTrue(!result.contains("\r"));
+        Assert.assertTrue(result.equals("select app1 %2Capp2 from table where id = a;"));
+    }
 }