From de06b1e28b996e8cb910d725c0b8970a007afcdf Mon Sep 17 00:00:00 2001
From: liningrui <liningrui@baidu.com>
Date: Thu, 3 Jan 2019 14:30:36 +0800
Subject: [PATCH] Limit the name of the gremlin request less than 256 bytes

Change-Id: I46ccdffd92087b50676e1612226424208f66dbee
---
 .../baidu/hugegraph/api/job/GremlinAPI.java   | 30 ++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/hugegraph-api/src/main/java/com/baidu/hugegraph/api/job/GremlinAPI.java b/hugegraph-api/src/main/java/com/baidu/hugegraph/api/job/GremlinAPI.java
index b6b1f20e9f..dc1b66ae8f 100644
--- a/hugegraph-api/src/main/java/com/baidu/hugegraph/api/job/GremlinAPI.java
+++ b/hugegraph-api/src/main/java/com/baidu/hugegraph/api/job/GremlinAPI.java
@@ -19,6 +19,11 @@
 
 package com.baidu.hugegraph.api.job;
 
+import java.nio.ByteBuffer;
+import java.nio.charset.CharacterCodingException;
+import java.nio.charset.Charset;
+import java.nio.charset.CharsetDecoder;
+import java.nio.charset.CodingErrorAction;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -32,6 +37,7 @@
 
 import org.slf4j.Logger;
 
+import com.baidu.hugegraph.HugeException;
 import com.baidu.hugegraph.HugeGraph;
 import com.baidu.hugegraph.api.API;
 import com.baidu.hugegraph.api.filter.StatusFilter.Status;
@@ -58,6 +64,8 @@ public class GremlinAPI extends API {
 
     private static final Logger LOG = Log.logger(RestServer.class);
 
+    private static final int MAX_NAME_LENGTH = 256;
+
     private static final Histogram gremlinJobInputHistogram =
             MetricsUtil.registerHistogram(GremlinAPI.class, "gremlin-input");
 
@@ -148,7 +156,27 @@ public Map<String, String> aliases() {
 
         public String name() {
             // Get the first line of script as the name
-            return this.gremlin.split("\r\n|\r|\n", 2)[0];
+            String firstLine = this.gremlin.split("\r\n|\r|\n", 2)[0];
+            final Charset charset = Charset.forName(CHARSET);
+            final byte[] bytes = firstLine.getBytes(charset);
+            if (bytes.length <= MAX_NAME_LENGTH) {
+                return firstLine;
+            }
+
+            /*
+             * Reference https://stackoverflow.com/questions/3576754/truncating-strings-by-bytes
+             */
+            CharsetDecoder decoder = charset.newDecoder();
+            decoder.onMalformedInput(CodingErrorAction.IGNORE);
+            decoder.reset();
+
+            ByteBuffer buffer = ByteBuffer.wrap(bytes, 0, MAX_NAME_LENGTH);
+            try {
+                return decoder.decode(buffer).toString();
+            } catch (CharacterCodingException e) {
+                throw new HugeException("Failed to decode truncated bytes of " +
+                                        "gremlin first line", e);
+            }
         }
 
         @Override