io.datakernel.http.MultipartParser Maven / Gradle / Ivy

Go to download
/*
 * Copyright (C) 2015-2019 SoftIndex LLC.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package io.datakernel.http;

import io.datakernel.bytebuf.ByteBuf;
import io.datakernel.bytebuf.ByteBufPool;
import io.datakernel.bytebuf.ByteBufQueue;
import io.datakernel.common.ApplicationSettings;
import io.datakernel.common.Recyclable;
import io.datakernel.common.exception.StacklessException;
import io.datakernel.common.ref.Ref;
import io.datakernel.csp.ChannelConsumer;
import io.datakernel.csp.ChannelConsumers;
import io.datakernel.csp.ChannelSupplier;
import io.datakernel.csp.binary.BinaryChannelSupplier;
import io.datakernel.csp.binary.ByteBufsParser;
import io.datakernel.http.MultipartParser.MultipartFrame;
import io.datakernel.promise.Promise;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.function.Function;

import static io.datakernel.bytebuf.ByteBufStrings.CR;
import static io.datakernel.bytebuf.ByteBufStrings.LF;
import static io.datakernel.common.MemSize.kilobytes;
import static io.datakernel.common.Utils.nullify;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.stream.Collectors.toMap;

/**
 * Util class that allows to parse some binary channel (mainly, the request body stream) into a channel of multipart frames.
 */
public final class MultipartParser implements ByteBufsParser {
	private static final int MAX_META_SIZE = ApplicationSettings.getMemSize(MultipartParser.class, "maxMetaBuffer", kilobytes(4)).toInt();

	@Nullable
	private List readingHeaders = null;

	private final byte[] boundary;
	private final byte[] lastBoundary;

	private MultipartParser(String boundary) {
		this.boundary = ("--" + boundary).getBytes(UTF_8);
		this.lastBoundary = ("--" + boundary + "--").getBytes(UTF_8);
	}

	public static MultipartParser create(String boundary) {
		return new MultipartParser(boundary);
	}

	/**
	 * Converts resulting channel of frames into a binary channel, ignoring any multipart headers.
	 */
	public ByteBufsParser ignoreHeaders() {
		return bufs -> {
			MultipartFrame frame = tryParse(bufs);
			if (frame == null || frame.isHeaders()) {
				return null;
			}
			return frame.getData();
		};
	}

	private Promise> getContentDispositionFields(MultipartFrame frame) {
		String header = frame.getHeaders().get("content-disposition");
		if (header == null) {
			return Promise.ofException(new StacklessException(MultipartParser.class, "Headers had no Content-Disposition"));
		}
		String[] headerParts = header.split(";");
		if (headerParts.length == 0 || !"form-data".equals(headerParts[0].trim())) {
			return Promise.ofException(new StacklessException(MultipartParser.class, "Content-Disposition type is not 'form-data'"));
		}
		return Promise.of(Arrays.stream(headerParts)
				.skip(1)
				.map(part -> part.trim().split("=", 2))
				.collect(toMap(s -> s[0], s -> {
					String value = s.length == 1 ? "" : s[1];
					// stripping double quotation
					return value.substring(1, value.length() - 1);
				})));
	}

	private Promise doSplit(MultipartFrame headerFrame, ChannelSupplier frames,
			MultipartDataHandler dataHandler) {
		return getContentDispositionFields(headerFrame)
				.then(contentDispositionFields -> {
					String fieldName = contentDispositionFields.get("name");
					String fileName = contentDispositionFields.get("filename");
					Ref lastRef = new Ref<>();
					return frames
							.until(f -> {
								if (f.isHeaders()) {
									lastRef.set(f);
									return true;
								}
								return false;
							})
							.filter(MultipartFrame::isData)
							.map(MultipartFrame::getData)
							.streamTo(ChannelConsumer.ofPromise(fileName == null ?
									dataHandler.handleField(fieldName) :
									dataHandler.handleFile(fieldName, fileName)
							))
							.then($ -> lastRef.get() != null ?
									doSplit(lastRef.get(), frames, dataHandler) :
									Promise.complete())
							.toVoid();
				});
	}

	/**
	 * Complex operation that streams this channel of multipart frames into multiple binary consumers,
	 * as specified by the Content-Disposition multipart header.
	 */
	public Promise split(ChannelSupplier source, MultipartDataHandler dataHandler) {
		ChannelSupplier frames = BinaryChannelSupplier.of(source).parseStream(this);
		return frames.get()
				.then(frame -> {
					if (frame.isHeaders()) {
						return doSplit(frame, frames, dataHandler);
					}
					StacklessException e = new StacklessException(MultipartParser.class, "First frame had no headers");
					frames.close(e);
					return Promise.ofException(e);
				});
	}

	private boolean sawCrlf = true;
	private boolean finished = false;

	@Nullable
	@Override
	public MultipartFrame tryParse(ByteBufQueue bufs) {
		if (finished) {
			return null;
		}
		for (int i = 0; i < bufs.remainingBytes() - 1; i++) {
			if (bufs.peekByte(i) != CR || bufs.peekByte(i + 1) != LF) {
				continue;
			}
			if (sawCrlf) {
				ByteBuf term = bufs.takeExactSize(i);
				if (readingHeaders == null) {
					if (term.isContentEqual(lastBoundary)) {
						bufs.skip(2);
						finished = true;
						term.recycle();
						return null;
					} else if (term.isContentEqual(boundary)) {
						bufs.skip(2);
						i = -1; // fix the index (so that it's 0 on next iteration) because we've taken bytes from queue
						term.recycle();
						readingHeaders = new ArrayList<>();
					} else {
						sawCrlf = false;
						return getFalseTermFrame(term);
					}
				} else {
					bufs.skip(2);
					if (i != 0) {
						i = -1; // see above comment
						readingHeaders.add(term.asString(UTF_8));
						continue;
					}
					sawCrlf = false;
					term.recycle();
					List readingHeaders = this.readingHeaders;
					this.readingHeaders = null;
					if (readingHeaders.isEmpty()) {
						break;
					}
					return MultipartFrame.of(readingHeaders.stream()
							.map(s -> s.split(":\\s?", 2))
							.collect(toMap(s -> s[0].toLowerCase(), s -> s[1])));
				}
			} else {
				sawCrlf = true;
				ByteBuf tail = bufs.takeExactSize(i);
				bufs.skip(2);
				return MultipartFrame.of(tail);
			}
		}

		int remaining = bufs.remainingBytes();
		if (sawCrlf) {
			if (remaining >= MAX_META_SIZE) {
				sawCrlf = false;
				return getFalseTermFrame(bufs.takeRemaining());
			}
			return null;
		}
		int toTake = remaining == 0 ? 0 : remaining - (bufs.peekByte(remaining - 1) == CR ? 1 : 0);
		if (toTake == 0) {
			return null;
		}
		ByteBuf data = bufs.takeExactSize(toTake);
		return MultipartFrame.of(data);
	}

	@NotNull
	private MultipartFrame getFalseTermFrame(ByteBuf term) {
		ByteBuf buf = ByteBufPool.allocate(term.readRemaining() + 2);
		buf.writeByte((byte) '\r');
		buf.writeByte((byte) '\n');
		term.drainTo(buf, term.readRemaining());
		term.recycle();
		return MultipartFrame.of(buf);
	}

	public static final class MultipartFrame implements Recyclable {
		@Nullable
		private ByteBuf data;
		@Nullable
		private final Map headers;

		private MultipartFrame(@Nullable ByteBuf data, @Nullable Map headers) {
			this.data = data;
			this.headers = headers;
		}

		public static MultipartFrame of(ByteBuf data) {
			return new MultipartFrame(data, null);
		}

		public static MultipartFrame of(Map headers) {
			return new MultipartFrame(null, headers);
		}

		public boolean isData() {
			return data != null;
		}

		public ByteBuf getData() {
			assert data != null : "Trying to get data out of header frame";
			return data;
		}

		public boolean isHeaders() {
			return headers != null;
		}

		public Map getHeaders() {
			assert headers != null : "Trying to get headers out of data frame";
			return headers;
		}

		@Override
		public void recycle() {
			data = nullify(data, ByteBuf::recycle);
		}

		@Override
		public String toString() {
			return isHeaders() ? "headers" + headers : "" + data;
		}
	}

	public interface MultipartDataHandler {
		Promise> handleField(String fieldName);

		Promise> handleFile(String fieldName, String fileName);

		static MultipartDataHandler fieldsToMap(Map fields) {
			return fieldsToMap(fields, ($1, $2) -> Promise.of(ChannelConsumers.recycling()));
		}

		static MultipartDataHandler fieldsToMap(Map fields,
				Function>> uploader) {
			return fieldsToMap(fields, ($, fileName) -> uploader.apply(fileName));
		}

		static MultipartDataHandler fieldsToMap(Map fields,
				BiFunction>> uploader) {
			return new MultipartDataHandler() {
				@Override
				public Promise> handleField(String fieldName) {
					return Promise.of(ChannelConsumer.ofSupplier(supplier -> supplier.toCollector(ByteBufQueue.collector())
							.map(value -> {
								fields.put(fieldName, value.asString(UTF_8));
								return (Void) null;
							})));
				}

				@Override
				public Promise> handleFile(String fieldName, String fileName) {
					return uploader.apply(fieldName, fileName);
				}
			};
		}

		static MultipartDataHandler file(Function>> uploader) {
			return files(($, fileName) -> uploader.apply(fileName));
		}

		static MultipartDataHandler files(BiFunction>> uploader) {
			return new MultipartDataHandler() {
				@Override
				public Promise> handleField(String fieldName) {
					return Promise.of(ChannelConsumers.recycling());
				}

				@Override
				public Promise> handleFile(String fieldName, String fileName) {
					return uploader.apply(fieldName, fileName);
				}
			};
		}

	}

}