Add new streaming mode for LargeText (#722)

Co-authored-by: Tim Jacomb <[email protected]>

Author: das7pad

core/src/main/java/org/kohsuke/stapler/framework/io/LargeText.java

@@ -40,9 +40,12 @@
 import java.nio.charset.Charset;
 import java.nio.file.Files;
 import java.nio.file.StandardOpenOption;
+import java.util.UUID;
 import java.util.zip.GZIPInputStream;
+import net.sf.json.JSONObject;
 import org.apache.commons.io.output.CountingOutputStream;
 import org.kohsuke.stapler.ReflectionUtils;
+import org.kohsuke.stapler.Stapler;
 import org.kohsuke.stapler.StaplerRequest;
 import org.kohsuke.stapler.StaplerRequest2;
 import org.kohsuke.stapler.StaplerResponse;
@@ -84,12 +87,19 @@ public interface Source {
         boolean exists();
     }
 
+    /**
+     * For multipart/form-data streaming mode: Enable searching for the first new line character after the provided ?start and abort searching at the given position.
+     */
+    protected static final String SEARCH_STOP_PARAMETER = "searchNewLineUntil";
+
     private final Source source;
 
     protected final Charset charset;
 
     private volatile boolean completed;
 
+    private JSONObject streamingMeta;
+
     public LargeText(File file, boolean completed) {
         this(file, Charset.defaultCharset(), completed);
     }
@@ -257,13 +267,24 @@ public long writeLogTo(long start, OutputStream out) throws IOException {
     }
 
     private void writeLogUncounted(long start, OutputStream os) throws IOException {
-
         try (Session f = source.open()) {
             if (f.skip(start) != start) {
                 throw new EOFException("Attempted to read past the end of the log");
             }
 
-            if (completed) {
+            if (isStreamingRequest(Stapler.getCurrentRequest2())) {
+                // write everything until the previously determined EOF in bulk
+                long end = source.length();
+                long pos = start;
+                byte[] buf = new byte[64 * 1024];
+                int n;
+                while (pos < end && (n = f.read(buf)) >= 0) {
+                    long remaining = end - pos;
+                    if (n > remaining) n = (int) remaining;
+                    os.write(buf, 0, n);
+                    pos += n;
+                }
+            } else if (completed) {
                 // write everything till EOF
                 byte[] buf = new byte[1024];
                 int sz;
@@ -310,7 +331,111 @@ public void doProgressText(StaplerRequest req, StaplerResponse rsp) throws IOExc
         doProgressTextImpl(StaplerRequest.toStaplerRequest2(req), StaplerResponse.toStaplerResponse2(rsp));
     }
 
+    /**
+     * Detect use of streaming mode.
+     * @param req The current request.
+     * @return true if the new streaming mode is requested.
+     */
+    public boolean isStreamingRequest(StaplerRequest2 req) {
+        if (req == null) return false;
+        String accept = req.getHeader("Accept");
+        if (accept == null || accept.isEmpty()) return false;
+        return accept.startsWith("multipart/form-data");
+    }
+
+    /**
+     * Add additional meta data to a streaming response.
+     * @param key The field to (over)write meta data for.
+     * @param value The meta data value.
+     */
+    protected void putStreamingMeta(String key, Object value) {
+        if (streamingMeta == null) streamingMeta = new JSONObject();
+        streamingMeta.put(key, value);
+    }
+
+    private long findNextLineStart(long start, long stop) throws IOException {
+        try (var f = source.open()) {
+            if (f.skip(start) != start) {
+                // The log file rolled over, send it in full.
+                return 0;
+            }
+            byte[] buf = new byte[64 * 1024];
+            long searchPosition = start;
+            int n;
+            while (searchPosition + 1 < stop && (n = f.read(buf)) > 0) {
+                for (int i = 0; i < n && searchPosition + 1 < stop; i++) {
+                    if (buf[i] == '\n') {
+                        // We have a match, send from after the \n.
+                        putStreamingMeta("startFromNewLine", true);
+                        return searchPosition + 1;
+                    }
+                    searchPosition++;
+                }
+            }
+        }
+        // fall back to original start.
+        return start;
+    }
+
+    private void doProgressTextStreaming(StaplerRequest2 req, StaplerResponse2 rsp) throws IOException {
+        setContentType(rsp);
+        String contentType = rsp.getContentType();
+        if (contentType == null || contentType.isEmpty()) {
+            contentType = "text/plain;charset=UTF-8";
+        }
+        if (contentType.contains("\r") || contentType.contains("\n")) {
+            throw new IOException("Found CR/LF in Content-Type. Aborting streaming mode");
+        }
+        String boundary = UUID.randomUUID().toString();
+        rsp.setContentType("multipart/form-data;boundary=" + boundary);
+        rsp.setStatus(HttpServletResponse.SC_OK);
+        putStreamingMeta("completed", completed);
+
+        try (var writer = rsp.getWriter()) {
+            writer.write("--" + boundary + "\r\n" // preamble for first part
+                    + "Content-Disposition: form-data;name=text\r\n"
+                    + "Content-Type: " + contentType + "\r\n"
+                    + "\r\n");
+            long length = source.exists() ? source.length() : 0;
+
+            String s = req.getParameter("start");
+            long start = (s != null) ? Long.parseLong(s) : 0;
+            if (start > length) {
+                // text rolled over, send in full
+                start = 0;
+            } else if (start < 0 && length <= -start) {
+                // tail on small file, send in full
+                start = 0;
+            } else if (start < 0) {
+                // tail on large file, start at first new line in tail
+                start = findNextLineStart(length + start, length);
+            } else if (req.getParameter(SEARCH_STOP_PARAMETER) != null) {
+                // fetch more, start at first new line before last start
+                long searchStop = Long.parseLong(req.getParameter(SEARCH_STOP_PARAMETER));
+                start = findNextLineStart(start, searchStop);
+            }
+            putStreamingMeta("start", start);
+
+            long end = length;
+            if (start != end) {
+                end = writeLogTo(start, writer);
+            }
+            putStreamingMeta("end", end);
+
+            writer.write("\r\n--" + boundary + "\r\n" // transition to next part
+                    + "Content-Disposition: form-data;name=meta\r\n"
+                    + "Content-Type: application/json;charset=utf-8\r\n"
+                    + "\r\n"
+                    + streamingMeta + "\r\n--" + boundary + "--");
+        }
+    }
+
     private void doProgressTextImpl(StaplerRequest2 req, StaplerResponse2 rsp) throws IOException {
+        if (isStreamingRequest(req)) {
+            doProgressTextStreaming(req, rsp);
+            return;
+        }
+
         setContentType(rsp);
         rsp.setStatus(HttpServletResponse.SC_OK);