diff --git a/backend/framework/sdk/src/main/java/io/metersphere/sdk/config/TOTPConfig.java b/backend/framework/sdk/src/main/java/io/metersphere/sdk/config/TOTPConfig.java new file mode 100644 index 0000000000..d43227e026 --- /dev/null +++ b/backend/framework/sdk/src/main/java/io/metersphere/sdk/config/TOTPConfig.java @@ -0,0 +1,25 @@ +package io.metersphere.sdk.config; + +import com.bastiaanjansen.otp.HMACAlgorithm; +import com.bastiaanjansen.otp.TOTPGenerator; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +import java.nio.charset.StandardCharsets; + +@Configuration +public class TOTPConfig { + @Value("${totp.secret:secret}") + private String secret; + + @Bean + public TOTPGenerator totpGenerator() { + return new TOTPGenerator.Builder(secret.getBytes(StandardCharsets.UTF_8)) + .withHOTPGenerator(builder -> { + builder.withPasswordLength(6); + builder.withAlgorithm(HMACAlgorithm.SHA256); // SHA256 and SHA512 are also supported + }) + .build(); + } +} diff --git a/backend/framework/sdk/src/main/java/io/metersphere/sdk/constants/MsHttpHeaders.java b/backend/framework/sdk/src/main/java/io/metersphere/sdk/constants/MsHttpHeaders.java new file mode 100644 index 0000000000..1987687139 --- /dev/null +++ b/backend/framework/sdk/src/main/java/io/metersphere/sdk/constants/MsHttpHeaders.java @@ -0,0 +1,5 @@ +package io.metersphere.sdk.constants; + +public interface MsHttpHeaders { + String OTP_TOKEN = "otp-token"; +} diff --git a/backend/framework/sdk/src/main/java/io/metersphere/sdk/service/NodeResourcePoolService.java b/backend/framework/sdk/src/main/java/io/metersphere/sdk/service/NodeResourcePoolService.java index f91175a58c..302e620c9d 100644 --- a/backend/framework/sdk/src/main/java/io/metersphere/sdk/service/NodeResourcePoolService.java +++ b/backend/framework/sdk/src/main/java/io/metersphere/sdk/service/NodeResourcePoolService.java @@ -1,15 +1,21 @@ package io.metersphere.sdk.service; +import com.bastiaanjansen.otp.TOTPGenerator; +import io.metersphere.sdk.constants.MsHttpHeaders; import io.metersphere.sdk.controller.handler.ResultHolder; import io.metersphere.sdk.dto.TestResourceDTO; import io.metersphere.sdk.dto.TestResourceNodeDTO; import io.metersphere.sdk.exception.MSException; import io.metersphere.sdk.util.LogUtils; import io.metersphere.sdk.util.Translator; +import jakarta.annotation.Resource; import org.apache.commons.collections.CollectionUtils; import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.tuple.ImmutablePair; +import org.springframework.http.HttpEntity; +import org.springframework.http.HttpHeaders; +import org.springframework.http.HttpMethod; import org.springframework.http.ResponseEntity; import org.springframework.http.client.HttpComponentsClientHttpRequestFactory; import org.springframework.stereotype.Service; @@ -21,10 +27,13 @@ import java.util.List; @Service @Transactional(rollbackFor = Exception.class) public class NodeResourcePoolService { + @Resource + private TOTPGenerator totpGenerator; private final static String nodeControllerUrl = "http://%s:%s/status"; private static final RestTemplate restTemplateWithTimeOut = new RestTemplate(); + static { HttpComponentsClientHttpRequestFactory httpRequestFactory = new HttpComponentsClientHttpRequestFactory(); httpRequestFactory.setConnectionRequestTimeout(2000); @@ -55,7 +64,7 @@ public class NodeResourcePoolService { } isValid = validateNode(testResourceNodeDTO); if (!isValid) { - break; + break; } } //校验节点 @@ -73,7 +82,11 @@ public class NodeResourcePoolService { private boolean validateNode(TestResourceNodeDTO node) { try { - ResponseEntity entity = restTemplateWithTimeOut.getForEntity(String.format(nodeControllerUrl, node.getIp(), node.getPort()), ResultHolder.class); + String token = totpGenerator.now(); + HttpHeaders headers = new HttpHeaders(); + headers.add(MsHttpHeaders.OTP_TOKEN, token); + HttpEntity httpEntity = new HttpEntity<>(headers); + ResponseEntity entity = restTemplateWithTimeOut.exchange(String.format(nodeControllerUrl, node.getIp(), node.getPort()), HttpMethod.GET, httpEntity, ResultHolder.class); ResultHolder body = entity.getBody(); if (body == null) { return false;