ProxyController

package foo.bar;

import foo.bar.RestTemplateFactory;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.InputStreamResource;
import org.springframework.http.*;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.ResponseBody;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.HttpStatusCodeException;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.multipart.MultipartHttpServletRequest;
import org.springframework.web.util.UriComponentsBuilder;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.URL;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Enumeration;
import java.util.Iterator;

@RestController
@RequestMapping("proxy")
public class ProxyController {

    private Logger log = LoggerFactory.getLogger(ProxyController.class);
    private final RestTemplateFactory restTemplateContextFactory;
    private final TargetUrlBuilder targetUrlBuilder;

    @Autowired
    public ProxyController(RestTemplateFactory restTemplateContextFactory, TargetUrlBuilder targetUrlBuilder) {
        this.restTemplateContextFactory = restTemplateContextFactory;
        this.targetUrlBuilder = targetUrlBuilder;
    }

    @RequestMapping("/**")
    @ResponseBody
    public void proxy(@RequestBody(required = false) byte[] body,
                      HttpMethod method,
                      HttpServletRequest servletRequest,
                      HttpServletResponse servletResponse) throws IOException {

        URL targetUrl = buildTargetUrl(servletRequest);
        URI uri = buildUri(servletRequest, targetUrl);
        HttpHeaders headers = copyHeaders(servletRequest);
        HttpEntity httpEntity = buildHttpEntity(body, servletRequest, headers);

        sendProxyRequest(method, servletResponse, targetUrl, uri, httpEntity);
    }

    private HttpEntity buildHttpEntity(@RequestBody(required = false) byte[] body, HttpServletRequest servletRequest, HttpHeaders headers) throws IOException {
        HttpEntity httpEntity;

        if(servletRequest instanceof MultipartHttpServletRequest) {
            httpEntity = buildHttpEntityFromMultipartRequest(servletRequest, headers);
        } else {
            httpEntity = new HttpEntity<>(body, headers);
        }

        return httpEntity;
    }

    private URL buildTargetUrl(HttpServletRequest servletRequest) {
        URL targetUrl = getTargetUrl(servletRequest);
        if(targetUrl == null) {
            throw new RuntimeException("Failed to build target url.");
        }
        return targetUrl;
    }

    private URI buildUri(HttpServletRequest servletRequest, URL targetUrl) {
        return UriComponentsBuilder
                    .fromUriString(targetUrl.toString())
                    .query(servletRequest.getQueryString())
                    .build(true).toUri();
    }

    private HttpHeaders copyHeaders(HttpServletRequest servletRequest) {
        HttpHeaders headers = new HttpHeaders();
        Enumeration<String> headerNames = servletRequest.getHeaderNames();
        while (headerNames.hasMoreElements()) {
            String headerName = headerNames.nextElement();
            headers.set(headerName, servletRequest.getHeader(headerName));
        }
        return headers;
    }

    private HttpEntity<MultiValueMap<String, Object>> buildHttpEntityFromMultipartRequest(HttpServletRequest servletRequest, HttpHeaders headers) throws IOException {
        MultipartHttpServletRequest multipartRequest = (MultipartHttpServletRequest) servletRequest;
        Iterator<String> itr = multipartRequest.getFileNames();
        String fileName = itr.next();
        MultipartFile file = multipartRequest.getFile(fileName);

        MultiValueMap<String, Object> parts = new LinkedMultiValueMap<>();
        parts.add(fileName, new MultipartInputStreamFileResource(new ByteArrayInputStream(file.getBytes()), file.getOriginalFilename()));

        ArrayList<String> requestParameters = Collections.list(multipartRequest.getParameterNames());
        for (String requestParameter : requestParameters) {
            parts.add(requestParameter, multipartRequest.getParameter(requestParameter));
        }

        headers.setContentType(MediaType.MULTIPART_FORM_DATA);

        return new HttpEntity<>(parts, headers);
    }

    private void sendProxyRequest(HttpMethod method, HttpServletResponse servletResponse, URL targetUrl, URI uri, HttpEntity httpEntity) throws IOException {
        RestTemplate restTemplate = restTemplateContextFactory.build();
        try {
            ResponseEntity<byte[]> exchange = restTemplate.exchange(uri, method, httpEntity, byte[].class);

            int responseStatusCode = exchange.getStatusCodeValue();
            byte[] responseBody = exchange.getBody();
            HttpHeaders responseHeaders = exchange.getHeaders();

            log.debug("{} responded with status code {}", targetUrl, responseStatusCode);

            copyProxyResponseToServletResponse(servletResponse, responseBody, responseHeaders, responseStatusCode);

        } catch(HttpStatusCodeException e) {
            int proxyResponseStatusCode = e.getRawStatusCode();
            byte[] responseBody = e.getResponseBodyAsByteArray();
            HttpHeaders proxyResonseHeaders = e.getResponseHeaders();

            log.debug("FAILED: {} responded with status code {}", targetUrl, proxyResponseStatusCode);

            copyProxyResponseToServletResponse(servletResponse, responseBody, proxyResonseHeaders, proxyResponseStatusCode);
        }
    }

    private void copyProxyResponseToServletResponse(HttpServletResponse response, byte[] bodyString, HttpHeaders proxyResonseHeaders, int proxyResponseStatusCode) throws IOException {
        if(proxyResonseHeaders != null) {
            proxyResonseHeaders.forEach((headerName, value) -> {
                if(shouldCopyResponseHeader(headerName)) {
                    response.addHeader(headerName, String.join(",", value));
                }
            });
        }

        response.setStatus(proxyResponseStatusCode);
        if(bodyString != null) {
            IOUtils.copy(new ByteArrayInputStream(bodyString), response.getOutputStream());
        }
    }

    private boolean shouldCopyResponseHeader(String headerName) {
        return headerName != null
                && !"Transfer-Encoding".equals(headerName);
    }

    private URL getTargetUrl(HttpServletRequest request) {
        try {
            return targetUrlBuilder.buildUrl(request);
        } catch (BlacklistedServiceException e) {
            log.info("Blocking request to blacklisted service '{}'", e.getService());
            return null;
        } catch (Exception e) {
            log.warn("Failed to determine target URL", e);
            return null;
        }
    }

    class MultipartInputStreamFileResource extends InputStreamResource {

        private final String filename;

        MultipartInputStreamFileResource(InputStream inputStream, String filename) {
            super(inputStream);
            this.filename = filename;
        }

        @Override
        public String getFilename() {
            return this.filename;
        }

        @Override
        public long contentLength() throws IOException {
            return -1;
        }
    }

}