diff --git a/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/CommandExecutor.java b/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/CommandExecutor.java index a94a7f77b99..a0cd351bcda 100644 --- a/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/CommandExecutor.java +++ b/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/CommandExecutor.java @@ -47,7 +47,7 @@ default Map executeSync(String commandLine, Object authSubject, Map interruptJob(String sessionId); - Map createSession(); + Map createSession(boolean quiet); Map closeSession(String sessionId); @@ -61,4 +61,3 @@ default Map executeSync(String commandLine, Object authSubject, */ void setSessionUserId(String sessionId, String userId); } - diff --git a/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManager.java b/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManager.java index 949a438a045..91a7c5617f8 100644 --- a/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManager.java +++ b/arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManager.java @@ -79,7 +79,7 @@ public void updateAccessTime() { } public CommandSessionBinding createCommandSession(String mcpSessionId) { - Map result = commandExecutor.createSession(); + Map result = commandExecutor.createSession(true); CommandSessionBinding binding = new CommandSessionBinding( mcpSessionId, @@ -177,7 +177,7 @@ public CommandSessionBinding createIsolatedTaskSession(String taskId) { .build(); } - Map result = commandExecutor.createSession(); + Map result = commandExecutor.createSession(true); CommandSessionBinding binding = new CommandSessionBinding( "task-" + taskId, // 使用 task ID 作为 MCP session ID diff --git a/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/session/ArthasCommandContextAuthTest.java b/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/session/ArthasCommandContextAuthTest.java index 2f2607d45c2..85292c05c63 100644 --- a/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/session/ArthasCommandContextAuthTest.java +++ b/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/session/ArthasCommandContextAuthTest.java @@ -66,7 +66,7 @@ public Map interruptJob(String sessionId) { } @Override - public Map createSession() { + public Map createSession(boolean quiet) { Map result = new HashMap(); result.put("sessionId", "created-session"); result.put("consumerId", "created-consumer"); diff --git a/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManagerQuietTest.java b/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManagerQuietTest.java new file mode 100644 index 00000000000..143aded4d5f --- /dev/null +++ b/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManagerQuietTest.java @@ -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 quietFlags = new ArrayList(); + + @Override + public Map executeSync(String commandLine, long timeout, String sessionId, + Object authSubject, String userId) { + return new HashMap(); + } + + @Override + public Map executeAsync(String commandLine, String sessionId) { + return new HashMap(); + } + + @Override + public Map pullResults(String sessionId, String consumerId) { + return new HashMap(); + } + + @Override + public Map interruptJob(String sessionId) { + return new HashMap(); + } + + @Override + public Map createSession(boolean quiet) { + quietFlags.add(quiet); + Map result = new HashMap(); + result.put("sessionId", "session-" + quietFlags.size()); + result.put("consumerId", "consumer-" + quietFlags.size()); + return result; + } + + @Override + public Map closeSession(String sessionId) { + return new HashMap(); + } + + @Override + public void setSessionAuth(String sessionId, Object authSubject) { + } + + @Override + public void setSessionUserId(String sessionId, String userId) { + } + } +} diff --git a/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/task/DefaultCreateTaskContextAuthTest.java b/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/task/DefaultCreateTaskContextAuthTest.java index 7c7ed04cc55..b0060a62f6d 100644 --- a/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/task/DefaultCreateTaskContextAuthTest.java +++ b/arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/task/DefaultCreateTaskContextAuthTest.java @@ -80,7 +80,7 @@ public Map interruptJob(String sessionId) { } @Override - public Map createSession() { + public Map createSession(boolean quiet) { Map result = new HashMap(); result.put("sessionId", "isolated-session"); result.put("consumerId", "isolated-consumer"); diff --git a/client/pom.xml b/client/pom.xml index 77fb6b703ab..3a22491dd83 100644 --- a/client/pom.xml +++ b/client/pom.xml @@ -61,6 +61,16 @@ jline jline + + org.junit.jupiter + junit-jupiter + test + + + org.assertj + assertj-core + test + diff --git a/client/src/main/java/com/taobao/arthas/client/TelnetConsole.java b/client/src/main/java/com/taobao/arthas/client/TelnetConsole.java index 2c3d173d879..e6141d9573e 100644 --- a/client/src/main/java/com/taobao/arthas/client/TelnetConsole.java +++ b/client/src/main/java/com/taobao/arthas/client/TelnetConsole.java @@ -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 void setheight(int height) { 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 static int process(String[] args, ActionListener eotEventCallback) throws } } - 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 void run() { // 检查到有 [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 Integer getheight() { return height; } + public boolean isQuiet() { + return quiet; + } + public boolean isHelp() { return help; } diff --git a/client/src/test/java/com/taobao/arthas/client/TelnetConsoleBatchModeTest.java b/client/src/test/java/com/taobao/arthas/client/TelnetConsoleBatchModeTest.java new file mode 100644 index 00000000000..a145021f061 --- /dev/null +++ b/client/src/test/java/com/taobao/arthas/client/TelnetConsoleBatchModeTest.java @@ -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 = 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; + } + } +} diff --git a/core/src/main/java/com/taobao/arthas/core/command/CommandExecutorImpl.java b/core/src/main/java/com/taobao/arthas/core/command/CommandExecutorImpl.java index 7f8fdb65113..d868c85d05e 100644 --- a/core/src/main/java/com/taobao/arthas/core/command/CommandExecutorImpl.java +++ b/core/src/main/java/com/taobao/arthas/core/command/CommandExecutorImpl.java @@ -279,17 +279,34 @@ public Map interruptJob(String sessionId) { } @Override - public Map createSession() { + public Map 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 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 Map createSession() { welcomeModel.setPid(PidUtils.currentPid()); welcomeModel.setTime(DateUtils.getCurrentDateTime()); resultDistributor.appendResult(welcomeModel); - - updateSessionInputStatus(session, InputStatus.ALLOW_INPUT); - - Map result = new TreeMap<>(); - result.put("success", true); - result.put("sessionId", session.getSessionId()); - result.put("consumerId", resultConsumer.getConsumerId()); - return result; } @Override diff --git a/core/src/main/java/com/taobao/arthas/core/shell/impl/ShellImpl.java b/core/src/main/java/com/taobao/arthas/core/shell/impl/ShellImpl.java index c37b430337f..41cdfb40d95 100644 --- a/core/src/main/java/com/taobao/arthas/core/shell/impl/ShellImpl.java +++ b/core/src/main/java/com/taobao/arthas/core/shell/impl/ShellImpl.java @@ -52,7 +52,7 @@ */ 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 ShellImpl(ShellServer server, Term term, InternalCommandManager commandMa 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 ShellImpl(ShellServer server, Term term, InternalCommandManager commandMa } } } + 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 ShellImpl init() { 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())) { diff --git a/core/src/main/java/com/taobao/arthas/core/shell/session/Session.java b/core/src/main/java/com/taobao/arthas/core/shell/session/Session.java index 4ad7154b539..bfa0369c22f 100644 --- a/core/src/main/java/com/taobao/arthas/core/shell/session/Session.java +++ b/core/src/main/java/com/taobao/arthas/core/shell/session/Session.java @@ -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. */ diff --git a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/ExtHttpTtyConnection.java b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/ExtHttpTtyConnection.java index c73c0f6fde2..ca5d431672f 100644 --- a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/ExtHttpTtyConnection.java +++ b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/ExtHttpTtyConnection.java @@ -6,6 +6,7 @@ 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 @@ */ 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 void close() { } public Map extSessions() { + Map result = new HashMap(); + if (quiet) { + result.put(Session.QUIET, Boolean.TRUE); + } if (context != null) { HttpSession httpSession = HttpSessionManager.getHttpSessionFromContext(context); if (httpSession != null) { - Map result = new HashMap(); Object subject = httpSession.getAttribute(ArthasConstants.SUBJECT_KEY); if (subject != null) { result.put(ArthasConstants.SUBJECT_KEY, subject); @@ -72,11 +82,11 @@ public Map extSessions() { if (userId != null) { result.put(ArthasConstants.USER_ID_KEY, userId); } - if (!result.isEmpty()) { - return result; - } } } + if (!result.isEmpty()) { + return result; + } return Collections.emptyMap(); } diff --git a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/HttpRequestHandler.java b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/HttpRequestHandler.java index dd6e2c08de4..aa78f4552fe 100644 --- a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/HttpRequestHandler.java +++ b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/HttpRequestHandler.java @@ -56,6 +56,7 @@ public HttpRequestHandler(String wsUri, File dir) { 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)) { diff --git a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandler.java b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandler.java index 527166592db..a416952f5aa 100644 --- a/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandler.java +++ b/core/src/main/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandler.java @@ -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 Julien Viet */ public class TtyWebSocketFrameHandler extends SimpleChannelInboundHandler { + private static final Logger logger = LoggerFactory.getLogger(TtyWebSocketFrameHandler.class); + static final AttributeKey REQUEST_URI = AttributeKey.valueOf("arthas.websocket.requestUri"); + private final ChannelGroup group; private final Consumer handler; - private ChannelHandlerContext context; private HttpTtyConnection conn; public TtyWebSocketFrameHandler(ChannelGroup group, Consumer handler) { @@ -43,19 +50,20 @@ public TtyWebSocketFrameHandler(ChannelGroup group, Consumer hand 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 void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exc @Override public void channelInactive(ChannelHandlerContext ctx) throws Exception { HttpTtyConnection tmp = conn; - context = null; conn = null; if (tmp != null) { Consumer closeHandler = tmp.getCloseHandler(); @@ -78,6 +85,45 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { @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 values = new QueryStringDecoder(requestUri).parameters().get("quiet"); + if (values == null) { + return false; + } + for (String value : values) { + if ("true".equalsIgnoreCase(value)) { + return true; + } + } + return false; } } diff --git a/core/src/test/java/com/taobao/arthas/core/command/CommandExecutorImplQuietSessionTest.java b/core/src/test/java/com/taobao/arthas/core/command/CommandExecutorImplQuietSessionTest.java new file mode 100644 index 00000000000..895b768d850 --- /dev/null +++ b/core/src/test/java/com/taobao/arthas/core/command/CommandExecutorImplQuietSessionTest.java @@ -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 result = commandExecutor.createSession(false); + List 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 result = commandExecutor.createSession(true); + List 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 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.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; + } + } +} diff --git a/core/src/test/java/com/taobao/arthas/core/shell/impl/ShellImplQuietTest.java b/core/src/test/java/com/taobao/arthas/core/shell/impl/ShellImplQuietTest.java new file mode 100644 index 00000000000..9ecbcab1c0f --- /dev/null +++ b/core/src/test/java/com/taobao/arthas/core/shell/impl/ShellImplQuietTest.java @@ -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.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 handler) { + return this; + } + + @Override + public Term stdinHandler(Handler handler) { + return this; + } + + @Override + public Term stdoutHandler(Function 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 lineHandler) { + this.readlinePrompt = prompt; + } + + @Override + public void readline(String prompt, Handler lineHandler, Handler completionHandler) { + this.readlinePrompt = prompt; + } + + @Override + public Term closeHandler(Handler 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; + } + } +} diff --git a/core/src/test/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandlerQuietTest.java b/core/src/test/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandlerQuietTest.java new file mode 100644 index 00000000000..fdb0176cf80 --- /dev/null +++ b/core/src/test/java/com/taobao/arthas/core/shell/term/impl/http/TtyWebSocketFrameHandlerQuietTest.java @@ -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 extSessions = new ExtHttpTtyConnection(null, true).extSessions(); + + assertThat(extSessions).containsEntry(Session.QUIET, Boolean.TRUE); + } +}