mirror of
https://github.com/alibaba/arthas.git
synced 2026-10-21 05:51:38 +00:00
add quiet terminal mode for agents (#3213)
This commit is contained in:
@@ -47,7 +47,7 @@ public interface CommandExecutor {
|
||||
|
||||
Map<String, Object> interruptJob(String sessionId);
|
||||
|
||||
Map<String, Object> createSession();
|
||||
Map<String, Object> createSession(boolean quiet);
|
||||
|
||||
Map<String, Object> closeSession(String sessionId);
|
||||
|
||||
@@ -61,4 +61,3 @@ public interface CommandExecutor {
|
||||
*/
|
||||
void setSessionUserId(String sessionId, String userId);
|
||||
}
|
||||
|
||||
|
||||
+2
-2
@@ -79,7 +79,7 @@ public class ArthasCommandSessionManager {
|
||||
}
|
||||
|
||||
public CommandSessionBinding createCommandSession(String mcpSessionId) {
|
||||
Map<String, Object> result = commandExecutor.createSession();
|
||||
Map<String, Object> result = commandExecutor.createSession(true);
|
||||
|
||||
CommandSessionBinding binding = new CommandSessionBinding(
|
||||
mcpSessionId,
|
||||
@@ -177,7 +177,7 @@ public class ArthasCommandSessionManager {
|
||||
.build();
|
||||
}
|
||||
|
||||
Map<String, Object> result = commandExecutor.createSession();
|
||||
Map<String, Object> result = commandExecutor.createSession(true);
|
||||
|
||||
CommandSessionBinding binding = new CommandSessionBinding(
|
||||
"task-" + taskId, // 使用 task ID 作为 MCP session ID
|
||||
|
||||
+1
-1
@@ -66,7 +66,7 @@ class ArthasCommandContextAuthTest {
|
||||
}
|
||||
|
||||
@Override
|
||||
public Map<String, Object> createSession() {
|
||||
public Map<String, Object> createSession(boolean quiet) {
|
||||
Map<String, Object> result = new HashMap<String, Object>();
|
||||
result.put("sessionId", "created-session");
|
||||
result.put("consumerId", "created-consumer");
|
||||
|
||||
+72
@@ -0,0 +1,72 @@
|
||||
package com.taobao.arthas.mcp.server.session;
|
||||
|
||||
import com.taobao.arthas.mcp.server.CommandExecutor;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
class ArthasCommandSessionManagerQuietTest {
|
||||
|
||||
@Test
|
||||
void commandSessionsShouldBeCreatedQuietly() {
|
||||
RecordingCommandExecutor executor = new RecordingCommandExecutor();
|
||||
ArthasCommandSessionManager sessionManager = new ArthasCommandSessionManager(executor);
|
||||
|
||||
sessionManager.createCommandSession("mcp-session");
|
||||
sessionManager.createIsolatedTaskSession("task-1");
|
||||
|
||||
assertThat(executor.quietFlags).containsExactly(Boolean.TRUE, Boolean.TRUE);
|
||||
}
|
||||
|
||||
private static final class RecordingCommandExecutor implements CommandExecutor {
|
||||
private final List<Boolean> quietFlags = new ArrayList<Boolean>();
|
||||
|
||||
@Override
|
||||
public Map<String, Object> executeSync(String commandLine, long timeout, String sessionId,
|
||||
Object authSubject, String userId) {
|
||||
return new HashMap<String, Object>();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Map<String, Object> executeAsync(String commandLine, String sessionId) {
|
||||
return new HashMap<String, Object>();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Map<String, Object> pullResults(String sessionId, String consumerId) {
|
||||
return new HashMap<String, Object>();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Map<String, Object> interruptJob(String sessionId) {
|
||||
return new HashMap<String, Object>();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Map<String, Object> createSession(boolean quiet) {
|
||||
quietFlags.add(quiet);
|
||||
Map<String, Object> result = new HashMap<String, Object>();
|
||||
result.put("sessionId", "session-" + quietFlags.size());
|
||||
result.put("consumerId", "consumer-" + quietFlags.size());
|
||||
return result;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Map<String, Object> closeSession(String sessionId) {
|
||||
return new HashMap<String, Object>();
|
||||
}
|
||||
|
||||
@Override
|
||||
public void setSessionAuth(String sessionId, Object authSubject) {
|
||||
}
|
||||
|
||||
@Override
|
||||
public void setSessionUserId(String sessionId, String userId) {
|
||||
}
|
||||
}
|
||||
}
|
||||
+1
-1
@@ -80,7 +80,7 @@ class DefaultCreateTaskContextAuthTest {
|
||||
}
|
||||
|
||||
@Override
|
||||
public Map<String, Object> createSession() {
|
||||
public Map<String, Object> createSession(boolean quiet) {
|
||||
Map<String, Object> result = new HashMap<String, Object>();
|
||||
result.put("sessionId", "isolated-session");
|
||||
result.put("consumerId", "isolated-consumer");
|
||||
|
||||
@@ -61,6 +61,16 @@
|
||||
<groupId>jline</groupId>
|
||||
<artifactId>jline</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.junit.jupiter</groupId>
|
||||
<artifactId>junit-jupiter</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.assertj</groupId>
|
||||
<artifactId>assertj-core</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
</project>
|
||||
|
||||
@@ -84,6 +84,7 @@ public class TelnetConsole {
|
||||
|
||||
private Integer width = null;
|
||||
private Integer height = null;
|
||||
private boolean quiet = false;
|
||||
|
||||
@Argument(argName = "target-ip", index = 0, required = false)
|
||||
@Description("Target ip")
|
||||
@@ -133,6 +134,12 @@ public class TelnetConsole {
|
||||
this.height = height;
|
||||
}
|
||||
|
||||
@Option(longName = "quiet", flag = true)
|
||||
@Description("Suppress connection welcome output")
|
||||
public void setQuiet(boolean quiet) {
|
||||
this.quiet = quiet;
|
||||
}
|
||||
|
||||
public TelnetConsole() {
|
||||
}
|
||||
|
||||
@@ -273,7 +280,9 @@ public class TelnetConsole {
|
||||
}
|
||||
}
|
||||
|
||||
final TelnetClient telnet = new TelnetClient();
|
||||
final TelnetClient telnet = telnetConsole.isQuiet()
|
||||
? new TelnetClient("arthas-agent")
|
||||
: new TelnetClient();
|
||||
telnet.setConnectTimeout(DEFAULT_CONNECTION_TIMEOUT);
|
||||
|
||||
// send init terminal size
|
||||
@@ -369,7 +378,7 @@ public class TelnetConsole {
|
||||
|
||||
// 检查到有 [arthas@ 时,意味着可以执行下一个命令了
|
||||
int index = line.indexOf(PROMPT);
|
||||
if (index > 0) {
|
||||
if (index >= 0) {
|
||||
line.delete(0, index + PROMPT.length());
|
||||
receviedPromptQueue.put("");
|
||||
}
|
||||
@@ -446,6 +455,10 @@ public class TelnetConsole {
|
||||
return height;
|
||||
}
|
||||
|
||||
public boolean isQuiet() {
|
||||
return quiet;
|
||||
}
|
||||
|
||||
public boolean isHelp() {
|
||||
return help;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
package com.taobao.arthas.client;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.io.OutputStream;
|
||||
import java.net.ServerSocket;
|
||||
import java.net.Socket;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.concurrent.Future;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
class TelnetConsoleBatchModeTest {
|
||||
|
||||
@Test
|
||||
void quietBatchModeShouldSendCommandWhenPromptStartsTheStream() throws Exception {
|
||||
ExecutorService executorService = Executors.newSingleThreadExecutor();
|
||||
try (ServerSocket serverSocket = new ServerSocket(0)) {
|
||||
Future<ServerResult> serverResult = executorService.submit(() -> runPromptFirstServer(serverSocket));
|
||||
|
||||
int status = TelnetConsole.process(new String[] {
|
||||
"--quiet",
|
||||
"-c",
|
||||
"version",
|
||||
"-t",
|
||||
"1000",
|
||||
"127.0.0.1",
|
||||
String.valueOf(serverSocket.getLocalPort())
|
||||
});
|
||||
|
||||
ServerResult result = serverResult.get(2, TimeUnit.SECONDS);
|
||||
assertThat(status).isEqualTo(TelnetConsole.STATUS_OK);
|
||||
assertThat(result.command).contains("version | plaintext");
|
||||
assertThat(result.quit).isEqualTo("quit");
|
||||
} finally {
|
||||
executorService.shutdownNow();
|
||||
}
|
||||
}
|
||||
|
||||
private static ServerResult runPromptFirstServer(ServerSocket serverSocket) throws IOException {
|
||||
try (Socket socket = serverSocket.accept()) {
|
||||
socket.setSoTimeout(2000);
|
||||
OutputStream outputStream = socket.getOutputStream();
|
||||
InputStream inputStream = socket.getInputStream();
|
||||
|
||||
write(outputStream, "[arthas@123]$ ");
|
||||
String command = readLine(inputStream);
|
||||
|
||||
write(outputStream, "ok\n[arthas@123]$ ");
|
||||
String quit = readLine(inputStream);
|
||||
|
||||
return new ServerResult(command, quit);
|
||||
}
|
||||
}
|
||||
|
||||
private static void write(OutputStream outputStream, String value) throws IOException {
|
||||
outputStream.write(value.getBytes(StandardCharsets.UTF_8));
|
||||
outputStream.flush();
|
||||
}
|
||||
|
||||
private static String readLine(InputStream inputStream) throws IOException {
|
||||
ByteArrayOutputStream buffer = new ByteArrayOutputStream();
|
||||
int value;
|
||||
while ((value = inputStream.read()) != -1) {
|
||||
if (value == '\n') {
|
||||
break;
|
||||
}
|
||||
buffer.write(value);
|
||||
}
|
||||
return new String(buffer.toByteArray(), StandardCharsets.UTF_8).replace("\r", "");
|
||||
}
|
||||
|
||||
private static class ServerResult {
|
||||
private final String command;
|
||||
private final String quit;
|
||||
|
||||
private ServerResult(String command, String quit) {
|
||||
this.command = command;
|
||||
this.quit = quit;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -279,17 +279,34 @@ public class CommandExecutorImpl implements CommandExecutor {
|
||||
}
|
||||
|
||||
@Override
|
||||
public Map<String, Object> createSession() {
|
||||
public Map<String, Object> createSession(boolean quiet) {
|
||||
Session session = sessionManager.createSession();
|
||||
if (session == null) {
|
||||
return createErrorResult(null, "create api session failed");
|
||||
}
|
||||
if (quiet) {
|
||||
session.put(Session.QUIET, Boolean.TRUE);
|
||||
}
|
||||
|
||||
SharingResultDistributorImpl resultDistributor = new SharingResultDistributorImpl(session);
|
||||
ResultConsumer resultConsumer = new ResultConsumerImpl();
|
||||
resultDistributor.addConsumer(resultConsumer);
|
||||
session.setResultDistributor(resultDistributor);
|
||||
|
||||
if (!quiet) {
|
||||
appendWelcomeResults(resultDistributor);
|
||||
}
|
||||
|
||||
updateSessionInputStatus(session, InputStatus.ALLOW_INPUT);
|
||||
|
||||
Map<String, Object> result = new TreeMap<>();
|
||||
result.put("success", true);
|
||||
result.put("sessionId", session.getSessionId());
|
||||
result.put("consumerId", resultConsumer.getConsumerId());
|
||||
return result;
|
||||
}
|
||||
|
||||
private void appendWelcomeResults(ResultDistributor resultDistributor) {
|
||||
resultDistributor.appendResult(new MessageModel("Welcome to arthas!"));
|
||||
|
||||
WelcomeModel welcomeModel = new WelcomeModel();
|
||||
@@ -300,14 +317,6 @@ public class CommandExecutorImpl implements CommandExecutor {
|
||||
welcomeModel.setPid(PidUtils.currentPid());
|
||||
welcomeModel.setTime(DateUtils.getCurrentDateTime());
|
||||
resultDistributor.appendResult(welcomeModel);
|
||||
|
||||
updateSessionInputStatus(session, InputStatus.ALLOW_INPUT);
|
||||
|
||||
Map<String, Object> result = new TreeMap<>();
|
||||
result.put("success", true);
|
||||
result.put("sessionId", session.getSessionId());
|
||||
result.put("consumerId", resultConsumer.getConsumerId());
|
||||
return result;
|
||||
}
|
||||
|
||||
@Override
|
||||
|
||||
@@ -52,7 +52,7 @@ import javax.security.auth.login.LoginException;
|
||||
*/
|
||||
public class ShellImpl implements Shell {
|
||||
private static final Logger logger = LoggerFactory.getLogger(ShellImpl.class);
|
||||
private SecurityAuthenticator securityAuthenticator = ArthasBootstrap.getInstance().getSecurityAuthenticator();
|
||||
private static final String ARTHAS_AGENT_TERMINAL_TYPE = "arthas-agent";
|
||||
|
||||
private JobControllerImpl jobController;
|
||||
final String id;
|
||||
@@ -78,6 +78,7 @@ public class ShellImpl implements Shell {
|
||||
Principal principal = AuthUtils.localPrincipal(handlerContext);
|
||||
if (principal != null) {
|
||||
try {
|
||||
SecurityAuthenticator securityAuthenticator = ArthasBootstrap.getInstance().getSecurityAuthenticator();
|
||||
Subject subject = securityAuthenticator.login(principal);
|
||||
if (subject != null) {
|
||||
session.put(ArthasConstants.SUBJECT_KEY, subject);
|
||||
@@ -98,6 +99,9 @@ public class ShellImpl implements Shell {
|
||||
}
|
||||
}
|
||||
}
|
||||
if (term != null && ARTHAS_AGENT_TERMINAL_TYPE.equalsIgnoreCase(term.type())) {
|
||||
session.put(Session.QUIET, Boolean.TRUE);
|
||||
}
|
||||
session.put(Session.COMMAND_MANAGER, commandManager);
|
||||
session.put(Session.INSTRUMENTATION, instrumentation);
|
||||
session.put(Session.PID, pid);
|
||||
@@ -168,12 +172,16 @@ public class ShellImpl implements Shell {
|
||||
term.suspendHandler(new SuspendHandler(this));
|
||||
term.closeHandler(new CloseHandler(this));
|
||||
|
||||
if (welcome != null && welcome.length() > 0) {
|
||||
if (!isQuietSession() && welcome != null && welcome.length() > 0) {
|
||||
term.write(welcome + "\n");
|
||||
}
|
||||
return this;
|
||||
}
|
||||
|
||||
private boolean isQuietSession() {
|
||||
return Boolean.TRUE.equals(session.get(Session.QUIET));
|
||||
}
|
||||
|
||||
public String statusLine(Job job, ExecStatus status) {
|
||||
StringBuilder sb = new StringBuilder("[").append(job.id()).append("]");
|
||||
if (this.session().equals(job.getSession())) {
|
||||
|
||||
@@ -20,6 +20,10 @@ public interface Session {
|
||||
String ID = "id";
|
||||
String SERVER = "server";
|
||||
String USER_ID = "userId";
|
||||
/**
|
||||
* 会话静默模式,不输出连接欢迎信息。
|
||||
*/
|
||||
String QUIET = "arthas-session-quiet";
|
||||
/**
|
||||
* The tty this session related to.
|
||||
*/
|
||||
|
||||
+14
-4
@@ -6,6 +6,7 @@ import java.util.Map;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
import com.taobao.arthas.common.ArthasConstants;
|
||||
import com.taobao.arthas.core.shell.session.Session;
|
||||
import com.taobao.arthas.core.shell.term.impl.http.session.HttpSession;
|
||||
import com.taobao.arthas.core.shell.term.impl.http.session.HttpSessionManager;
|
||||
|
||||
@@ -23,9 +24,15 @@ import io.termd.core.http.HttpTtyConnection;
|
||||
*/
|
||||
public class ExtHttpTtyConnection extends HttpTtyConnection {
|
||||
private ChannelHandlerContext context;
|
||||
private final boolean quiet;
|
||||
|
||||
public ExtHttpTtyConnection(ChannelHandlerContext context) {
|
||||
this(context, false);
|
||||
}
|
||||
|
||||
public ExtHttpTtyConnection(ChannelHandlerContext context, boolean quiet) {
|
||||
this.context = context;
|
||||
this.quiet = quiet;
|
||||
}
|
||||
|
||||
@Override
|
||||
@@ -59,10 +66,13 @@ public class ExtHttpTtyConnection extends HttpTtyConnection {
|
||||
}
|
||||
|
||||
public Map<String, Object> extSessions() {
|
||||
Map<String, Object> result = new HashMap<String, Object>();
|
||||
if (quiet) {
|
||||
result.put(Session.QUIET, Boolean.TRUE);
|
||||
}
|
||||
if (context != null) {
|
||||
HttpSession httpSession = HttpSessionManager.getHttpSessionFromContext(context);
|
||||
if (httpSession != null) {
|
||||
Map<String, Object> result = new HashMap<String, Object>();
|
||||
Object subject = httpSession.getAttribute(ArthasConstants.SUBJECT_KEY);
|
||||
if (subject != null) {
|
||||
result.put(ArthasConstants.SUBJECT_KEY, subject);
|
||||
@@ -72,11 +82,11 @@ public class ExtHttpTtyConnection extends HttpTtyConnection {
|
||||
if (userId != null) {
|
||||
result.put(ArthasConstants.USER_ID_KEY, userId);
|
||||
}
|
||||
if (!result.isEmpty()) {
|
||||
return result;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!result.isEmpty()) {
|
||||
return result;
|
||||
}
|
||||
return Collections.emptyMap();
|
||||
}
|
||||
|
||||
|
||||
@@ -56,6 +56,7 @@ public class HttpRequestHandler extends SimpleChannelInboundHandler<FullHttpRequ
|
||||
protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest request) throws Exception {
|
||||
String path = new URI(request.uri()).getPath();
|
||||
if (wsUri.equalsIgnoreCase(path)) {
|
||||
ctx.channel().attr(TtyWebSocketFrameHandler.REQUEST_URI).set(request.uri());
|
||||
ctx.fireChannelRead(request.retain());
|
||||
} else {
|
||||
if (HttpUtil.is100ContinueExpected(request)) {
|
||||
|
||||
+59
-13
@@ -16,26 +16,33 @@
|
||||
|
||||
package com.taobao.arthas.core.shell.term.impl.http;
|
||||
|
||||
import com.alibaba.arthas.deps.org.slf4j.Logger;
|
||||
import com.alibaba.arthas.deps.org.slf4j.LoggerFactory;
|
||||
import io.netty.channel.ChannelHandlerContext;
|
||||
import io.netty.channel.SimpleChannelInboundHandler;
|
||||
import io.netty.channel.group.ChannelGroup;
|
||||
import io.netty.handler.codec.http.QueryStringDecoder;
|
||||
import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
|
||||
import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
|
||||
import io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
|
||||
import io.netty.handler.timeout.IdleStateEvent;
|
||||
import io.netty.util.AttributeKey;
|
||||
import io.termd.core.function.Consumer;
|
||||
import io.termd.core.http.HttpTtyConnection;
|
||||
import io.termd.core.tty.TtyConnection;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* @author <a href="mailto:julien@julienviet.com">Julien Viet</a>
|
||||
*/
|
||||
public class TtyWebSocketFrameHandler extends SimpleChannelInboundHandler<TextWebSocketFrame> {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(TtyWebSocketFrameHandler.class);
|
||||
static final AttributeKey<String> REQUEST_URI = AttributeKey.valueOf("arthas.websocket.requestUri");
|
||||
|
||||
private final ChannelGroup group;
|
||||
private final Consumer<TtyConnection> handler;
|
||||
private ChannelHandlerContext context;
|
||||
private HttpTtyConnection conn;
|
||||
|
||||
public TtyWebSocketFrameHandler(ChannelGroup group, Consumer<TtyConnection> handler) {
|
||||
@@ -43,19 +50,20 @@ public class TtyWebSocketFrameHandler extends SimpleChannelInboundHandler<TextWe
|
||||
this.handler = handler;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void channelActive(ChannelHandlerContext ctx) throws Exception {
|
||||
super.channelActive(ctx);
|
||||
context = ctx;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
|
||||
if (evt == WebSocketServerProtocolHandler.ServerHandshakeStateEvent.HANDSHAKE_COMPLETE) {
|
||||
ctx.pipeline().remove(HttpRequestHandler.class);
|
||||
group.add(ctx.channel());
|
||||
conn = new ExtHttpTtyConnection(context);
|
||||
handler.accept(conn);
|
||||
// Netty 会先发旧事件,再发带 requestUri 的 HandshakeComplete;这里延迟兜底,优先读取 query。
|
||||
ctx.executor().execute(new Runnable() {
|
||||
@Override
|
||||
public void run() {
|
||||
handleHandshakeComplete(ctx, null);
|
||||
}
|
||||
});
|
||||
} else if (evt instanceof WebSocketServerProtocolHandler.HandshakeComplete) {
|
||||
WebSocketServerProtocolHandler.HandshakeComplete handshakeComplete =
|
||||
(WebSocketServerProtocolHandler.HandshakeComplete) evt;
|
||||
handleHandshakeComplete(ctx, handshakeComplete.requestUri());
|
||||
} else if (evt instanceof IdleStateEvent) {
|
||||
ctx.writeAndFlush(new PingWebSocketFrame());
|
||||
} else {
|
||||
@@ -66,7 +74,6 @@ public class TtyWebSocketFrameHandler extends SimpleChannelInboundHandler<TextWe
|
||||
@Override
|
||||
public void channelInactive(ChannelHandlerContext ctx) throws Exception {
|
||||
HttpTtyConnection tmp = conn;
|
||||
context = null;
|
||||
conn = null;
|
||||
if (tmp != null) {
|
||||
Consumer<Void> closeHandler = tmp.getCloseHandler();
|
||||
@@ -78,6 +85,45 @@ public class TtyWebSocketFrameHandler extends SimpleChannelInboundHandler<TextWe
|
||||
|
||||
@Override
|
||||
public void channelRead0(ChannelHandlerContext ctx, TextWebSocketFrame msg) throws Exception {
|
||||
conn.writeToDecoder(msg.text());
|
||||
HttpTtyConnection tmp = conn;
|
||||
if (tmp == null) {
|
||||
logger.warn("websocket frame received before handshake completed, closing channel");
|
||||
ctx.close();
|
||||
return;
|
||||
}
|
||||
tmp.writeToDecoder(msg.text());
|
||||
}
|
||||
|
||||
private void handleHandshakeComplete(ChannelHandlerContext ctx, String requestUri) {
|
||||
if (conn != null) {
|
||||
return;
|
||||
}
|
||||
ctx.pipeline().remove(HttpRequestHandler.class);
|
||||
group.add(ctx.channel());
|
||||
conn = new ExtHttpTtyConnection(ctx, isQuietRequest(ctx, requestUri));
|
||||
handler.accept(conn);
|
||||
}
|
||||
|
||||
static boolean isQuietRequest(ChannelHandlerContext ctx, String requestUri) {
|
||||
if (requestUri == null && ctx != null) {
|
||||
requestUri = ctx.channel().attr(REQUEST_URI).get();
|
||||
}
|
||||
return isQuietRequest(requestUri);
|
||||
}
|
||||
|
||||
static boolean isQuietRequest(String requestUri) {
|
||||
if (requestUri == null) {
|
||||
return false;
|
||||
}
|
||||
List<String> values = new QueryStringDecoder(requestUri).parameters().get("quiet");
|
||||
if (values == null) {
|
||||
return false;
|
||||
}
|
||||
for (String value : values) {
|
||||
if ("true".equalsIgnoreCase(value)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
+125
@@ -0,0 +1,125 @@
|
||||
package com.taobao.arthas.core.command;
|
||||
|
||||
import com.taobao.arthas.core.command.model.InputStatusModel;
|
||||
import com.taobao.arthas.core.command.model.MessageModel;
|
||||
import com.taobao.arthas.core.command.model.ResultModel;
|
||||
import com.taobao.arthas.core.command.model.WelcomeModel;
|
||||
import com.taobao.arthas.core.distribution.ResultConsumer;
|
||||
import com.taobao.arthas.core.distribution.SharingResultDistributor;
|
||||
import com.taobao.arthas.core.shell.command.CommandResolver;
|
||||
import com.taobao.arthas.core.shell.session.Session;
|
||||
import com.taobao.arthas.core.shell.session.SessionManager;
|
||||
import com.taobao.arthas.core.shell.session.impl.SessionImpl;
|
||||
import com.taobao.arthas.core.shell.system.JobController;
|
||||
import com.taobao.arthas.core.shell.system.impl.InternalCommandManager;
|
||||
import com.taobao.arthas.core.shell.system.impl.JobControllerImpl;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.lang.instrument.Instrumentation;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.UUID;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
class CommandExecutorImplQuietSessionTest {
|
||||
private final RecordingSessionManager sessionManager = new RecordingSessionManager();
|
||||
private final CommandExecutorImpl commandExecutor = new CommandExecutorImpl(sessionManager);
|
||||
|
||||
@AfterEach
|
||||
void tearDown() {
|
||||
sessionManager.close();
|
||||
}
|
||||
|
||||
@Test
|
||||
void createSessionShouldAppendWelcomeByDefault() {
|
||||
Map<String, Object> result = commandExecutor.createSession(false);
|
||||
List<ResultModel> results = pollResults((String) result.get("consumerId"));
|
||||
|
||||
assertThat(results).anyMatch(MessageModel.class::isInstance);
|
||||
assertThat(results).anyMatch(WelcomeModel.class::isInstance);
|
||||
assertThat(results).anyMatch(InputStatusModel.class::isInstance);
|
||||
assertThat((Object) sessionManager.lastSession.get(Session.QUIET)).isNull();
|
||||
}
|
||||
|
||||
@Test
|
||||
void createQuietSessionShouldSkipWelcomeModelsAndKeepInputStatus() {
|
||||
Map<String, Object> result = commandExecutor.createSession(true);
|
||||
List<ResultModel> results = pollResults((String) result.get("consumerId"));
|
||||
|
||||
assertThat(results).noneMatch(MessageModel.class::isInstance);
|
||||
assertThat(results).noneMatch(WelcomeModel.class::isInstance);
|
||||
assertThat(results).anyMatch(InputStatusModel.class::isInstance);
|
||||
assertThat((Object) sessionManager.lastSession.get(Session.QUIET)).isEqualTo(Boolean.TRUE);
|
||||
}
|
||||
|
||||
private List<ResultModel> pollResults(String consumerId) {
|
||||
SharingResultDistributor distributor = sessionManager.lastSession.getResultDistributor();
|
||||
ResultConsumer consumer = distributor.getConsumer(consumerId);
|
||||
return consumer.pollResults();
|
||||
}
|
||||
|
||||
private static final class RecordingSessionManager implements SessionManager {
|
||||
private final InternalCommandManager commandManager =
|
||||
new InternalCommandManager(Collections.<CommandResolver>emptyList());
|
||||
private final JobControllerImpl jobController = new JobControllerImpl();
|
||||
private Session lastSession;
|
||||
|
||||
@Override
|
||||
public Session createSession() {
|
||||
Session session = new SessionImpl();
|
||||
session.put(Session.COMMAND_MANAGER, commandManager);
|
||||
session.put(Session.PID, 123L);
|
||||
session.put(Session.ID, UUID.randomUUID().toString());
|
||||
this.lastSession = session;
|
||||
return session;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Session getSession(String sessionId) {
|
||||
if (lastSession != null && lastSession.getSessionId().equals(sessionId)) {
|
||||
return lastSession;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Session removeSession(String sessionId) {
|
||||
Session session = getSession(sessionId);
|
||||
if (session != null && session.getResultDistributor() != null) {
|
||||
session.getResultDistributor().close();
|
||||
}
|
||||
lastSession = null;
|
||||
return session;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void updateAccessTime(Session session) {
|
||||
session.setLastAccessTime(System.currentTimeMillis());
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
if (lastSession != null && lastSession.getResultDistributor() != null) {
|
||||
lastSession.getResultDistributor().close();
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public InternalCommandManager getCommandManager() {
|
||||
return commandManager;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Instrumentation getInstrumentation() {
|
||||
return null;
|
||||
}
|
||||
|
||||
@Override
|
||||
public JobController getJobController() {
|
||||
return jobController;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
package com.taobao.arthas.core.shell.impl;
|
||||
|
||||
import com.taobao.arthas.core.shell.ShellServer;
|
||||
import com.taobao.arthas.core.shell.cli.Completion;
|
||||
import com.taobao.arthas.core.shell.command.CommandResolver;
|
||||
import com.taobao.arthas.core.shell.handlers.Handler;
|
||||
import com.taobao.arthas.core.shell.session.Session;
|
||||
import com.taobao.arthas.core.shell.system.impl.InternalCommandManager;
|
||||
import com.taobao.arthas.core.shell.system.impl.JobControllerImpl;
|
||||
import com.taobao.arthas.core.shell.term.SignalHandler;
|
||||
import com.taobao.arthas.core.shell.term.Term;
|
||||
import io.termd.core.function.Function;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.util.Collections;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
class ShellImplQuietTest {
|
||||
|
||||
@Test
|
||||
void initShouldWriteWelcomeForNormalSession() {
|
||||
RecordingTerm term = new RecordingTerm("xterm");
|
||||
ShellImpl shell = newShell(term);
|
||||
|
||||
shell.setWelcome("welcome");
|
||||
shell.init();
|
||||
|
||||
assertThat(term.output()).isEqualTo("welcome\n");
|
||||
}
|
||||
|
||||
@Test
|
||||
void initShouldSkipWelcomeForArthasAgentTerminalType() {
|
||||
RecordingTerm term = new RecordingTerm("arthas-agent");
|
||||
ShellImpl shell = newShell(term);
|
||||
|
||||
shell.setWelcome("welcome");
|
||||
shell.init();
|
||||
shell.readline();
|
||||
|
||||
assertThat(term.output()).isEmpty();
|
||||
assertThat(term.readlinePrompt()).isEqualTo("[arthas@123]$ ");
|
||||
assertThat((Object) shell.session().get(Session.QUIET)).isEqualTo(Boolean.TRUE);
|
||||
}
|
||||
|
||||
private ShellImpl newShell(Term term) {
|
||||
InternalCommandManager commandManager =
|
||||
new InternalCommandManager(Collections.<CommandResolver>emptyList());
|
||||
return new ShellImpl((ShellServer) null, term, commandManager, null, 123L, new JobControllerImpl());
|
||||
}
|
||||
|
||||
private static final class RecordingTerm implements Term {
|
||||
private final String type;
|
||||
private final StringBuilder output = new StringBuilder();
|
||||
private Session session;
|
||||
private String readlinePrompt;
|
||||
|
||||
private RecordingTerm(String type) {
|
||||
this.type = type;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Term resizehandler(Handler<Void> handler) {
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Term stdinHandler(Handler<String> handler) {
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Term stdoutHandler(Function<String, String> handler) {
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Term write(String data) {
|
||||
output.append(data);
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public long lastAccessedTime() {
|
||||
return System.currentTimeMillis();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Term echo(String text) {
|
||||
output.append(text);
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Term setSession(Session session) {
|
||||
this.session = session;
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Term interruptHandler(SignalHandler handler) {
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Term suspendHandler(SignalHandler handler) {
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void readline(String prompt, Handler<String> lineHandler) {
|
||||
this.readlinePrompt = prompt;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void readline(String prompt, Handler<String> lineHandler, Handler<Completion> completionHandler) {
|
||||
this.readlinePrompt = prompt;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Term closeHandler(Handler<Void> handler) {
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
}
|
||||
|
||||
@Override
|
||||
public String type() {
|
||||
return type;
|
||||
}
|
||||
|
||||
@Override
|
||||
public int width() {
|
||||
return 80;
|
||||
}
|
||||
|
||||
@Override
|
||||
public int height() {
|
||||
return 24;
|
||||
}
|
||||
|
||||
private String output() {
|
||||
return output.toString();
|
||||
}
|
||||
|
||||
private String readlinePrompt() {
|
||||
return readlinePrompt;
|
||||
}
|
||||
}
|
||||
}
|
||||
+55
@@ -0,0 +1,55 @@
|
||||
package com.taobao.arthas.core.shell.term.impl.http;
|
||||
|
||||
import com.taobao.arthas.core.shell.session.Session;
|
||||
import io.netty.channel.ChannelHandlerContext;
|
||||
import io.netty.channel.ChannelInboundHandlerAdapter;
|
||||
import io.netty.channel.embedded.EmbeddedChannel;
|
||||
import io.netty.channel.group.DefaultChannelGroup;
|
||||
import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
|
||||
import io.netty.util.concurrent.GlobalEventExecutor;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.util.Map;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
class TtyWebSocketFrameHandlerQuietTest {
|
||||
|
||||
@Test
|
||||
void isQuietRequestShouldOnlyAcceptQuietTrue() {
|
||||
assertThat(TtyWebSocketFrameHandler.isQuietRequest("/ws?quiet=true")).isTrue();
|
||||
assertThat(TtyWebSocketFrameHandler.isQuietRequest("/ws?quiet=TRUE")).isTrue();
|
||||
assertThat(TtyWebSocketFrameHandler.isQuietRequest("/ws?quiet=false")).isFalse();
|
||||
assertThat(TtyWebSocketFrameHandler.isQuietRequest("/ws")).isFalse();
|
||||
assertThat(TtyWebSocketFrameHandler.isQuietRequest(null)).isFalse();
|
||||
}
|
||||
|
||||
@Test
|
||||
void isQuietRequestShouldReadFallbackUriFromChannelAttribute() {
|
||||
EmbeddedChannel channel = new EmbeddedChannel(new ChannelInboundHandlerAdapter());
|
||||
ChannelHandlerContext context = channel.pipeline().firstContext();
|
||||
channel.attr(TtyWebSocketFrameHandler.REQUEST_URI).set("/ws?quiet=true");
|
||||
|
||||
assertThat(TtyWebSocketFrameHandler.isQuietRequest(context, null)).isTrue();
|
||||
assertThat(TtyWebSocketFrameHandler.isQuietRequest(context, "/ws?quiet=false")).isFalse();
|
||||
}
|
||||
|
||||
@Test
|
||||
void channelReadShouldCloseChannelWhenHandshakeIsNotCompleted() {
|
||||
EmbeddedChannel channel = new EmbeddedChannel(new TtyWebSocketFrameHandler(
|
||||
new DefaultChannelGroup(GlobalEventExecutor.INSTANCE),
|
||||
connection -> {
|
||||
}));
|
||||
|
||||
channel.writeInbound(new TextWebSocketFrame("help"));
|
||||
|
||||
assertThat(channel.isOpen()).isFalse();
|
||||
}
|
||||
|
||||
@Test
|
||||
void extSessionsShouldCarryQuietFlag() {
|
||||
Map<String, Object> extSessions = new ExtHttpTtyConnection(null, true).extSessions();
|
||||
|
||||
assertThat(extSessions).containsEntry(Session.QUIET, Boolean.TRUE);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user