Command Injection in Ballcat Codegen

@@ -1,6 +1,7 @@
 package com.hccake.ballcat.codegen.engine;
 
 import com.hccake.ballcat.codegen.exception.TemplateRenderException;
+import freemarker.core.TemplateClassResolver;
 import freemarker.template.Configuration;
 import freemarker.template.Template;
 import lombok.SneakyThrows;
@@ -23,6 +24,9 @@ public class FreemarkerTemplateEngine implements TemplateEngine {
 	public FreemarkerTemplateEngine() {
 		this.configuration = new Configuration(Configuration.DEFAULT_INCOMPATIBLE_IMPROVEMENTS);
 		configuration.setDefaultEncoding(StandardCharsets.UTF_8.name());
+
+		// 安全处理 https://ackcent.com/blog/in-depth-freemarker-template-injection/
+		configuration.setNewBuiltinClassResolver(TemplateClassResolver.SAFER_RESOLVER);
 	}
 
 	@Override

@@ -2,12 +2,13 @@
 
 import cn.hutool.core.lang.Assert;
 import cn.hutool.core.util.StrUtil;
+import com.hccake.ballcat.codegen.exception.TemplateRenderException;
 
 import java.util.Map;
 
 /**
  * 模板引擎的委托者
- * 
+ *
  * @author hccake
  */
 public class TemplateEngineDelegator {
@@ -24,7 +25,8 @@ public TemplateEngineDelegator(Map<TemplateEngineTypeEnum, TemplateEngine> templ
 	 * @param context 渲染使用的上下文
 	 * @return 渲染完成后的字符串
 	 */
-	public String render(TemplateEngineTypeEnum engineType, String templateContent, Map<String, Object> context) {
+	public String render(TemplateEngineTypeEnum engineType, String templateContent, Map<String, Object> context)
+			throws TemplateRenderException {
 		if (StrUtil.isEmpty(templateContent)) {
 			return StrUtil.EMPTY;
 		}

@@ -1,5 +1,7 @@
 package com.hccake.ballcat.codegen.engine;
 
+import com.hccake.ballcat.codegen.exception.TemplateRenderException;
+
 import java.util.Map;
 
 /**
@@ -21,6 +23,6 @@ public interface TemplateEngine {
 	 * @param context 渲染使用的上下文
 	 * @return 渲染完成后的字符串
 	 */
-	String render(String templateContent, Map<String, Object> context);
+	String render(String templateContent, Map<String, Object> context) throws TemplateRenderException;
 
 }

@@ -4,6 +4,7 @@
 import org.apache.velocity.VelocityContext;
 import org.apache.velocity.app.Velocity;
 import org.apache.velocity.runtime.resource.loader.ClasspathResourceLoader;
+import org.apache.velocity.util.introspection.SecureUberspector;
 import org.springframework.stereotype.Component;
 
 import java.io.StringWriter;
@@ -22,6 +23,8 @@ public class VelocityTemplateEngine implements TemplateEngine {
 		// 设置velocity资源加载器, 保留文件加载loader
 		Properties prop = new Properties();
 		prop.put("file.resource.loader.class", ClasspathResourceLoader.class.getName());
+		// 安全处理
+		prop.put("runtime.introspector.uberspect", SecureUberspector.class.getName());
 		Velocity.init(prop);
 	}
 
@@ -31,7 +34,7 @@ public TemplateEngineTypeEnum type() {
 	}
 
 	@Override
-	public String render(String templateContent, Map<String, Object> context) {
+	public String render(String templateContent, Map<String, Object> context) throws TemplateRenderException {
 		VelocityContext velocityContext = new VelocityContext(context);
 		try (StringWriter sw = new StringWriter()) {
 			Velocity.evaluate(velocityContext, sw, "velocityTemplateEngine", templateContent);

@@ -5,7 +5,7 @@
  *
  * @author hccake
  */
-public class TemplateRenderException extends RuntimeException {
+public class TemplateRenderException extends Exception {
 
 	public TemplateRenderException(Exception e) {
 		super(e);

@@ -6,6 +6,7 @@
 import com.hccake.ballcat.codegen.constant.TemplateEntryTypeEnum;
 import com.hccake.ballcat.codegen.engine.TemplateEngineDelegator;
 import com.hccake.ballcat.codegen.engine.TemplateEngineTypeEnum;
+import com.hccake.ballcat.codegen.exception.TemplateRenderException;
 import com.hccake.ballcat.codegen.model.bo.FileEntry;
 import com.hccake.ballcat.codegen.model.bo.TableDetails;
 import com.hccake.ballcat.codegen.model.bo.TemplateFile;
@@ -15,6 +16,8 @@
 import com.hccake.ballcat.codegen.service.TableInfoQuery;
 import com.hccake.ballcat.codegen.service.TemplateEntryService;
 import com.hccake.ballcat.codegen.util.GenUtils;
+import com.hccake.ballcat.common.core.exception.BusinessException;
+import com.hccake.ballcat.common.model.result.SystemResultCode;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.stereotype.Service;
@@ -133,22 +136,32 @@ public Map<String, FileEntry> generatorCode(TableDetails tableDetails, String ta
 			fileEntry.setType(templateFile.getType());
 
 			// 替换路径中的占位符
-			String filename = StrUtil.format(templateFile.getFilename(), context);
+			String templateFilename = templateFile.getFilename();
+			String filename = StrUtil.format(templateFilename, context);
 			fileEntry.setFilename(filename);
 
 			String parentFilePath = GenUtils.evaluateRealPath(templateFile.getParentFilePath(), context);
 			fileEntry.setParentFilePath(parentFilePath);
 
 			// 如果是文件
 			if (TemplateEntryTypeEnum.FILE.getType().equals(fileEntry.getType())) {
-				fileEntry.setFilePath(GenUtils.concatFilePath(parentFilePath, filename));
+				String filePath = GenUtils.concatFilePath(parentFilePath, filename);
+				fileEntry.setFilePath(filePath);
 				// 文件内容渲染
 				TemplateEngineTypeEnum engineTypeEnum = TemplateEngineTypeEnum.of(templateFile.getEngineType());
-				String content = templateEngineDelegator.render(engineTypeEnum, templateFile.getContent(), context);
-				fileEntry.setContent(content);
+
+				try {
+					String content = templateEngineDelegator.render(engineTypeEnum, templateFile.getContent(), context);
+					fileEntry.setContent(content);
+				}
+				catch (TemplateRenderException ex) {
+					String errorMessage = StrUtil.format("模板渲染异常，模板文件名：【{}】，错误详情：{}", templateFilename,
+							ex.getMessage());
+					throw new BusinessException(SystemResultCode.SERVER_ERROR.getCode(), errorMessage);
+				}
 			}
 			else {
-				String currentPath = GenUtils.evaluateRealPath(templateFile.getFilename(), context);
+				String currentPath = GenUtils.evaluateRealPath(templateFilename, context);
 				fileEntry.setFilePath(GenUtils.concatFilePath(parentFilePath, currentPath));
 			}