view mercurial-common/src/jetbrains/buildServer/buildTriggers/vcs/mercurial/command/CommandUtil.java @ 1011:17bd4181c476

TW-55105 Add "teamcity.hg.command.readBytesLimit" to avoid OutOfMemoryError on input reading
author Nikolai.Kulakov@UNIT-1322.Labs.IntelliJ.Net
date Thu, 31 Jan 2019 19:14:24 +0300
parents 7bf4d943d5bb
children 10dc26b32c35
line wrap: on
line source
/*
 * Copyright 2000-2018 JetBrains s.r.o.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package jetbrains.buildServer.buildTriggers.vcs.mercurial.command;

import jetbrains.buildServer.ExecResult;
import jetbrains.buildServer.LineAwareByteArrayOutputStream;
import jetbrains.buildServer.SimpleCommandLineProcessRunner;
import jetbrains.buildServer.buildTriggers.vcs.mercurial.HgFileUtil;
import jetbrains.buildServer.log.Loggers;
import jetbrains.buildServer.util.FileUtil;
import jetbrains.buildServer.util.StringUtil;
import jetbrains.buildServer.vcs.VcsException;
import org.jetbrains.annotations.NotNull;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.nio.charset.Charset;
import java.util.Set;

public class CommandUtil {

  @NotNull
  private static File extractCommandPy(@NotNull final File root, @NotNull final String commandPy) throws VcsException {
    try {
      final File py = new File(root, commandPy);
      FileUtil.copyResource(CommandUtil.class, "/python/" + commandPy, py);

      if (py.length() < 100) throw new IOException("Failed to unpack command resource");
      return py;
    } catch (IOException e) {
      throw new VcsException("Failed to extract .py file: " + e.getMessage(), e);
    }
  }

  public static void setupExtensionsFromResource(@NotNull final MercurialCommandLine cli,
                                                 @NotNull final File tempDir,
                                                 @NotNull final String commandPy) throws VcsException {
    final File file = extractCommandPy(tempDir, commandPy);
    final String extName = commandPy.replaceAll("[^a-zA-Z]+", "");
    cli.addParameters("--config", "extensions." + extName + "=" + file);
  }

  @NotNull
  public static CommandResult runWrappedCommand(@NotNull final MercurialCommandLine originalCommandLine,
                                                @NotNull final CommandSettings settings) throws VcsException {
    final String realCommand = logRunCommand(originalCommandLine, settings);

    final File tempDir = createTempDir();
    try {
      final File commands = writeCommandArguments(originalCommandLine, tempDir);

      final MercurialCommandLine fork = originalCommandLine.forkWithoutCommandlineArguments();
      setupExtensionsFromResource(fork, tempDir, "load-commands-command.py");
      fork.addParameters("CMD", commands.getAbsolutePath());

      String forkCommand = logRunCommand(fork, settings);

      return runCommandWithName(fork, settings, forkCommand + " \n|| for command: " + realCommand);
    } finally {
      FileUtil.delete(tempDir);
    }
  }

  @NotNull
  private static File writeCommandArguments(@NotNull MercurialCommandLine originalCommandLine,
                                            @NotNull File tempDir) throws VcsException {
    try {
      final File commands = new File(tempDir, "command.args");
      FileUtil.writeFile(commands, StringUtil.join("\n", originalCommandLine.getArguments()), "utf-8");
      return commands;
    } catch (IOException e) {
      throw new VcsException("Failed to generate commands file. " + e.getMessage(), e);
    }

  }

  @NotNull
  private static File createTempDir() throws VcsException {
    try {
      return HgFileUtil.createTempDir();
    } catch (IOException e) {
      throw new VcsException("Failed to create temp file. " + e.getMessage(), e);
    }
  }

  @NotNull
  public static CommandResult runCommand(@NotNull final MercurialCommandLine cli,
                                         @NotNull final CommandSettings settings) throws VcsException {
    return runCommandWithName(cli, settings, logRunCommand(cli, settings));
  }

  private static CommandResult runCommandWithName(MercurialCommandLine cli, CommandSettings settings, String command) throws VcsException {
    CommandResult res = run(cli, settings.getTimeout(), settings.getReadBytesLimit(), command, settings.getPrivateData(), settings);
    if (settings.isCheckForFailure() || settings.isFailWithNonEmptyStderr())
      res.checkFailure(settings.isFailWithNonEmptyStderr());
    logCommandOutput(command, res, settings);
    return res;
  }

  @NotNull
  private static CommandResult run(@NotNull final MercurialCommandLine cli,
                                   final int executionTimeout,
                                   final int readBytesLimit,
                                   @NotNull final String command,
                                   @NotNull final Set<String> privateData,
                                   @NotNull CommandSettings settings) {
    final long start = System.currentTimeMillis();
    ByteArrayOutputStream stdoutBuffer = new ByteArrayOutputStream();
    ProgressParser.ProgressConsumer progressConsumer = settings.getProgressConsumer();
    ByteArrayOutputStream stderrBuffer;
    if (progressConsumer != null && cli.hasProgress()) {
      stderrBuffer =  new LineAwareByteArrayOutputStream(Charset.forName("UTF-8"), new ProgressParser(progressConsumer));
      ((LineAwareByteArrayOutputStream) stderrBuffer).setCREndsLine(true);
    } else {
      stderrBuffer = new ByteArrayOutputStream();
    }
    ExecResult res = SimpleCommandLineProcessRunner.runCommandSecure(
            cli.toGeneralCommandLine(), command, null,
            new SimpleCommandLineProcessRunner.ProcessRunCallback() {
              @Override
              public void onProcessStarted(@NotNull Process ps) {
              }

              @Override
              public void onProcessFinished(@NotNull Process ps) {
                long duration = System.currentTimeMillis() - start;
                Loggers.VCS.debug("Command " + command + " took " + duration + "ms");
              }

              @Override
              public Integer getOutputIdleSecondsTimeout() {
                return executionTimeout;
              }

              @Override
              public Integer getMaxAcceptedOutputSize() {
                return readBytesLimit;
              }
            },
            stdoutBuffer, stderrBuffer);
    return new CommandResult(Loggers.VCS, command, res, privateData, settings, cli.getWorkingDirectory());
  }

  public static String removePrivateData(final String str, final Set<String> privateData) {
    String result = str;
    for (String data: privateData) {
      if (data == null || data.length() == 0) continue;
      result = result.replace(data, "******");
    }
    return result;
  }

  private static void logRunCommand(@NotNull MercurialCommandLine cmd, @NotNull String command, @NotNull CommandSettings settings) {
    String workingDir = cmd.getWorkingDirectory();
    String dir = workingDir != null ? workingDir : new File(".").getAbsolutePath();
    String message = "[" + dir + "] " + command;
    if ("debug".equals(settings.getLogLevel())) {
      Loggers.VCS.debug(message);
    } else {
      Loggers.VCS.info(message);
    }
  }

  @NotNull
  private static String logRunCommand(@NotNull final MercurialCommandLine cli,
                                      @NotNull final CommandSettings settings) {
    final String command = removePrivateData(cli.getCommandLineString(), settings.getPrivateData());
    logRunCommand(cli, command, settings);
    return command;
  }

  private static void logCommandOutput(@NotNull String command, @NotNull CommandResult result, @NotNull CommandSettings settings) {
    int limit = settings.getLogOutputLimit();
    if (limit == -1) {
      Loggers.VCS.debug("Command " + command + " output:\n" + result.getSecureStdout());
    } else {
      Loggers.VCS.debug("Command " + command + " output:\n" + StringUtil.truncateStringValueWithDotsAtEnd(result.getSecureStdout(), limit));
    }
  }
}