archive/MemShellParty
1
0
Fork 0
mirror of https://github.com/ReaJason/MemShellParty.git synced 2026-09-25 00:11:52 +08:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

fix: disconnected

Browse Source
This commit is contained in:
ReaJason
2025-05-28 01:22:54 +08:00
parent dba1b68c23
commit 7b2329de45
1 changed files with 5 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/godzilla/src/main/java/com/reajason/javaweb/godzilla/BlockingJavaWebSocketClient.java
+5 -34
View File
@@ -13,7 +13,7 @@ import java.util.concurrent.atomic.AtomicReference;
public class BlockingJavaWebSocketClient extends WebSocketClient { public class BlockingJavaWebSocketClient extends WebSocketClient {
private CountDownLatch connectLatch = new CountDownLatch(1); private final CountDownLatch connectLatch = new CountDownLatch(1);
private volatile CountDownLatch responseLatch; private volatile CountDownLatch responseLatch;
private final AtomicReference<String> responseMessage = new AtomicReference<>(); private final AtomicReference<String> responseMessage = new AtomicReference<>();
private final AtomicReference<byte[]> responseBytesMessage = new AtomicReference<>(); private final AtomicReference<byte[]> responseBytesMessage = new AtomicReference<>();
@@ -31,7 +31,7 @@ public class BlockingJavaWebSocketClient extends WebSocketClient {
} }
public void onMessage(String message) { public void onMessage(String message) {
System.out.println("收到消息: " + message); System.out.println("收到消息：" + message);
responseMessage.set(message); responseMessage.set(message);
if (responseLatch != null) { if (responseLatch != null) {
responseLatch.countDown(); responseLatch.countDown();
@@ -40,7 +40,7 @@ public class BlockingJavaWebSocketClient extends WebSocketClient {
} }
public void onMessage(ByteBuffer byteBuffer) { public void onMessage(ByteBuffer byteBuffer) {
System.out.println("收到字节消息: " + byteBuffer); System.out.println("收到字节消息：" + byteBuffer);
responseBytesMessage.set(byteBuffer.array()); responseBytesMessage.set(byteBuffer.array());
if (responseLatch != null) { if (responseLatch != null) {
responseLatch.countDown(); responseLatch.countDown();
@@ -49,9 +49,8 @@ public class BlockingJavaWebSocketClient extends WebSocketClient {
} }
public void onClose(int code, String reason, boolean remote) { public void onClose(int code, String reason, boolean remote) {
System.out.println("连接关闭: " + code + " - " + reason); System.out.println("连接关闭：" + code + " - " + reason);
connected = false; connected = false;
// Signal any waiting threads
if (responseLatch != null) { if (responseLatch != null) {
responseLatch.countDown(); responseLatch.countDown();
} }
@@ -59,9 +58,8 @@ public class BlockingJavaWebSocketClient extends WebSocketClient {
} }
public void onError(Exception ex) { public void onError(Exception ex) {
System.out.println("连接错误: " + ex.getMessage()); System.out.println("连接错误：" + ex.getMessage());
connected = false; connected = false;
// Signal any waiting threads
if (responseLatch != null) { if (responseLatch != null) {
responseLatch.countDown(); responseLatch.countDown();
} }
@@ -70,7 +68,6 @@ public class BlockingJavaWebSocketClient extends WebSocketClient {
} }
public String sendRequest(String message) throws InterruptedException { public String sendRequest(String message) throws InterruptedException {
// Connect if not already connected
if (!connected && !isOpen()) { if (!connected && !isOpen()) {
connect(); connect();
if (!connectLatch.await(5, TimeUnit.SECONDS)) { if (!connectLatch.await(5, TimeUnit.SECONDS)) {
@@ -82,29 +79,19 @@ public class BlockingJavaWebSocketClient extends WebSocketClient {
throw new IllegalStateException("WebSocket connection is not open."); throw new IllegalStateException("WebSocket connection is not open.");
} }
// Reset response data and create new response latch for this request
responseMessage.set(null); responseMessage.set(null);
responseBytesMessage.set(null); responseBytesMessage.set(null);
responseLatch = new CountDownLatch(1); responseLatch = new CountDownLatch(1);
// Send the message
send(message); send(message);
// Wait for response
if (!responseLatch.await(10, TimeUnit.SECONDS)) { if (!responseLatch.await(10, TimeUnit.SECONDS)) {
throw new InterruptedException("Timeout waiting for WebSocket response."); throw new InterruptedException("Timeout waiting for WebSocket response.");
} }
// Check if connection was closed during wait
if (!connected) {
throw new IllegalStateException("WebSocket connection was closed while waiting for response.");
}
return responseMessage.get(); return responseMessage.get();
} }
public byte[] sendRequest(ByteBuffer message) throws InterruptedException { public byte[] sendRequest(ByteBuffer message) throws InterruptedException {
// Connect if not already connected
if (!connected && !isOpen()) { if (!connected && !isOpen()) {
connect(); connect();
if (!connectLatch.await(5, TimeUnit.SECONDS)) { if (!connectLatch.await(5, TimeUnit.SECONDS)) {
@@ -116,34 +103,18 @@ public class BlockingJavaWebSocketClient extends WebSocketClient {
throw new IllegalStateException("WebSocket connection is not open."); throw new IllegalStateException("WebSocket connection is not open.");
} }
// Reset response data and create new response latch for this request
responseMessage.set(null); responseMessage.set(null);
responseBytesMessage.set(null); responseBytesMessage.set(null);
responseLatch = new CountDownLatch(1); responseLatch = new CountDownLatch(1);
// Send the message
send(message); send(message);
// Wait for response
if (!responseLatch.await(10, TimeUnit.SECONDS)) { if (!responseLatch.await(10, TimeUnit.SECONDS)) {
throw new InterruptedException("Timeout waiting for WebSocket response."); throw new InterruptedException("Timeout waiting for WebSocket response.");
} }
// Check if connection was closed during wait
if (!connected) {
throw new IllegalStateException("WebSocket connection was closed while waiting for response.");
}
return responseBytesMessage.get(); return responseBytesMessage.get();
} }
public void disconnect() {
if (connected && isOpen()) {
close();
}
}
@SneakyThrows @SneakyThrows
public static String sendRequestWaitResponse(String entrypoint, String message) { public static String sendRequestWaitResponse(String entrypoint, String message) {
BlockingJavaWebSocketClient blockingJavaWebSocketClient = new BlockingJavaWebSocketClient(URI.create(entrypoint)); BlockingJavaWebSocketClient blockingJavaWebSocketClient = new BlockingJavaWebSocketClient(URI.create(entrypoint));