From 12f2529be50a7e4bd3ed3ed184f94289463049c9 Mon Sep 17 00:00:00 2001
From: Andy Wilkinson <wilkinsona@vmware.com>
Date: Tue, 17 Nov 2020 11:27:45 +0000
Subject: [PATCH] Allow actuator endpoint to have responses with custom HTTP
 statuses

Fixes gh-24123
---
 ...AbstractWebFluxEndpointHandlerMapping.java |  3 +-
 .../AbstractWebMvcEndpointHandlerMapping.java |  2 +-
 .../AbstractWebEndpointIntegrationTests.java  | 29 ++++++++++++++++++-
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/reactive/AbstractWebFluxEndpointHandlerMapping.java b/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/reactive/AbstractWebFluxEndpointHandlerMapping.java
index 64b073f14f58..29bf30fbbd70 100644
--- a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/reactive/AbstractWebFluxEndpointHandlerMapping.java
+++ b/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/reactive/AbstractWebFluxEndpointHandlerMapping.java
@@ -361,8 +361,7 @@ private ResponseEntity<Object> toResponseEntity(Object response) {
 				return new ResponseEntity<>(response, HttpStatus.OK);
 			}
 			WebEndpointResponse<?> webEndpointResponse = (WebEndpointResponse<?>) response;
-			return new ResponseEntity<>(webEndpointResponse.getBody(),
-					HttpStatus.valueOf(webEndpointResponse.getStatus()));
+			return ResponseEntity.status(webEndpointResponse.getStatus()).body(webEndpointResponse.getBody());
 		}
 
 		@Override
diff --git a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/servlet/AbstractWebMvcEndpointHandlerMapping.java b/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/servlet/AbstractWebMvcEndpointHandlerMapping.java
index 372077b737a9..a14371baf75c 100644
--- a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/servlet/AbstractWebMvcEndpointHandlerMapping.java
+++ b/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/web/servlet/AbstractWebMvcEndpointHandlerMapping.java
@@ -367,7 +367,7 @@ private Object handleResult(Object result, HttpMethod httpMethod) {
 				return result;
 			}
 			WebEndpointResponse<?> response = (WebEndpointResponse<?>) result;
-			return new ResponseEntity<Object>(response.getBody(), HttpStatus.valueOf(response.getStatus()));
+			return ResponseEntity.status(response.getStatus()).body(response.getBody());
 		}
 
 	}
diff --git a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/web/annotation/AbstractWebEndpointIntegrationTests.java b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/web/annotation/AbstractWebEndpointIntegrationTests.java
index dcef4df5b5b7..9fa8f4586fa2 100644
--- a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/web/annotation/AbstractWebEndpointIntegrationTests.java
+++ b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/web/annotation/AbstractWebEndpointIntegrationTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2012-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -379,6 +379,12 @@ void userInRoleReturnsTrueWhenUserIsInRole() {
 				.expectStatus().isOk().expectBody(String.class).isEqualTo("ACTUATOR: true"));
 	}
 
+	@Test
+	void endpointCanProduceAResponseWithACustomStatus() {
+		load((context) -> context.register(CustomResponseStatusEndpointConfiguration.class),
+				(client) -> client.get().uri("/customstatus").exchange().expectStatus().isEqualTo(234));
+	}
+
 	protected abstract int getPort(T context);
 
 	protected void validateErrorBody(WebTestClient.BodyContentSpec body, HttpStatus status, String path,
@@ -621,6 +627,17 @@ UserInRoleEndpoint userInRoleEndpoint() {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	@Import(BaseConfiguration.class)
+	static class CustomResponseStatusEndpointConfiguration {
+
+		@Bean
+		CustomResponseStatusEndpoint customResponseStatusEndpoint() {
+			return new CustomResponseStatusEndpoint();
+		}
+
+	}
+
 	@Endpoint(id = "test")
 	static class TestEndpoint {
 
@@ -847,6 +864,16 @@ String read(SecurityContext securityContext, String role) {
 
 	}
 
+	@Endpoint(id = "customstatus")
+	static class CustomResponseStatusEndpoint {
+
+		@ReadOperation
+		WebEndpointResponse<String> read() {
+			return new WebEndpointResponse<>("Custom status", 234);
+		}
+
+	}
+
 	interface EndpointDelegate {
 
 		void write();