[incubator-servicecomb-java-chassis] 03/04: [SCB-203] set upload parameters to servlet

[li...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

liubao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-servicecomb-java-chassis.git

commit ba4d889c6b5714ee91a811d115039ffad528cc9d
Author: wujimin <wu...@huawei.com>
AuthorDate: Thu Jun 7 12:53:00 2018 +0800

    [SCB-203] set upload parameters to servlet
---
 .../starter/transport/RestServletInitializer.java  |   5 +-
 .../rest/servlet/CseXmlWebApplicationContext.java  |   3 +-
 .../transport/rest/servlet/ServletUtils.java       |  68 ++++++++++--
 .../transport/rest/servlet/TestServletUtils.java   | 117 +++++++++++++++++++++
 4 files changed, 177 insertions(+), 16 deletions(-)

diff --git a/spring-boot-starter/spring-boot-starter-transport/src/main/java/org/apache/servicecomb/springboot/starter/transport/RestServletInitializer.java b/spring-boot-starter/spring-boot-starter-transport/src/main/java/org/apache/servicecomb/springboot/starter/transport/RestServletInitializer.java
index 0055c5b..175f567 100644
--- a/spring-boot-starter/spring-boot-starter-transport/src/main/java/org/apache/servicecomb/springboot/starter/transport/RestServletInitializer.java
+++ b/spring-boot-starter/spring-boot-starter-transport/src/main/java/org/apache/servicecomb/springboot/starter/transport/RestServletInitializer.java
@@ -24,7 +24,6 @@ import javax.servlet.ServletContext;
 import javax.servlet.ServletException;
 
 import org.apache.commons.configuration.Configuration;
-import org.apache.servicecomb.transport.rest.servlet.RestServletInjector;
 import org.apache.servicecomb.transport.rest.servlet.ServletConfig;
 import org.apache.servicecomb.transport.rest.servlet.ServletUtils;
 import org.slf4j.Logger;
@@ -60,8 +59,8 @@ public class RestServletInitializer extends AbstractConfigurableEmbeddedServletC
         Configuration configuration = (Configuration) DynamicPropertyFactory.getBackingConfigurationSource();
         configuration.setProperty(ServletConfig.KEY_SERVLET_URL_PATTERN, ServletConfig.DEFAULT_URL_PATTERN);
       }
-      RestServletInjector.defaultInject(servletContext);
-      ServletUtils.saveUrlPrefix(servletContext);
+
+      ServletUtils.init(servletContext);
     } catch (IOException e) {
       throw new ServletException(e);
     }
diff --git a/transports/transport-rest/transport-rest-servlet/src/main/java/org/apache/servicecomb/transport/rest/servlet/CseXmlWebApplicationContext.java b/transports/transport-rest/transport-rest-servlet/src/main/java/org/apache/servicecomb/transport/rest/servlet/CseXmlWebApplicationContext.java
index ee2fa2c..59918d5 100644
--- a/transports/transport-rest/transport-rest-servlet/src/main/java/org/apache/servicecomb/transport/rest/servlet/CseXmlWebApplicationContext.java
+++ b/transports/transport-rest/transport-rest-servlet/src/main/java/org/apache/servicecomb/transport/rest/servlet/CseXmlWebApplicationContext.java
@@ -53,8 +53,7 @@ public class CseXmlWebApplicationContext extends XmlWebApplicationContext {
     super.invokeBeanFactoryPostProcessors(beanFactory);
 
     // inject servlet after config installed and before transport init
-    RestServletInjector.defaultInject(getServletContext());
-    ServletUtils.saveUrlPrefix(getServletContext());
+    ServletUtils.init(getServletContext());
   }
 
   @Override
diff --git a/transports/transport-rest/transport-rest-servlet/src/main/java/org/apache/servicecomb/transport/rest/servlet/ServletUtils.java b/transports/transport-rest/transport-rest-servlet/src/main/java/org/apache/servicecomb/transport/rest/servlet/ServletUtils.java
index a73b149..94c726f 100644
--- a/transports/transport-rest/transport-rest-servlet/src/main/java/org/apache/servicecomb/transport/rest/servlet/ServletUtils.java
+++ b/transports/transport-rest/transport-rest-servlet/src/main/java/org/apache/servicecomb/transport/rest/servlet/ServletUtils.java
@@ -17,14 +17,18 @@
 
 package org.apache.servicecomb.transport.rest.servlet;
 
+import java.io.File;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.stream.Collectors;
 
+import javax.servlet.MultipartConfigElement;
 import javax.servlet.ServletContext;
 import javax.servlet.ServletRegistration;
+import javax.servlet.ServletRegistration.Dynamic;
 
+import org.apache.servicecomb.common.rest.UploadConfig;
 import org.apache.servicecomb.foundation.common.exceptions.ServiceCombException;
 import org.apache.servicecomb.foundation.common.net.IpPort;
 import org.apache.servicecomb.foundation.common.net.NetUtils;
@@ -75,9 +79,7 @@ public class ServletUtils {
 
   static String[] filterUrlPatterns(Collection<String> urlPatterns) {
     return urlPatterns.stream()
-        .filter(pattern -> {
-          return !pattern.trim().isEmpty();
-        })
+        .filter(pattern -> !pattern.trim().isEmpty())
         .filter(pattern -> {
           checkUrlPattern(pattern.trim());
           return true;
@@ -86,13 +88,7 @@ public class ServletUtils {
   }
 
   static String[] collectUrlPatterns(ServletContext servletContext, Class<?> servletCls) {
-    List<ServletRegistration> servlets = servletContext.getServletRegistrations()
-        .values()
-        .stream()
-        .filter(predicate -> {
-          return predicate.getClassName().equals(servletCls.getName());
-        })
-        .collect(Collectors.toList());
+    List<ServletRegistration> servlets = findServletRegistrations(servletContext, servletCls);
     if (servlets.isEmpty()) {
       return new String[] {};
     }
@@ -108,13 +104,22 @@ public class ServletUtils {
     return filterUrlPatterns(mappings);
   }
 
+  static List<ServletRegistration> findServletRegistrations(ServletContext servletContext,
+      Class<?> servletCls) {
+    return servletContext.getServletRegistrations()
+        .values()
+        .stream()
+        .filter(predicate -> predicate.getClassName().equals(servletCls.getName()))
+        .collect(Collectors.toList());
+  }
+
   static String collectUrlPrefix(ServletContext servletContext, Class<?> servletCls) {
     String[] urlPatterns = collectUrlPatterns(servletContext, servletCls);
     if (urlPatterns.length == 0) {
       return null;
     }
 
-    // even have multiple urlPattern, we only choose a to set as urlPrefix
+    // even have multiple urlPattern, we only choose one to set as urlPrefix
     // only make sure sdk can invoke
     String urlPattern = urlPatterns[0];
     return servletContext.getContextPath() + urlPattern.substring(0, urlPattern.length() - 2);
@@ -130,4 +135,45 @@ public class ServletUtils {
     System.setProperty(Const.URL_PREFIX, urlPrefix);
     LOGGER.info("UrlPrefix of this instance is \"{}\".", urlPrefix);
   }
+
+  static File createUploadDir(ServletContext servletContext, String location) {
+    // If relative, it is relative to TEMPDIR
+    File dir = new File(location);
+    if (!dir.isAbsolute()) {
+      dir = new File((File) servletContext.getAttribute(ServletContext.TEMPDIR), location).getAbsoluteFile();
+    }
+
+    if (!dir.exists()) {
+      dir.mkdirs();
+    }
+
+    return dir;
+  }
+
+  static void setServletParameters(ServletContext servletContext) {
+    UploadConfig uploadConfig = new UploadConfig();
+    MultipartConfigElement multipartConfig = uploadConfig.toMultipartConfigElement();
+    if (multipartConfig == null) {
+      return;
+    }
+
+    File dir = createUploadDir(servletContext, multipartConfig.getLocation());
+    LOGGER.info("set uploads directory to {}.", dir.getAbsolutePath());
+
+    List<ServletRegistration> servlets = findServletRegistrations(servletContext, RestServlet.class);
+    for (ServletRegistration servletRegistration : servlets) {
+      if (!Dynamic.class.isInstance(servletRegistration)) {
+        continue;
+      }
+
+      Dynamic dynamic = (Dynamic) servletRegistration;
+      dynamic.setMultipartConfig(multipartConfig);
+    }
+  }
+
+  public static void init(ServletContext servletContext) {
+    RestServletInjector.defaultInject(servletContext);
+    ServletUtils.saveUrlPrefix(servletContext);
+    setServletParameters(servletContext);
+  }
 }
diff --git a/transports/transport-rest/transport-rest-servlet/src/test/java/org/apache/servicecomb/transport/rest/servlet/TestServletUtils.java b/transports/transport-rest/transport-rest-servlet/src/test/java/org/apache/servicecomb/transport/rest/servlet/TestServletUtils.java
index bf17198..2850694 100644
--- a/transports/transport-rest/transport-rest-servlet/src/test/java/org/apache/servicecomb/transport/rest/servlet/TestServletUtils.java
+++ b/transports/transport-rest/transport-rest-servlet/src/test/java/org/apache/servicecomb/transport/rest/servlet/TestServletUtils.java
@@ -17,25 +17,50 @@
 
 package org.apache.servicecomb.transport.rest.servlet;
 
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 
+import javax.servlet.MultipartConfigElement;
 import javax.servlet.ServletContext;
 import javax.servlet.ServletRegistration;
+import javax.servlet.ServletRegistration.Dynamic;
+import javax.servlet.http.HttpServlet;
 
+import org.apache.servicecomb.common.rest.RestConst;
 import org.apache.servicecomb.foundation.common.exceptions.ServiceCombException;
+import org.apache.servicecomb.foundation.test.scaffolding.config.ArchaiusUtils;
 import org.apache.servicecomb.serviceregistry.api.Const;
 import org.hamcrest.Matchers;
+import org.junit.After;
 import org.junit.Assert;
+import org.junit.Before;
 import org.junit.Test;
 
 import mockit.Expectations;
+import mockit.Mock;
+import mockit.MockUp;
 import mockit.Mocked;
 
 public class TestServletUtils {
+  @Before
+  public void setUp() {
+    ArchaiusUtils.resetConfig();
+  }
+
+  @After
+  public void tearDown() {
+    ArchaiusUtils.resetConfig();
+  }
+
   @Test
   public void testCheckUrlPatternNormal() {
     ServletUtils.checkUrlPattern("/*");
@@ -179,4 +204,96 @@ public class TestServletUtils {
     Assert.assertThat(System.getProperty(Const.URL_PREFIX), Matchers.is("/root/rest"));
     System.clearProperty(Const.URL_PREFIX);
   }
+
+  @Test
+  public void createUploadDir_relative(@Mocked ServletContext servletContext) throws IOException {
+    File tempDir = Files.createTempDirectory("temp").toFile();
+    new Expectations() {
+      {
+        servletContext.getAttribute(ServletContext.TEMPDIR);
+        result = tempDir;
+      }
+    };
+
+    File expectDir = new File(tempDir, "upload");
+    Assert.assertFalse(expectDir.exists());
+
+    File dir = ServletUtils.createUploadDir(servletContext, "upload");
+    Assert.assertTrue(expectDir.exists());
+    Assert.assertEquals(expectDir.getAbsolutePath(), dir.getAbsolutePath());
+
+    dir.delete();
+    Assert.assertFalse(expectDir.exists());
+
+    tempDir.delete();
+    Assert.assertFalse(tempDir.exists());
+  }
+
+  @Test
+  public void createUploadDir_absolute(@Mocked ServletContext servletContext) throws IOException {
+    File tempDir = Files.createTempDirectory("temp").toFile();
+
+    File expectDir = new File(tempDir, "upload");
+    Assert.assertFalse(expectDir.exists());
+
+    File dir = ServletUtils.createUploadDir(servletContext, expectDir.getAbsolutePath());
+    Assert.assertTrue(expectDir.exists());
+    Assert.assertEquals(expectDir.getAbsolutePath(), dir.getAbsolutePath());
+
+    dir.delete();
+    Assert.assertFalse(expectDir.exists());
+
+    tempDir.delete();
+    Assert.assertFalse(tempDir.exists());
+  }
+
+  @Test
+  public void setServletParameters_notSupportUpload() {
+    // not support upload will not set parameters to servlet, so servletContext is null will not throw exception
+    ServletUtils.setServletParameters(null);
+  }
+
+  @Test
+  public void setServletParameters_supportUpload(@Mocked ServletContext servletContext, @Mocked Dynamic d1,
+      @Mocked ServletRegistration d2) throws IOException {
+    Map<String, ServletRegistration> servletRegistrations = new HashMap<>();
+    servletRegistrations.put("d1", d1);
+    servletRegistrations.put("d2", d2);
+    new Expectations() {
+      {
+        servletContext.getServletRegistrations();
+        result = servletRegistrations;
+        d1.getClassName();
+        result = RestServlet.class.getName();
+        d2.getClassName();
+        result = HttpServlet.class.getName();
+      }
+    };
+
+    List<MultipartConfigElement> multipartConfigs = new ArrayList<>();
+    new MockUp<Dynamic>(d1) {
+      @Mock
+      void setMultipartConfig(MultipartConfigElement multipartConfig) {
+        multipartConfigs.add(multipartConfig);
+      }
+    };
+
+    File tempDir = Files.createTempDirectory("temp").toFile();
+    File uploadDir = new File(tempDir, "upload");
+    ArchaiusUtils.setProperty(RestConst.UPLOAD_DIR, uploadDir.getAbsolutePath());
+
+    ServletUtils.setServletParameters(servletContext);
+
+    Assert.assertEquals(1, multipartConfigs.size());
+
+    MultipartConfigElement multipartConfigElement = multipartConfigs.get(0);
+    Assert.assertEquals(uploadDir.getAbsolutePath(), multipartConfigElement.getLocation());
+    Assert.assertEquals(-1, multipartConfigElement.getMaxFileSize());
+    Assert.assertEquals(-1, multipartConfigElement.getMaxRequestSize());
+    Assert.assertEquals(0, multipartConfigElement.getFileSizeThreshold());
+
+    uploadDir.delete();
+    tempDir.delete();
+    Assert.assertFalse(tempDir.exists());
+  }
 }

-- 
To stop receiving notification emails like this one, please contact
liubao@apache.org.