#!/bin/bash
# Copyright 2017 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

#### CATEGORY=Other
### serve from a remote workstation

## usage: fx serve-remote [--no-serve] [--tunnel-ports=NNNN,..] HOSTNAME [REMOTE-PATH]
##
## HOSTNAME     the hostname of the workstation you want to serve from
## REMOTE-PATH  defaults to ~/fuchsia. The path on the to FUCHSIA_DIR on the workstation.
##
##  --no-serve                    only tunnel, do not start a package server
##  --no-check-ssh-keys  do not verify that the default SSH
##                                credentials are the same before serving.
##  --[no-]persist                enable or disable persistence of repository metadata.
##                                Disabled by default.
##  --tunnel-ports=NNN1,NNN2,NNN3 comma-separated list of additional ports to
##                                tunnel. This is used for e2e tests running on
##                                remote host that needs to reach the local device.
##
## HOST and DIR are persisted in the file //.fx-remote-config and are reused as
## defaults in future invocations of any 'fx *-remote' tools.

source "$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"/lib/vars.sh || exit $?
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"/lib/fx-remote.sh || exit $?
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"/lib/verify-default-keys.sh || exit $?
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"/lib/updates.sh || exit $?
fx-config-read

fx-standard-switches "$@"
set -- "${FX_ARGV[@]}"

serve=true
check_ssh_keys=true
verbose=false
host=""
dir=""
serve_persist_arg="--no-persist"
package_server_port="8083"
has_tunnel_ports=false

while [[ $# -ne 0 ]]; do
  case "$1" in
  --help|-h)
      fx-command-help
      exit 0
      ;;
  --no-serve)
    serve=false
    ;;
  --no-check-ssh-keys)
    check_ssh_keys=false
    ;;
  --no-persist)
    ;;
  --persist)
    serve_persist_arg="--persist"
    ;;
  -v)
    verbose=true
    ;;
  --repo-port)
    if [[ $# -lt 2 ]]; then
      fx-error Invalid syntax
      fx-command-help
      exit 1
    fi
    package_server_port=$2
    shift
    ;;
  --tunnel-ports)
    if [[ $# -lt 2 ]]; then
      fx-error Invalid syntax
      fx-command-help
      exit 1
    fi
    has_tunnel_ports=true
    # Split comma-separated list of ports to an array.
    tunnel_ports=(${2//,/ })
    shift
    ;;
  -*)
    fx-error "Unknown flag: $1"
    fx-command-help
    exit 1
    ;;
  *)
    if [[ -z "${host}" ]]; then
      host="$1"
    elif [[ -z "${dir}" ]]; then
      dir="$1"
    else
      fx-error "unexpected argument: '$1'"
      exit 1
    fi
    ;;
  esac
  shift
done

if cached=( $(load_remote_info "${host}") ); then
  host="${cached[0]}"
  dir="${cached[1]}"
fi

if [[ -z "${host}" ]]; then
  fx-error "HOSTNAME must be specified"
  fx-command-help
  exit 1
fi

if "${serve}"; then
  if [[ -z "${dir}" ]]; then
    if ssh "${host}" ls \~/fuchsia/.jiri_root/bin/fx > /dev/null; then
      dir="~/fuchsia"
    else
      fx-error "failed to find ~/fuchsia on $host, please specify REMOTE-DIR"
      fx-command-help
      exit 1
    fi
  fi
fi

save_remote_info "${host}" "${dir}"

fx-export-device-address
if [[ $? -ne 0 || -z "${FX_DEVICE_ADDR}" ]]; then
  fx-error "unable to discover device. Is the target up?"
  exit 1
fi

if [[ -z "${FX_SSH_PORT}" ]]; then
  FX_SSH_PORT=22
fi

echo "Using remote ${host}:${dir}"
echo "Using target device ${FX_DEVICE_NAME} (${FX_SSH_ADDR}:${FX_SSH_PORT})"


# Use a dedicated ControlPath so script can manage a connection seperately from the user's. We
# intentionally do not use %h/%p in the control path because there can only be one forwarding
# session at a time (due to the local forward of the package server port).
ssh_base_args=(
  -S "${HOME}/.ssh/control-fuchsia-fx-remote"
  -o "ControlMaster=auto"
)

ssh_exit() {
  # Failure to end existing multiplexed SSH connections is acceptable.
  ssh "${ssh_base_args[@]}" "${host}" -O exit > /dev/null 2>&1 || true
  wait # for ssh to exit
}

# If there is already control master then exit it. We can't be sure its to the right host and it
# also could be stale.
ssh_exit

# When we exit the script, close the background ssh connection.
trap_exit() {
  ssh_exit
  exit
}
trap trap_exit EXIT

ssh_tunnel_args=(
  -6 # We want ipv6 binds for the port forwards
  -L "\*:${package_server_port}:localhost:${package_server_port}"       # fx serve
  -R "8022:${FX_SSH_ADDR}:${FX_SSH_PORT}"                               # fx shell
  -o ExitOnForwardFailure=yes
  # Match google default server timeout so in spotty network situations the client doesn't timeout
  # before server (and leave the server process still holding on to tunneling port).
  -o ServerAliveInterval=30
  -o ServerAliveCountMax=20
)

# This is a rudimentary assumption that the device is reachable at FX_SSH_ADDR
# iff it is listening on the default SSH port, which is often true, but in no
# way guaranteed to be true. A better approach here would be to adjust to later
# perform these forwards inside the device fowrarded link rather than at the
# host link level.
if [[ "${FX_SSH_PORT}" == 22 ]]; then
  ssh_tunnel_args+=(
    -R "2345:${FX_SSH_ADDR}:2345"             # fx debug
    -R "8007:${FX_SSH_ADDR}:8007"             # Google-specific
    -R "8443:${FX_SSH_ADDR}:8443"             # Google-specific
    -R "9080:${FX_SSH_ADDR}:80"               # SL4F_HTTP_PORT
  )
else
  echo >&2 "Note: tunnelled targets will not receive forwards for SL4F or debug"
fi

# Add additional ports to tunnel if specified.
if "${has_tunnel_ports}"; then
  for port in "${tunnel_ports[@]}"; do
    ssh_tunnel_args+=(-R "${port}:${FX_SSH_ADDR}:${port}")
  done
fi

# Start tunneling session in background. It's started seperately from the command invocations below
# to allow the script to be consistent on how it is exited for both serve and non-serve cases. It
# also allows script to explicitly close the control session (to better avoid stale sshd sessions).

# Verify that keys match.
if "${check_ssh_keys}"; then
  verify_default_keys "${FUCHSIA_DIR}" "${host}" "${dir}" "${ssh_base_args[@]}" || exit $?
fi

# XXX: had to stop using -Nf because of b/160269794.
ssh "${ssh_base_args[@]}" "${ssh_tunnel_args[@]}" "${host}" -nT sleep infinity &
# Attempt to assert that the backgrounded ssh is alive and kicking, emulating -f as best we can.
ssh_pid=$!
# If there's a 2fa prompt, we may need a "human time" number of tries, which is why this is high.
tries=30
until ssh -q -O check ${ssh_base_args[@]} "${host}"; do
  if ! kill -0 ${ssh_pid}; then
    fx-error "SSH tunnel terminated prematurely"
    exit 1
  fi
  if ! ((tries--)); then
    fx-error "SSH tunnel appears not to have succeeded"
    kill -TERM $ssh_pid
    exit 1
  fi
  sleep 1
done

if "${serve}"; then
  remote_server_mode=$(ssh "${host}" "${ssh_base_args[@]}" "cd ${dir} && ./.jiri_root/bin/ffx config get repository.server.mode") || err=$?
  if [[ "${err}" -ne 0 ]]; then
    # Warn if we can't look up the remote server mode, which might mean the
    # remote checkout is out of date.
    fx-warn "It appears the remote host does not have a value defined for"
    fx-warn "'ffx config get repository.server.mode', which suggests it might be"
    fx-warn "out of date."
    fx-warn ""
    fx-warn "Run '$ jiri update && fx build ffx' on both the local and remote hosts, then"
    fx-warn "confirm that '$ ffx config get repository.server.mode' returns a string on"
    fx-warn "both the local and remote hosts."

    # Clear the error so later checks won't fail.
    err=0
  fi

  if [[ "${remote_server_mode}" == \"pm\" ]]; then
    # If the user requested serving, then we'll check to see if there's a
    # remote server already running and kill it, this prevents most cases where
    # signal propagation seems to sometimes not make it to "pm".
    # TODO(drees) This can be combined with the serve-updates call later to reduce ssh calls.
    if ssh "${host}" "${ssh_base_args[@]}" "ss -ln | grep :${package_server_port}" > /dev/null; then
      ssh "${ssh_base_args[@]}" "${host}" 'pkill -x -u $USER pm' || true
    fi
  else
    # FIXME(http://fxbug.dev/82788): Currently ffx will log, but otherwise
    # ignore, when both a local ffx and a remote ffx are both trying to serve
    # repositories over the same tunnel port. The first setup will win, even if
    # that's not what the user expects.
    #
    # To work around this, we will error out if both the local and remote ffx
    # repository servers are trying to use the same port.
    local_addr=$(fx-command-run ffx config get repository.server.listen) || err=$?
    if [[ "${err}" -ne 0 ]]; then
      fx-error "Unable to get the local configured repository server address."
      exit "${err}"
    fi

    remote_addr=$(ssh "${host}" "${ssh_base_args[@]}" "cd ${dir} && ./.jiri_root/bin/ffx config get repository.server.listen") || err=$?
    if [[ "${err}" -ne 0 ]]; then
      fx-error "Unable to get the remote configured repository server address."
      exit "${err}"
    fi

    # If the local server is running, error out if we're using the same port as the remote server.
    if [[ "${local_addr}" != "null" && "${local_addr}" = "${remote_addr}" ]]; then
      fx-error "The local and remote ffx repository servers cannot both be $local_addr."
      fx-error "Either change the ffx repository port with:"
      fx-error ""
      fx-error "$ ffx config set repository.server.listen \"[::]:\$SOME_OTHER_PORT\""
      fx-error "$ ffx doctor --restart-daemon"
      fx-error ""
      fx-error "Or disable the ffx repository server with:"
      fx-error ""
      fx-error "$ ffx config set repository.server.mode pm"
      fx-error "$ ffx doctor --restart-daemon"
      fx-error ""
      fx-error "See http://fxbug.dev/82788 for more details."
      exit 1
    fi
  fi
fi

if "${serve}"; then
  # Ctrl-C will exit the ssh remote command and this ssh session. Then script exit will trigger
  # `trap_exit` to close the ssh connection.
  echo -e "Press Ctrl-C to stop remote serving and tunneling.\n"

  # Set the experimental environment variables to ensure they're passed to `fx
  # set` and `fx serve-updates`.
  experimental_vars=$(get_env_vars_non_default_features)

  if "${verbose}"; then
    serve_verbose_arg=" -v"
  else
    serve_verbose_arg=""
  fi
  ssh_serve_args=(
    "-tt" # explicitly force a pty, for HUP'ing on the remote
    "cd ${dir} && FX_REMOTE_INVOCATION=1 ${experimental_vars} ./.jiri_root/bin/fx set-device '${_FX_REMOTE_WORKFLOW_DEVICE_ADDR}' && FX_REMOTE_INVOCATION=1 ${experimental_vars} ./.jiri_root/bin/fx serve-updates ${serve_verbose_arg} ${serve_persist_arg}"
  )
  ssh "${host}" "${ssh_base_args[@]}" "${ssh_serve_args[@]}"
else
  echo "Press Ctrl-C to stop tunneling."
  # Wait for user Ctrl-C. Then script exit will trigger trap_exit to close the
  # ssh connection. Use a command on the remote side over ssh with a -tt
  # parameter to create a pseudo tty on the remote host. That way, the Ctrl-C
  # hup will go to the remote host which will close the remote ssh connections.
  ssh "${host}" "${ssh_base_args[@]}" -tt sleep infinity
fi