JShellClient.java

   1 package org.argeo.cms.jshell;
   2
   3 import static java.lang.System.Logger.Level.ERROR;
   4 import static java.lang.System.Logger.Level.TRACE;
   5 import static java.net.StandardProtocolFamily.UNIX;
   6 import static java.nio.charset.StandardCharsets.UTF_8;
   7
   8 import java.io.BufferedReader;
   9 import java.io.File;
  10 import java.io.IOException;
  11 import java.io.InputStream;
  12 import java.io.InputStreamReader;
  13 import java.io.OutputStream;
  14 import java.lang.System.Logger;
  15 import java.lang.management.ManagementFactory;
  16 import java.net.StandardSocketOptions;
  17 import java.net.UnixDomainSocketAddress;
  18 import java.nio.ByteBuffer;
  19 import java.nio.channels.Channels;
  20 import java.nio.channels.ClosedByInterruptException;
  21 import java.nio.channels.ReadableByteChannel;
  22 import java.nio.channels.SocketChannel;
  23 import java.nio.channels.WritableByteChannel;
  24 import java.nio.file.Files;
  25 import java.nio.file.Path;
  26 import java.nio.file.Paths;
  27 import java.util.ArrayList;
  28 import java.util.HashMap;
  29 import java.util.List;
  30 import java.util.Map;
  31 import java.util.UUID;
  32
  33 public class JShellClient {
  34         private final static Logger logger = System.getLogger(JShellClient.class.getName());
  35
  36         public final static String STD = "std";
  37         public final static String CTL = "ctl";
  38
  39         public final static String JSH = "jsh";
  40         public final static String JTERM = "jterm";
  41
  42         private static String sttyExec = "/usr/bin/stty";
  43
  44         /** Benchmark based on uptime. */
  45         private static boolean benchmark = false;
  46
  47         /**
  48          * The real path (following symbolic links) to the directory were to create
  49          * sessions.
  50          */
  51         private Path localBase;
  52
  53         /** The symbolic name of the bundle from which to run. */
  54         private String symbolicName;
  55
  56         /** The script to run. */
  57         private Path script;
  58         /** Additional arguments of the script */
  59         private List<String> scriptArgs;
  60
  61         private String ttyConfig;
  62         private boolean terminal;
  63
  64         /** Workaround to be able to test in Eclipse console */
  65         private boolean inEclipse = false;
  66
  67         public JShellClient(Path targetStateDirectory, String symbolicName, Path script, List<String> scriptArgs) {
  68                 try {
  69                         this.terminal = System.console() != null && script == null;
  70                         if (inEclipse && script == null)
  71                                 terminal = true;
  72                         if (terminal) {
  73                                 localBase = targetStateDirectory.resolve(JTERM);
  74                         } else {
  75                                 localBase = targetStateDirectory.resolve(JSH);
  76                         }
  77                         if (Files.isSymbolicLink(localBase)) {
  78                                 localBase = localBase.toRealPath();
  79                         }
  80                         this.symbolicName = symbolicName;
  81                         this.script = script;
  82                         this.scriptArgs = scriptArgs == null ? new ArrayList<>() : scriptArgs;
  83                 } catch (IOException e) {
  84                         throw new IllegalStateException("Cannot initialise client", e);
  85                 }
  86         }
  87
  88         public void run() {
  89                 try {
  90                         if (terminal)
  91                                 toRawTerminal();
  92                         SocketPipeSource std = new SocketPipeSource(STD, script != null);
  93                         std.setInputStream(System.in);
  94                         std.setOutputStream(System.out);
  95
  96                         SocketPipeSource ctl = new SocketPipeSource(CTL, false);
  97                         ctl.setOutputStream(System.err);
  98
  99                         Runtime.getRuntime().addShutdownHook(new Thread(() -> {
 100                                 System.out.println("\nShutting down...");
 101                                 toOriginalTerminal();
 102                                 std.shutdown();
 103                                 ctl.shutdown();
 104                         }, "Shut down JShell client"));
 105
 106                         Path bundleSnDir = localBase.resolve(symbolicName);
 107                         if (!Files.exists(bundleSnDir))
 108                                 Files.createDirectory(bundleSnDir);
 109                         UUID uuid = UUID.randomUUID();
 110                         Path sessionDir = bundleSnDir.resolve(uuid.toString());
 111
 112                         // creating the directory will trigger opening of the session on server side
 113                         Files.createDirectory(sessionDir);
 114
 115                         Path stdPath = sessionDir.resolve(JShellClient.STD);
 116                         Path ctlPath = sessionDir.resolve(JShellClient.CTL);
 117
 118                         while (!(Files.exists(stdPath) && Files.exists(ctlPath))) {
 119                                 // TODO timeout
 120                                 try {
 121                                         Thread.sleep(1);
 122                                 } catch (InterruptedException e) {
 123                                         // silent
 124                                 }
 125                         }
 126
 127                         UnixDomainSocketAddress stdSocketAddress = UnixDomainSocketAddress.of(stdPath.toRealPath());
 128                         UnixDomainSocketAddress ctlSocketAddress = UnixDomainSocketAddress.of(ctlPath.toRealPath());
 129
 130                         try (SocketChannel stdChannel = SocketChannel.open(UNIX);
 131                                         SocketChannel ctlChannel = SocketChannel.open(UNIX);) {
 132                                 ctlChannel.connect(ctlSocketAddress);
 133                                 ctl.process(ctlChannel);
 134                                 if (script != null) {
 135                                         new ScriptThread(ctlChannel).start();
 136                                 }
 137                                 stdChannel.connect(stdSocketAddress);
 138                                 std.process(stdChannel);
 139
 140                                 while (!std.isCompleted() && !ctl.isCompleted()) {
 141                                         // isCompleted() will block
 142                                 }
 143                         }
 144                         if (benchmark)
 145                                 System.err.println(ManagementFactory.getRuntimeMXBean().getUptime());
 146                         std.shutdown();
 147                         ctl.shutdown();
 148                 } catch (IOException e) {
 149                         e.printStackTrace();
 150                 } finally {
 151                         toOriginalTerminal();
 152                 }
 153
 154         }
 155
 156         public static void main(String[] args) throws IOException, InterruptedException {
 157                 if (benchmark)
 158                         System.err.println(ManagementFactory.getRuntimeMXBean().getUptime());
 159                 List<String> plainArgs = new ArrayList<>();
 160                 Map<String, List<String>> options = new HashMap<>();
 161                 String currentOption = null;
 162                 for (int i = 0; i < args.length; i++) {
 163                         if (args[i].startsWith("-")) {
 164                                 currentOption = args[i];
 165                                 if (!options.containsKey(currentOption))
 166                                         options.put(currentOption, new ArrayList<>());
 167                                 i++;
 168                                 options.get(currentOption).add(args[i]);
 169                         } else {
 170                                 plainArgs.add(args[i]);
 171                         }
 172                 }
 173
 174                 Path targetStateDirectory = Paths.get(options.get("-d").get(0));
 175                 String symbolicName = options.get("-b").get(0);
 176
 177                 Path script = plainArgs.isEmpty() ? null : Paths.get(plainArgs.get(0));
 178                 List<String> scriptArgs = new ArrayList<>();
 179                 for (int i = 1; i < plainArgs.size(); i++)
 180                         scriptArgs.add(plainArgs.get(i));
 181
 182                 JShellClient client = new JShellClient(targetStateDirectory, symbolicName, script, scriptArgs);
 183                 client.run();
 184         }
 185
 186         /*
 187          * TERMINAL
 188          */
 189         /** Set the terminal to raw mode. */
 190         protected synchronized void toRawTerminal() {
 191                 boolean isWindows = File.separatorChar == '\\';
 192                 if (isWindows)
 193                         return;
 194                 if (inEclipse)
 195                         return;
 196                 // save current configuration
 197                 ttyConfig = stty("-g");
 198                 if (ttyConfig == null)
 199                         return;
 200                 ttyConfig.trim();
 201                 // set the console to be character-buffered instead of line-buffered
 202                 stty("-icanon min 1");
 203                 // disable character echoing
 204                 stty("-echo");
 205         }
 206
 207         /** Restore original terminal configuration. */
 208         protected synchronized void toOriginalTerminal() {
 209                 if (ttyConfig == null)
 210                         return;
 211                 try {
 212                         stty(ttyConfig);
 213                 } catch (Exception e) {
 214                         e.printStackTrace();
 215                 }
 216                 ttyConfig = null;
 217         }
 218
 219         /**
 220          * Execute the stty command with the specified arguments against the current
 221          * active terminal.
 222          */
 223         protected String stty(String args) {
 224                 List<String> cmd = new ArrayList<>();
 225                 cmd.add("/bin/sh");
 226                 cmd.add("-c");
 227                 cmd.add(sttyExec + " " + args + " < /dev/tty");
 228
 229                 logger.log(TRACE, () -> cmd.toString());
 230
 231                 try {
 232                         ProcessBuilder pb = new ProcessBuilder(cmd);
 233                         Process p = pb.start();
 234                         String firstLine = new BufferedReader(new InputStreamReader(p.getInputStream())).readLine();
 235                         p.waitFor();
 236                         logger.log(TRACE, () -> firstLine);
 237                         return firstLine;
 238                 } catch (IOException | InterruptedException e) {
 239                         e.printStackTrace();
 240                         return null;
 241                 }
 242         }
 243
 244         /*
 245          * SCRIPT
 246          */
 247         private class ScriptThread extends Thread {
 248                 private SocketChannel channel;
 249
 250                 public ScriptThread(SocketChannel channel) {
 251                         super("JShell script writer");
 252                         this.channel = channel;
 253                 }
 254
 255                 @Override
 256                 public void run() {
 257                         try {
 258                                 if (benchmark)
 259                                         System.err.println(ManagementFactory.getRuntimeMXBean().getUptime());
 260                                 StringBuilder sb = new StringBuilder();
 261 //                              sb.append("/set feedback silent\n");
 262                                 if (!scriptArgs.isEmpty()) {
 263                                         // additional arguments as $1, $2, etc.
 264                                         for (String arg : scriptArgs)
 265                                                 sb.append('\"').append(arg).append('\"').append(";\n");
 266                                 }
 267                                 if (sb.length() > 0)
 268                                         writeLine(sb);
 269
 270                                 ByteBuffer buffer = ByteBuffer.allocate(1024);
 271                                 try (BufferedReader reader = Files.newBufferedReader(script)) {
 272                                         String line;
 273                                         lines: while ((line = reader.readLine()) != null) {
 274                                                 if (line.startsWith("#"))
 275                                                         continue lines;
 276                                                 buffer.put((line + "\n").getBytes(UTF_8));
 277                                                 buffer.flip();
 278                                                 channel.write(buffer);
 279                                                 buffer.rewind();
 280                                         }
 281                                 }
 282
 283 //                              ByteBuffer buffer = ByteBuffer.allocate(1024);
 284 //                              try (SeekableByteChannel scriptChannel = Files.newByteChannel(script, StandardOpenOption.READ)) {
 285 //                                      while (channel.isConnected()) {
 286 //                                              if (scriptChannel.read(buffer) < 0)
 287 //                                                      break;
 288 //                                              buffer.flip();
 289 //                                              channel.write(buffer);
 290 //                                              buffer.rewind();
 291 //                                      }
 292 //                              }
 293
 294                                 // exit
 295                                 if (channel.isConnected())
 296                                         writeLine("/exit");
 297                         } catch (IOException e) {
 298                                 logger.log(ERROR, "Cannot execute " + script, e);
 299                         }
 300                 }
 301
 302                 private void writeLine(Object obj) throws IOException {
 303                         channel.write(ByteBuffer.wrap((obj + "\n").getBytes(UTF_8)));
 304                 }
 305         }
 306 }
 307
 308 /** Pipe streams to a channel. */
 309 class SocketPipeSource {
 310         private ReadableByteChannel inChannel;
 311         private WritableByteChannel outChannel;
 312
 313         private Thread readThread;
 314         private Thread forwardThread;
 315
 316         private int inBufferSize = 1;
 317         private int outBufferSize = 1024;
 318
 319         private final String id;
 320         private final boolean batch;
 321
 322         private boolean completed = false;
 323
 324         public SocketPipeSource(String id, boolean batch) {
 325                 this.id = id;
 326                 this.batch = batch;
 327         }
 328
 329         public void process(SocketChannel channel) throws IOException {
 330                 if (batch) {
 331                         Integer socketRcvBuf = channel.getOption(StandardSocketOptions.SO_RCVBUF);
 332                         inBufferSize = socketRcvBuf;
 333                         outBufferSize = socketRcvBuf;
 334                 }
 335
 336                 readThread = new Thread(() -> {
 337
 338                         try {
 339                                 ByteBuffer buffer = ByteBuffer.allocate(outBufferSize);
 340                                 while (true) {
 341                                         if (channel.read(buffer) < 0)
 342                                                 break;
 343                                         buffer.flip();
 344                                         outChannel.write(buffer);
 345                                         buffer.rewind();
 346                                 }
 347                         } catch (ClosedByInterruptException e) {
 348                                 // silent
 349                         } catch (IOException e) {
 350                                 e.printStackTrace();
 351                         }
 352                         markCompleted();
 353                 }, "JShell read " + id);
 354                 readThread.start();
 355
 356                 // TODO make it smarter than a 1 byte buffer
 357                 // we should recognize control characters
 358                 // e.g ^C
 359 //              int c = System.in.read();
 360 //              if (c == 0x1B) {
 361 //                      break;
 362 //              }
 363
 364                 if (inChannel != null) {
 365                         forwardThread = new Thread(() -> {
 366                                 try {
 367                                         ByteBuffer buffer = ByteBuffer.allocate(inBufferSize);
 368                                         while (channel.isConnected()) {
 369                                                 if (inChannel.read(buffer) < 0)
 370                                                         break;
 371 //                      int b = (int) buffer.get(0);
 372 //                      if (b == 0x1B) {
 373 //                              System.out.println("Ctrl+C");
 374 //                      }
 375
 376                                                 buffer.flip();
 377                                                 channel.write(buffer);
 378                                                 buffer.rewind();
 379                                         }
 380                                 } catch (IOException e) {
 381                                         e.printStackTrace();
 382                                 }
 383                         }, "JShell write " + id);
 384                         forwardThread.setDaemon(true);
 385                         forwardThread.start();
 386                         // end
 387                         // TODO make it more robust
 388                         // we want to be asynchronous when read only
 389 //                      try {
 390 //                              // TODO add timeout
 391 //                              readThread.join();
 392 //                      } catch (InterruptedException e) {
 393 //                              e.printStackTrace();
 394 //                      }
 395
 396                 }
 397         }
 398
 399         public synchronized boolean isCompleted() {
 400                 if (!completed)
 401                         try {
 402                                 wait();
 403                         } catch (InterruptedException e) {
 404                                 // silent
 405                         }
 406                 return completed;
 407         }
 408
 409         protected synchronized void markCompleted() {
 410                 completed = true;
 411                 notifyAll();
 412         }
 413
 414         public void shutdown() {
 415                 if (inChannel != null)
 416                         try {
 417                                 inChannel.close();
 418                         } catch (IOException e) {
 419                                 e.printStackTrace();
 420                         }
 421                 try {
 422                         outChannel.close();
 423                 } catch (IOException e) {
 424                         e.printStackTrace();
 425                 }
 426                 if (inChannel != null)
 427                         forwardThread.interrupt();
 428                 readThread.interrupt();
 429         }
 430
 431         public void setInputStream(InputStream in) {
 432                 inChannel = Channels.newChannel(in);
 433         }
 434
 435         public void setOutputStream(OutputStream out) {
 436                 outChannel = Channels.newChannel(out);
 437         }
 438 }