Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 96 additions & 33 deletions taier-common/src/main/java/com/dtstack/taier/common/util/ZipUtil.java
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

package com.dtstack.taier.common.util;

import com.dtstack.taier.common.exception.TaierDefineException;
import org.apache.tools.zip.ZipFile;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand Down Expand Up @@ -54,6 +55,12 @@ public class ZipUtil {

private static byte[] _byte = new byte[1024];

private static final int MAX_ZIP_ENTRY_COUNT = 1000;
private static final int MAX_ZIP_RECURSION_DEPTH = 3;
private static final long MAX_ZIP_TOTAL_UNCOMPRESSED_SIZE = 100L * 1024 * 1024;
private static final long MAX_ZIP_ENTRY_UNCOMPRESSED_SIZE = 50L * 1024 * 1024;
private static final long MAX_ZIP_COMPRESSION_RATIO = 100L;

public static byte[] compress(byte[] rowData) {
byte[] backData = null;
ZipOutputStream zip = null;
Expand Down Expand Up @@ -224,68 +231,124 @@ public static List<File> upzipFile(String zipPath, String descDir) {
*/
@SuppressWarnings("rawtypes")
public static List<File> upzipFile(File zipFile, String descDir) {
try {
return upzipFile(zipFile, descDir, new UnzipContext(), 0);
} catch (IOException e) {
throw new TaierDefineException(String.format("Unzip exception : %s", e.getMessage()), e);
}
}

@SuppressWarnings("rawtypes")
private static List<File> upzipFile(File zipFile, String descDir, UnzipContext context, int depth) throws IOException {
if (depth > MAX_ZIP_RECURSION_DEPTH) {
throw new IOException(String.format("zip recursion depth exceeds limit: %s", MAX_ZIP_RECURSION_DEPTH));
}
List<File> _list = new ArrayList<>();
File baseDir = new File(descDir);
String basePath = getCanonicalDirPath(baseDir);
ZipFile _zipFile = null;
OutputStream _out = null;
InputStream _in = null;
try {
_zipFile = new ZipFile(zipFile, "GBK");
for (Enumeration entries = _zipFile.getEntries(); entries.hasMoreElements(); ) {
org.apache.tools.zip.ZipEntry entry = (org.apache.tools.zip.ZipEntry) entries.nextElement();
File _file = new File(descDir + File.separator + entry.getName());
context.addEntry(entry.getName());
File _file = resolveZipEntryFile(baseDir, basePath, entry.getName());
if (_file.isHidden()) {
continue;
}
if (entry.isDirectory()) {
_file.mkdirs();
makeDirs(_file);
} else {
File _parent = _file.getParentFile();
if (!_parent.exists()) {
_parent.mkdirs();
}
_in = _zipFile.getInputStream(entry);
_out = new FileOutputStream(_file);
makeDirs(_parent);
byte[] buffer = new byte[4];
int length = _in.read(buffer, 0, 4);
int len = 0;
_out.write(buffer);
while ((len = _in.read(_byte)) > 0) {
_out.write(_byte, 0, len);
byte[] fileBuffer = new byte[1024];
int length;
try (InputStream _in = _zipFile.getInputStream(entry);
OutputStream _out = new FileOutputStream(_file)) {
length = _in.read(buffer, 0, 4);
long written = 0L;
if (length > 0) {
_out.write(buffer, 0, length);
written += length;
context.addUncompressedSize(length);
validateEntrySize(entry, written);
}
int len = 0;
while ((len = _in.read(fileBuffer)) > 0) {
_out.write(fileBuffer, 0, len);
written += len;
context.addUncompressedSize(len);
validateEntrySize(entry, written);
}
_out.flush();
}
_out.flush();

if (length == 4 && (Arrays.equals(ZIP_HEADER_1, buffer) || Arrays.equals(ZIP_HEADER_2, buffer))) {
_list.addAll(upzipFile(_file, _file.getPath() + "tmp"));
_list.addAll(upzipFile(_file, _file.getPath() + "tmp", context, depth + 1));
} else {
_list.add(_file);
}

}
}
} catch (IOException e) {
} finally {
if (_out != null) {
try {
_out.close();
} catch (IOException e) {
}
}
if (_in != null) {
try {
_in.close();
} catch (IOException e) {
}
}
if (_zipFile != null) {
try {
_zipFile.close();
} catch (IOException e) {
}
_zipFile.close();
}
}
return _list;
}

private static File resolveZipEntryFile(File baseDir, String basePath, String entryName) throws IOException {
File targetFile = new File(baseDir, entryName);
String targetPath = targetFile.getCanonicalPath();
if (!targetPath.equals(basePath) && !targetPath.startsWith(basePath + File.separator)) {
throw new IOException(String.format("zip entry is outside of target dir: %s", entryName));
}
return targetFile;
}

private static String getCanonicalDirPath(File dir) throws IOException {
makeDirs(dir);
return dir.getCanonicalPath();
}

private static void makeDirs(File dir) throws IOException {
if (dir != null && !dir.exists() && !dir.mkdirs()) {
throw new IOException(String.format("failed to create directory: %s", dir));
}
}

private static void validateEntrySize(org.apache.tools.zip.ZipEntry entry, long written) throws IOException {
if (written > MAX_ZIP_ENTRY_UNCOMPRESSED_SIZE) {
throw new IOException(String.format("zip entry size exceeds limit: %s", entry.getName()));
}
long compressedSize = entry.getCompressedSize();
if (compressedSize > 0 && written > compressedSize * MAX_ZIP_COMPRESSION_RATIO) {
throw new IOException(String.format("zip entry compression ratio exceeds limit: %s", entry.getName()));
}
}

private static class UnzipContext {
private int entryCount;
private long totalUncompressedSize;

private void addEntry(String entryName) throws IOException {
entryCount++;
if (entryCount > MAX_ZIP_ENTRY_COUNT) {
throw new IOException(String.format("zip entry count exceeds limit: %s", entryName));
}
}

private void addUncompressedSize(long size) throws IOException {
totalUncompressedSize += size;
if (totalUncompressedSize > MAX_ZIP_TOTAL_UNCOMPRESSED_SIZE) {
throw new IOException("zip total uncompressed size exceeds limit");
}
}
}

/**
* 对临时生成的文件夹和文件夹下的文件进行删除
*/
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,141 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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 com.dtstack.taier.common.util;

import com.dtstack.taier.common.exception.TaierDefineException;
import org.junit.Assert;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.List;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

public class ZipUtilTest {

@Rule
public TemporaryFolder temporaryFolder = new TemporaryFolder();

@Test
public void testUpzipFile() throws Exception {
File zipFile = temporaryFolder.newFile("normal.zip");
writeZip(zipFile, new ZipItem("conf/core-site.xml", "content"));
File targetDir = temporaryFolder.newFolder("normal");

List<File> files = ZipUtil.upzipFile(zipFile, targetDir.getAbsolutePath());

Assert.assertEquals(1, files.size());
File extractedFile = new File(targetDir, "conf/core-site.xml");
Assert.assertTrue(extractedFile.isFile());
Assert.assertEquals("content", new String(Files.readAllBytes(extractedFile.toPath()), StandardCharsets.UTF_8));
}

@Test
public void testRejectZipSlipEntry() throws Exception {
File zipFile = temporaryFolder.newFile("slip.zip");
writeZip(zipFile, new ZipItem("../evil.txt", "evil"));
File targetDir = temporaryFolder.newFolder("slip");
File escapedFile = new File(targetDir.getParentFile(), "evil.txt");

try {
ZipUtil.upzipFile(zipFile, targetDir.getAbsolutePath());
Assert.fail("Zip Slip entry should be rejected");
} catch (TaierDefineException e) {
Assert.assertTrue(e.getMessage().contains("outside of target dir"));
}
Assert.assertFalse(escapedFile.exists());
}

@Test
public void testRejectTooManyEntries() throws Exception {
File zipFile = temporaryFolder.newFile("too-many.zip");
writeZip(zipFile, buildZipItems(1001));
File targetDir = temporaryFolder.newFolder("too-many");

try {
ZipUtil.upzipFile(zipFile, targetDir.getAbsolutePath());
Assert.fail("Zip with too many entries should be rejected");
} catch (TaierDefineException e) {
Assert.assertTrue(e.getMessage().contains("entry count exceeds limit"));
}
}

@Test
public void testRejectRecursiveZipBomb() throws Exception {
File zipFile = temporaryFolder.newFile("nested.zip");
writeNestedZip(zipFile, 4);
File targetDir = temporaryFolder.newFolder("nested");

try {
ZipUtil.upzipFile(zipFile, targetDir.getAbsolutePath());
Assert.fail("Recursive zip should be rejected");
} catch (TaierDefineException e) {
Assert.assertTrue(e.getMessage().contains("recursion depth exceeds limit"));
}
}

private static ZipItem[] buildZipItems(int count) {
ZipItem[] items = new ZipItem[count];
for (int i = 0; i < count; i++) {
items[i] = new ZipItem("file-" + i + ".txt", "a");
}
return items;
}

private static void writeNestedZip(File zipFile, int depth) throws IOException {
if (depth == 0) {
writeZip(zipFile, new ZipItem("leaf.txt", "leaf"));
return;
}
File innerZip = File.createTempFile("inner", ".zip", zipFile.getParentFile());
writeNestedZip(innerZip, depth - 1);
try (ZipOutputStream zipOutputStream = new ZipOutputStream(new FileOutputStream(zipFile))) {
zipOutputStream.putNextEntry(new ZipEntry("inner-" + depth + ".zip"));
Files.copy(innerZip.toPath(), zipOutputStream);
zipOutputStream.closeEntry();
}
Files.delete(innerZip.toPath());
}

private static void writeZip(File zipFile, ZipItem... items) throws IOException {
try (ZipOutputStream zipOutputStream = new ZipOutputStream(new FileOutputStream(zipFile))) {
for (ZipItem item : items) {
zipOutputStream.putNextEntry(new ZipEntry(item.name));
zipOutputStream.write(item.content.getBytes(StandardCharsets.UTF_8));
zipOutputStream.closeEntry();
}
}
}

private static class ZipItem {
private final String name;
private final String content;

private ZipItem(String name, String content) {
this.name = name;
this.content = content;
}
}
}
Loading
Loading