请求代理转发(四)

请求代理转发(四)

使用 webclient 代理请求和支持文件上传

代理 Api

@Slf4j
@RestController
public class ProxyApi {


    private WebClient.Builder builder = WebClient.builder();

    public WebClient client() {
        return Webclient.create();
    }

    private static final String CLIENT_RESPONSE_ATTR = "CLIENT_RESPONSE_ATTR";

    private static final String PERCENTAGE_SIGN = "%";

    //目标服务器
    @Value("${hq.host:http://localhost:7777}")
    private String host;

    private static final String prefix = "/proxy";

    @RequestMapping(prefix + "/**")
    public Mono<Void> proxy(ServerWebExchange exchange) {
            URI requestUrl = rebuildURI(exchange);
            log.info(">>>>> proxy url: {}", requestUrl);
            ServerHttpRequest request = exchange.getRequest();
            HttpMethod method = request.getMethod();
            WebClient.RequestBodySpec bodySpec = client().method(method).uri(requestUrl).headers(httpHeaders -> {
                httpHeaders.addAll(exchange.getRequest().getHeaders());
                httpHeaders.remove("Host");
            });
            if (isUploadRequest(request)) {
                //处理上传
                return exchange.getMultipartData().flatMap(mmp -> {
                    WebClient.RequestHeadersSpec<?> headersSpec = bodySpec.body(BodyInserters.fromMultipartData(mmp));
                    return doProxy(exchange, headersSpec);
                });
            }
            WebClient.RequestHeadersSpec<?> headersSpec;
            if (requiresBody(method)) {
                headersSpec = bodySpec.body(BodyInserters.fromDataBuffers(request.getBody()));
            } else {
                headersSpec = bodySpec;
            }
            return doProxy(exchange, headersSpec);
        
    }

    private static boolean isUploadRequest(ServerHttpRequest request) {
        return request.getHeaders().getContentType() != null && request.getHeaders().getContentType().equalsTypeAndSubtype(MediaType.MULTIPART_FORM_DATA);
    }

    public Mono<? extends Void> doProxy(ServerWebExchange exchange, WebClient.RequestHeadersSpec<?> headersSpec) {
        return headersSpec.exchange().log().flatMap(clientResponse -> {
            ServerHttpResponse response = exchange.getResponse();
            response.setStatusCode(clientResponse.statusCode());
            response.getHeaders().setContentType(MediaType.APPLICATION_JSON_UTF8);
            exchange.getAttributes().put(CLIENT_RESPONSE_ATTR, clientResponse);
            return Mono.justOrEmpty(1);
        }).flatMap(a -> {
            ClientResponse clientResponse = exchange.getAttribute(CLIENT_RESPONSE_ATTR);
            if (clientResponse == null) {
                return Mono.empty();
            }
            ServerHttpResponse response = exchange.getResponse();
            return clientResponse.bodyToMono(String.class).flatMap(str -> {
                try {
                    log.info(">>>>>>>> response {}", str);
                    //自定义返回json
                    DataBuffer wrap = response.bufferFactory().wrap(str.getBytes(StandardCharsets.UTF_8));
                    return response.writeWith(Mono.just(wrap));
                } catch (Exception e) {
                    return Mono.error(e);
                }
            });
        });
    }

    private URI rebuildURI(ServerWebExchange exchange) {
        URI original = exchange.getRequest().getURI();
        String path = path(prefix, exchange);
        URI uri = URI.create(host);
        //这个方式可以 copy get 请求的参数
        return UriComponentsBuilder.fromUri(original).scheme(uri.getScheme()).host(uri.getHost())
                .port(uri.getPort()).replacePath(path).build(containsEncodedParts(original)).toUri();
    }

    /**
     * @param prefix
     * @param exchange
     * @return
     */
    private String path(String prefix, ServerWebExchange exchange) {
        //获取 完整的请求 url
        String path = exchange.getRequest().getURI().getPath();
        return path.substring(prefix.length());
    }

    private boolean requiresBody(HttpMethod method) {
        switch (method) {
            case PUT:
            case POST:
            case PATCH:
                return true;
            default:
                return false;
        }
    }


    public static boolean containsEncodedParts(URI uri) {
        boolean encoded = (uri.getRawQuery() != null && uri.getRawQuery().contains(PERCENTAGE_SIGN))
                || (uri.getRawPath() != null && uri.getRawPath().contains(PERCENTAGE_SIGN));
        if (encoded) {
            try {
                UriComponentsBuilder.fromUri(uri).build(true);
                return true;
            } catch (IllegalArgumentException ignored) {
                if (log.isTraceEnabled()) {
                    log.trace("Error in containsEncodedParts", ignored);
                }
            }
            return false;
        }
        return encoded;
    }
}

服务api:

@PostMapping(value = "/body")
public Mono<R<String>> body(@RequestBody String name) {
    log.info("name: {}", name);
    return Mono.just(R.ok("success"));
}


@PostMapping(value = "/form")
public Mono<R<String>> form(String name) {
    log.info("name: {}", name);
    return Mono.just(R.ok("success"));
}

@RequestMapping(value = "/get")
public Mono<R<String>> get(String name) {
    log.info("name: {}", name);
    return Mono.just(R.ok("success"));
}


@PostMapping(value = "/upload")
public Mono<R<String>> index(@RequestPart("file") FilePart filePart,
		 ServerWebExchange exchange) throws IOException {
    System.out.println(filePart.filename());
    Path file = Files.createFile(Paths.get("test".concat(filePart.filename())));
    filePart.transferTo(file.toFile());
    return Mono.just(R.ok("success"));
}

good luck!