org.springframework.core.io.buffer.DataBuffer Java Examples

@Test
public void toEntityList() {
	DefaultDataBufferFactory factory = new DefaultDataBufferFactory();
	DefaultDataBuffer dataBuffer =
			factory.wrap(ByteBuffer.wrap("foo".getBytes(StandardCharsets.UTF_8)));
	Flux<DataBuffer> body = Flux.just(dataBuffer);
	mockTextPlainResponse(body);

	List<HttpMessageReader<?>> messageReaders = Collections
			.singletonList(new DecoderHttpMessageReader<>(StringDecoder.allMimeTypes()));
	given(mockExchangeStrategies.messageReaders()).willReturn(messageReaders);

	ResponseEntity<List<String>> result = defaultClientResponse.toEntityList(String.class).block();
	assertEquals(Collections.singletonList("foo"), result.getBody());
	assertEquals(HttpStatus.OK, result.getStatusCode());
	assertEquals(HttpStatus.OK.value(), result.getStatusCodeValue());
	assertEquals(MediaType.TEXT_PLAIN, result.getHeaders().getContentType());
}
 

@Override
public Mono<Void> write(Publisher<? extends MultiValueMap<String, String>> inputStream,
		ResolvableType elementType, @Nullable MediaType mediaType, ReactiveHttpOutputMessage message,
		Map<String, Object> hints) {

	mediaType = getMediaType(mediaType);
	message.getHeaders().setContentType(mediaType);

	Charset charset = mediaType.getCharset() != null ? mediaType.getCharset() : getDefaultCharset();

	return Mono.from(inputStream).flatMap(form -> {
		logFormData(form, hints);
		String value = serializeForm(form, charset);
		ByteBuffer byteBuffer = charset.encode(value);
		DataBuffer buffer = message.bufferFactory().wrap(byteBuffer); // wrapping only, no allocation
		message.getHeaders().setContentLength(byteBuffer.remaining());
		return message.writeWith(Mono.just(buffer));
	});
}
 

@Override
public Flux<Object> decode(Publisher<DataBuffer> input, ResolvableType elementType,
		@Nullable MimeType mimeType, @Nullable Map<String, Object> hints) {

	Flux<TokenBuffer> tokens = Jackson2Tokenizer.tokenize(
			Flux.from(input), this.jsonFactory, getObjectMapper(), true);

	ObjectReader reader = getObjectReader(elementType, hints);

	return tokens.handle((tokenBuffer, sink) -> {
		try {
			Object value = reader.readValue(tokenBuffer.asParser(getObjectMapper()));
			logValue(value, hints);
			if (value != null) {
				sink.next(value);
			}
		}
		catch (IOException ex) {
			sink.error(processException(ex));
		}
	});
}
 

@Test
public void limitResponseSize() {
	DefaultDataBufferFactory bufferFactory = new DefaultDataBufferFactory();
	DataBuffer b1 = dataBuffer("foo", bufferFactory);
	DataBuffer b2 = dataBuffer("bar", bufferFactory);
	DataBuffer b3 = dataBuffer("baz", bufferFactory);

	ClientRequest request = ClientRequest.create(HttpMethod.GET, DEFAULT_URL).build();
	ClientResponse response = ClientResponse.create(HttpStatus.OK).body(Flux.just(b1, b2, b3)).build();

	Mono<ClientResponse> result = ExchangeFilterFunctions.limitResponseSize(5)
			.filter(request, req -> Mono.just(response));

	StepVerifier.create(result.flatMapMany(res -> res.body(BodyExtractors.toDataBuffers())))
			.consumeNextWith(buffer -> assertEquals("foo", string(buffer)))
			.consumeNextWith(buffer -> assertEquals("ba", string(buffer)))
			.expectComplete()
			.verify();

}
 

private Mono<Void> sendDeferResponse(ServerHttpResponse serverHttpResponse, ExecutionResult executionResult, Publisher<DeferredExecutionResult> deferredResults) {
    // this implements this apollo defer spec: https://github.com/apollographql/apollo-server/blob/defer-support/docs/source/defer-support.md
    // the spec says CRLF + "-----" + CRLF is needed at the end, but it works without it and with it we get client
    // side errors with it, so we skp it
    serverHttpResponse.setStatusCode(HttpStatus.OK);
    HttpHeaders headers = serverHttpResponse.getHeaders();
    headers.set("Content-Type", "multipart/mixed; boundary=\"-\"");
    headers.set("Connection", "keep-alive");

    Flux<Mono<DataBuffer>> deferredDataBuffers = Flux.from(deferredResults).map(deferredExecutionResult -> {
        DeferPart deferPart = new DeferPart(deferredExecutionResult.toSpecification());
        StringBuilder builder = new StringBuilder();
        String body = deferPart.write();
        builder.append(CRLF).append("---").append(CRLF);
        builder.append(body);
        return strToDataBuffer(builder.toString());
    });
    Flux<Mono<DataBuffer>> firstResult = Flux.just(firstResult(executionResult));


    return serverHttpResponse.writeAndFlushWith(Flux.mergeSequential(firstResult, deferredDataBuffers));
}
 

protected Mono<Void> handleAuthenticationFailure(ServerWebExchange exchange, String errorMsg) {
    CommonResponseDto responseDto = CommonResponseDto.error(errorMsg);
    ServerHttpResponse response = exchange.getResponse();
    try {
        byte[] bits = objectMapper.writeValueAsBytes(responseDto);
        DataBuffer buffer = response.bufferFactory().wrap(bits);
        response.setStatusCode(HttpStatus.UNAUTHORIZED);
        response.getHeaders().add("Content-Type", "application/json;charset=UTF-8");
        response.getHeaders().add(HttpHeaders.WWW_AUTHENTICATE, headerValue);

        return response.writeWith(Mono.just(buffer));
    } catch (JsonProcessingException e) {
        log.debug("failed to process json", e);
        response.setStatusCode(HttpStatus.INTERNAL_SERVER_ERROR);
        return response.setComplete();
    }
}
 

@Test
public void toFormData() {
	DefaultDataBufferFactory factory = new DefaultDataBufferFactory();
	String text = "name+1=value+1&name+2=value+2%2B1&name+2=value+2%2B2&name+3";
	DefaultDataBuffer dataBuffer = factory.wrap(ByteBuffer.wrap(text.getBytes(StandardCharsets.UTF_8)));
	Flux<DataBuffer> body = Flux.just(dataBuffer);

	MockServerHttpRequest request = MockServerHttpRequest.post("/")
			.contentType(MediaType.APPLICATION_FORM_URLENCODED)
			.body(body);

	Mono<MultiValueMap<String, String>> result = BodyExtractors.toFormData().extract(request, this.context);

	StepVerifier.create(result)
			.consumeNextWith(form -> {
				assertEquals("Invalid result", 3, form.size());
				assertEquals("Invalid result", "value 1", form.getFirst("name 1"));
				List<String> values = form.get("name 2");
				assertEquals("Invalid result", 2, values.size());
				assertEquals("Invalid result", "value 2+1", values.get(0));
				assertEquals("Invalid result", "value 2+2", values.get(1));
				assertNull("Invalid result", form.getFirst("name 3"));
			})
			.expectComplete()
			.verify();
}
 

private Mono<Void> writeMultipart(
		MultiValueMap<String, ?> map, ReactiveHttpOutputMessage outputMessage, Map<String, Object> hints) {

	byte[] boundary = generateMultipartBoundary();

	Map<String, String> params = new HashMap<>(2);
	params.put("boundary", new String(boundary, StandardCharsets.US_ASCII));
	params.put("charset", getCharset().name());

	outputMessage.getHeaders().setContentType(new MediaType(MediaType.MULTIPART_FORM_DATA, params));

	LogFormatUtils.traceDebug(logger, traceOn -> Hints.getLogPrefix(hints) + "Encoding " +
			(isEnableLoggingRequestDetails() ?
					LogFormatUtils.formatValue(map, !traceOn) :
					"parts " + map.keySet() + " (content masked)"));

	Flux<DataBuffer> body = Flux.fromIterable(map.entrySet())
			.concatMap(entry -> encodePartValues(boundary, entry.getKey(), entry.getValue()))
			.concatWith(Mono.just(generateLastLine(boundary)));

	return outputMessage.writeWith(body);
}
 

@Override
public Mono<ClientHttpResponse> connect(HttpMethod method, URI uri,
		Function<? super ClientHttpRequest, Mono<Void>> requestCallback) {

	if (!uri.isAbsolute()) {
		return Mono.error(new IllegalArgumentException("URI is not absolute: " + uri));
	}

	if (!this.httpClient.isStarted()) {
		try {
			this.httpClient.start();
		}
		catch (Exception ex) {
			return Mono.error(ex);
		}
	}

	JettyClientHttpRequest clientHttpRequest = new JettyClientHttpRequest(
			this.httpClient.newRequest(uri).method(method.toString()), this.bufferFactory);

	return requestCallback.apply(clientHttpRequest).then(Mono.from(
			clientHttpRequest.getReactiveRequest().response((response, chunks) -> {
				Flux<DataBuffer> content = Flux.from(chunks).map(this::toDataBuffer);
				return Mono.just(new JettyClientHttpResponse(response, content));
			})));
}
 

@Test
public void toFlux() {
	BodyExtractor<Flux<String>, ReactiveHttpInputMessage> extractor = BodyExtractors.toFlux(String.class);

	DefaultDataBufferFactory factory = new DefaultDataBufferFactory();
	DefaultDataBuffer dataBuffer =
			factory.wrap(ByteBuffer.wrap("foo".getBytes(StandardCharsets.UTF_8)));
	Flux<DataBuffer> body = Flux.just(dataBuffer);

	MockServerHttpRequest request = MockServerHttpRequest.post("/").body(body);
	Flux<String> result = extractor.extract(request, this.context);

	StepVerifier.create(result)
			.expectNext("foo")
			.expectComplete()
			.verify();
}
 

/**
 * Read from the request body InputStream and return a DataBuffer.
 * Invoked only when {@link ServletInputStream#isReady()} returns "true".
 * @return a DataBuffer with data read, or {@link #EOF_BUFFER} if the input
 * stream returned -1, or null if 0 bytes were read.
 */
@Nullable
DataBuffer readFromInputStream() throws IOException {
	int read = this.request.getInputStream().read(this.buffer);
	logBytesRead(read);

	if (read > 0) {
		DataBuffer dataBuffer = this.bufferFactory.allocateBuffer(read);
		dataBuffer.write(this.buffer, 0, read);
		return dataBuffer;
	}

	if (read == -1) {
		return EOF_BUFFER;
	}

	return null;
}
 

@NotNull
@Override
public Mono<Object> decodeToMono(@NotNull Publisher<DataBuffer> inputStream, @NotNull ResolvableType elementType, MimeType mimeType, Map<String, Object> hints) {
    return Mono.from(inputStream).handle((dataBuffer, sink) -> {
        try {
            sink.next(decode(dataBuffer));
        } catch (Exception e) {
            sink.error(e);
        }
    });
}
 

@Test
public void invalidData() {
	Flux<DataBuffer> input =
			Flux.from(stringBuffer("{\"foofoo\": \"foofoo\", \"barbar\": \"barbar\""));
	testDecode(input, Pojo.class, step -> step
			.verifyError(DecodingException.class));
}
 

@Override
public void onNext(final ClientResponse clientResponse) {
    try {
        // decorate response body
        subscriber.onNext(ClientResponse.from(clientResponse)
                .body(clientResponse.bodyToFlux(DataBuffer.class).subscriberContext(context))
                .build());
    } finally {
        spanDecorators.forEach(spanDecorator ->
                safelyCall(() -> spanDecorator.onResponse(clientRequest, clientResponse, span)));
    }
}
 

private Flux<DataBuffer> toDataBuffers(String s, int length, Charset charset) {
	byte[] bytes = s.getBytes(charset);
	List<byte[]> chunks = new ArrayList<>();
	for (int i = 0; i < bytes.length; i += length) {
		chunks.add(Arrays.copyOfRange(bytes, i, i + length));
	}
	return Flux.fromIterable(chunks)
			.map(chunk -> {
				DataBuffer dataBuffer = this.bufferFactory.allocateBuffer(length);
				dataBuffer.write(chunk, 0, chunk.length);
				return dataBuffer;
			});
}
 

@Override
public Flux<DataBuffer> content() {
	byte[] bytes = this.content.getBytes(getCharset());
	DataBuffer buffer = getBufferFactory().allocateBuffer(bytes.length);
	buffer.write(bytes);
	return Flux.just(buffer);
}
 

@SuppressWarnings("unchecked")
private static <S, T> Function<Flux<DataBuffer>, Flux<DataBuffer>> convertUsing(
		ParameterizedTypeReference<S> sourceType, ParameterizedTypeReference<T> targetType,
		Function<S, T> converterFn) {
	return (input) -> DECODER.decodeToMono(input, ResolvableType.forType(sourceType), null, null)
			.map((body) -> converterFn.apply((S) body)).flatMapMany((output) -> ENCODER.encode(Mono.just(output),
					new DefaultDataBufferFactory(), ResolvableType.forType(targetType), null, null));
}
 

@Test
public void shouldCreateHeaderFromString() {
    DataBuffer value = new DefaultDataBufferFactory().wrap("value".getBytes(StandardCharsets.UTF_8));
    SnowdropHeader header = new SnowdropHeader("key", "value");

    assertThat(header.key()).isEqualTo("key");
    assertThat(header.value()).isEqualTo(value);
}
 

private MockServerHttpRequest(HttpMethod httpMethod, URI uri, @Nullable String contextPath,
		HttpHeaders headers, MultiValueMap<String, HttpCookie> cookies,
		@Nullable InetSocketAddress remoteAddress, @Nullable SslInfo sslInfo,
		Publisher<? extends DataBuffer> body) {

	super(uri, contextPath, headers);
	this.httpMethod = httpMethod;
	this.cookies = cookies;
	this.remoteAddress = remoteAddress;
	this.sslInfo = sslInfo;
	this.body = Flux.from(body);
}
 

@Test
public void toMonoVoidAsClientWithEmptyBody() {
	TestPublisher<DataBuffer> body = TestPublisher.create();

	BodyExtractor<Mono<Void>, ReactiveHttpInputMessage> extractor = BodyExtractors.toMono(Void.class);
	MockClientHttpResponse response = new MockClientHttpResponse(HttpStatus.OK);
	response.setBody(body.flux());

	StepVerifier.create(extractor.extract(response, this.context))
			.then(() -> {
				body.assertWasSubscribed();
				body.complete();
			})
			.verifyComplete();
}
 

protected final Consumer<DataBuffer> expect(Msg msg) {
	return dataBuffer -> {
		try {
			assertEquals(msg, Msg.parseDelimitedFrom(dataBuffer.asInputStream()));

		}
		catch (IOException ex) {
			throw new UncheckedIOException(ex);
		}
		finally {
			DataBufferUtils.release(dataBuffer);
		}
	};
}
 

public Consumer<DataBuffer> pojoConsumer(Pojo expected) {
	return dataBuffer -> {
		try {
			Pojo actual = this.mapper.reader().forType(Pojo.class)
					.readValue(DataBufferTestUtils.dumpBytes(dataBuffer));
			assertEquals(expected, actual);
			release(dataBuffer);
		}
		catch (IOException ex) {
			throw new UncheckedIOException(ex);
		}
	};
}
 

/**
 * Constructor for a WebSocketMessage.
 * <p>See static factory methods in {@link WebSocketSession} or alternatively
 * use {@link WebSocketSession#bufferFactory()} to create the payload and
 * then invoke this constructor.
 */
public WebSocketMessage(Type type, DataBuffer payload) {
	Assert.notNull(type, "'type' must not be null");
	Assert.notNull(payload, "'payload' must not be null");
	this.type = type;
	this.payload = payload;
}
 

private Mono<Void> encodeAndWriteRegions(Publisher<? extends ResourceRegion> publisher,
		@Nullable MediaType mediaType, ReactiveHttpOutputMessage message, Map<String, Object> hints) {

	Flux<DataBuffer> body = this.regionEncoder.encode(
			publisher, message.bufferFactory(), REGION_TYPE, mediaType, hints);

	return message.writeWith(body);
}
 

@Test
public void should_convert_threaddump() {
    LegacyEndpointConverter converter = LegacyEndpointConverters.threaddump();
    assertThat(converter.canConvert("threaddump")).isTrue();
    assertThat(converter.canConvert("foo")).isFalse();

    Flux<DataBuffer> legacyInput = this.read("threaddump-legacy.json");

    Flux<Object> converted = converter.convert(legacyInput).transform(this::unmarshal);
    Flux<Object> expected = this.read("threaddump-expected.json").transform(this::unmarshal);

    StepVerifier.create(Flux.zip(converted, expected))
                .assertNext(t -> assertThat(t.getT1()).isEqualTo(t.getT2()))
                .verifyComplete();
}
 

@Override
public byte[] encode(DataBuffer original) {
	try {
		ByteArrayOutputStream bis = new ByteArrayOutputStream();
		GZIPOutputStream gos = new GZIPOutputStream(bis);
		FileCopyUtils.copy(original.asInputStream(), gos);
		return bis.toByteArray();
	}
	catch (IOException e) {
		throw new IllegalStateException("couldn't encode body to gzip", e);
	}
}
 

@Override
@Nullable
protected DataBuffer read() throws IOException {
	if (this.inputStream.isReady()) {
		DataBuffer dataBuffer = readFromInputStream();
		if (dataBuffer == EOF_BUFFER) {
			// No need to wait for container callback...
			onAllDataRead();
			dataBuffer = null;
		}
		return dataBuffer;
	}
	return null;
}
 

private DataBuffer randomBuffer(int size) {
	byte[] bytes = new byte[size];
	rnd.nextBytes(bytes);
	DataBuffer buffer = dataBufferFactory.allocateBuffer(size);
	buffer.write(bytes);
	return buffer;
}
 

@Override
public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
    String token = exchange.getRequest().getQueryParams().getFirst("token");

    if (token == null || token.isEmpty()) {
        ServerHttpResponse response = exchange.getResponse();

        // 封装错误信息
        Map<String, Object> responseData = Maps.newHashMap();
        responseData.put("code", 401);
        responseData.put("message", "非法请求");
        responseData.put("cause", "Token is empty");

        try {
            // 将信息转换为 JSON
            ObjectMapper objectMapper = new ObjectMapper();
            byte[] data = objectMapper.writeValueAsBytes(responseData);

            // 输出错误信息到页面
            DataBuffer buffer = response.bufferFactory().wrap(data);
            response.setStatusCode(HttpStatus.UNAUTHORIZED);
            response.getHeaders().add("Content-Type", "application/json;charset=UTF-8");
            return response.writeWith(Mono.just(buffer));
        } catch (JsonProcessingException e) {
            log.error("{}", e);
        }
    }

    return chain.filter(exchange);
}
 

@Override
public Mono<T> decodeToMono(Publisher<DataBuffer> input, ResolvableType elementType,
		@Nullable MimeType mimeType, @Nullable Map<String, Object> hints) {

	return DataBufferUtils.join(input)
			.map(buffer -> decodeDataBuffer(buffer, elementType, mimeType, hints));
}