Support changing the fwup task via SSH

README.md

@@ -95,6 +95,44 @@ equivalent:
 cat $firmware | ssh -s $nerves_device fwup
 ```
 
+To run a specific task (like `complete` instead of the default `upgrade`), use
+SSH's `SendEnv` option to pass the `FWUP_TASK` environment variable:
+
+```shell
+FWUP_TASK=complete cat $firmware | ssh -o SendEnv=FWUP_TASK -s $nerves_device fwup
+```
+
+## Running different fwup tasks
+
+By default, uploads run the `upgrade` fwup task. For advanced use cases, you
+may want to run different tasks like:
+
+* `complete` - Completely re-image the device (useful for recovering or initial setup)
+* `ops` - Run operations firmware for partition validation, data erasure, or U-Boot updates
+
+No special server configuration is needed to support different tasks. The
+client specifies the task using SSH's `SendEnv` option.
+
+### Uploading with a specific task
+
+Using `mix upload`:
+
+```shell
+mix upload nerves.local --task complete
+```
+
+Using `upload.sh`:
+
+```shell
+./upload.sh --task complete nerves.local
+```
+
+Using raw ssh:
+
+```shell
+FWUP_TASK=complete cat my_firmware.fw | ssh -o SendEnv=FWUP_TASK -s nerves.local fwup
+```
+
 ## Configuration
 
 The default options should satisfy most use cases, but it's possible to alter
@@ -136,7 +174,9 @@ The following options are available:
   value closes the connection.
 * `:success_callback` - an MFArgs to call when a firmware update completes
   successfully. Defaults to `{Nerves.Runtime, :reboot, []}`.
-* `:task` - the task to run in the firmware update. Defaults to `"upgrade"`
+* `:task` - the task to run in the firmware update. Defaults to `"upgrade"`.
+  This can be overridden by clients using SSH's `SendEnv` option to pass the
+  `FWUP_TASK` environment variable.
 
 ## License
 

lib/mix/tasks/firmware.gen.script.ex

@@ -25,7 +25,11 @@ defmodule Mix.Tasks.Firmware.Gen.Script do
   # Upload new firmware to a device running ssh_subsystem_fwup
   #
   # Usage:
-  #   upload.sh [destination IP] [Path to .fw file]
+  #   upload.sh [options] [destination IP] [Path to .fw file]
+  #
+  # Options:
+  #   --task <task>    Specify the fwup task to run (e.g., "upgrade", "complete", "ops")
+  #                    Default: No task specified (uses server default, typically "upgrade")
   #
   # If unspecified, the destination is nerves.local and the .fw file is naively
   # guessed
@@ -52,21 +56,76 @@ defmodule Mix.Tasks.Firmware.Gen.Script do
 
   set -e
 
-  DESTINATION=$1
-  FILENAME="$2"
+  TASK=""
+  DESTINATION=""
+  FILENAME=""
 
-  help() {
+  show_help() {
+    echo
+    echo "Usage: upload.sh [options] [destination IP] [Path to .fw file]"
     echo
-    echo "upload.sh [destination IP] [Path to .fw file]"
+    echo "Options:"
+    echo "  --task <task>    Specify the fwup task to run (e.g., upgrade, complete, ops)"
+    echo "                   Default: No task specified (uses server default)"
+    echo "  --help           Show this help message"
     echo
     echo "Default destination IP is 'nerves.local'"
     echo "Default firmware bundle is the first .fw file in '_build/\\${MIX_TARGET}_\\${MIX_ENV}/nerves/images'"
     echo
-    echo "MIX_TARGET=$MIX_TARGET"
-    echo "MIX_ENV=$MIX_ENV"
-    exit 1
+    echo "Examples:"
+    echo "  ./upload.sh                                    # Upload to nerves.local"
+    echo "  ./upload.sh 192.168.1.100                      # Upload to specific IP"
+    echo "  ./upload.sh --task complete nerves.local       # Run complete task"
+    echo "  ./upload.sh --task ops 192.168.1.100 my.fw     # Run ops task with specific firmware"
+    echo
+    echo "Environment:"
+    echo "  MIX_TARGET=$MIX_TARGET"
+    echo "  MIX_ENV=$MIX_ENV"
   }
 
+  # Parse arguments
+  while [ $# -gt 0 ]; do
+    case "$1" in
+      --task)
+        # Check that $2 exists and doesn't start with a dash
+        case "$2" in
+          ""|-)
+            echo "Error: --task requires a value"
+            exit 1
+            ;;
+          -*)
+            echo "Error: --task requires a value"
+            exit 1
+            ;;
+          *)
+            TASK="$2"
+            shift 2
+            ;;
+        esac
+        ;;
+      --help|-h)
+        show_help
+        exit 0
+        ;;
+      -*)
+        echo "Error: Unknown option $1"
+        exit 1
+        ;;
+      *)
+        # Positional arguments
+        if [ -z "$DESTINATION" ]; then
+          DESTINATION="$1"
+        elif [ -z "$FILENAME" ]; then
+          FILENAME="$1"
+        else
+          echo "Error: Too many arguments"
+          exit 1
+        fi
+        shift
+        ;;
+    esac
+  done
+
   [ -n "$DESTINATION" ] || DESTINATION=nerves.local
   if [ -z "$FILENAME" ]; then
     [ -n "$MIX_TARGET" ] || MIX_TARGET=rpi0
@@ -93,10 +152,10 @@ defmodule Mix.Tasks.Firmware.Gen.Script do
     fi
 
     FILENAME=$(ls "$FIRMWARE_PATH/"*.fw 2> /dev/null | head -n 1)
-    [ -n "$FILENAME" ] || (echo "Error: error determining firmware bundle."; help)
+    [ -n "$FILENAME" ] || (echo "Error: error determining firmware bundle."; show_help; exit 1)
   fi
 
-  [ -f "$FILENAME" ] || (echo "Error: can't find '$FILENAME'"; help)
+  [ -f "$FILENAME" ] || (echo "Error: can't find '$FILENAME'"; show_help; exit 1)
 
   FIRMWARE_METADATA=$(fwup -m -i "$FILENAME" || echo "meta-product=Error reading metadata!")
   FIRMWARE_PRODUCT=$(echo "$FIRMWARE_METADATA" | grep -E "^meta-product=" -m 1 2>/dev/null | cut -d '=' -f 2- | tr -d '"')
@@ -108,10 +167,20 @@ defmodule Mix.Tasks.Firmware.Gen.Script do
   echo "Product: $FIRMWARE_PRODUCT $FIRMWARE_VERSION"
   echo "UUID: $FIRMWARE_UUID"
   echo "Platform: $FIRMWARE_PLATFORM"
+
+  # Set up task environment variable and SSH options if task is specified
+  if [ -n "$TASK" ]; then
+    export FWUP_TASK="$TASK"
+    SEND_ENV_OPT="-o SendEnv=FWUP_TASK"
+    echo "Task: $TASK"
+  else
+    SEND_ENV_OPT=""
+  fi
+
   echo
   echo "Uploading to $DESTINATION..."
 
-  cat "$FILENAME" | ssh -s $SSH_OPTIONS $DESTINATION fwup
+  cat "$FILENAME" | ssh -s $SEND_ENV_OPT $SSH_OPTIONS $DESTINATION fwup
   """
 
   @spec run(keyword()) :: :ok

lib/mix/tasks/upload.ex

@@ -26,6 +26,9 @@ defmodule Mix.Tasks.Upload do
 
    * `--firmware` - The path to a fw file
    * `--port` - An alternative TCP port to use for the upload (defaults to 22)
+   * `--task` - The fwup task to run on the device (defaults to "upgrade").
+     Use this to run alternative tasks like "complete" or custom tasks defined
+     in your firmware.
 
   ## Examples
 
@@ -37,11 +40,16 @@ defmodule Mix.Tasks.Upload do
 
       mix upload 192.168.1.120 --firmware _build/rpi0_prod/nerves/images/app.fw
 
+  Run the complete task instead of upgrade:
+
+      MIX_TARGET=rpi0 mix upload nerves.local --task complete
+
   """
 
   @switches [
     firmware: :string,
-    port: :integer
+    port: :integer,
+    task: :string
   ]
 
   @doc false
@@ -66,18 +74,22 @@ defmodule Mix.Tasks.Upload do
     port = opts[:port] || 22
     validate_port!(port)
 
+    task = opts[:task]
+    task_env = if task, do: "FWUP_TASK=#{task} ", else: ""
+    send_env_opt = if task, do: "-o SendEnv=FWUP_TASK ", else: ""
+
     firmware_path = firmware(opts)
 
     Mix.shell().info("""
     Path: #{firmware_path}
     #{maybe_print_firmware_uuid(firmware_path)}
-    Uploading to #{ip}:#{port}...
+    #{maybe_print_task(task)}Uploading to #{ip}:#{port}...
     """)
 
     # LD_LIBRARY_PATH is unset to avoid errors with host ssl (see commit 9b1df471)
     {_, status} =
       InteractiveCmd.shell(
-        "cat #{shell_quote(firmware_path)} | ssh -p #{port} -s -- #{shell_quote(ip)} fwup",
+        "#{task_env}cat #{shell_quote(firmware_path)} | ssh #{send_env_opt}-p #{port} -s -- #{shell_quote(ip)} fwup",
         env: [{"LD_LIBRARY_PATH", false}]
       )
 
@@ -178,5 +190,8 @@ defmodule Mix.Tasks.Upload do
     _, _ -> ""
   end
 
+  defp maybe_print_task(nil), do: ""
+  defp maybe_print_task(task), do: "Task: #{task}\n"
+
   defp shell_quote(str), do: "'" <> String.replace(str, "'", "'\"'\"'") <> "'"
 end

lib/ssh_subsystem_fwup.ex

@@ -56,7 +56,9 @@ defmodule SSHSubsystemFwup do
     return value closes the connection.
   * `:success_callback` - an MFArgs to call when a firmware update completes
     successfully. Defaults to `{Nerves.Runtime, :reboot, []}`.
-  * `:task` - the task to run in the firmware update. Defaults to `"upgrade"`
+  * `:task` - the task to run in the firmware update. Defaults to `"upgrade"`.
+    This can be overridden by clients using SSH's `SendEnv` option to pass the
+    `FWUP_TASK` environment variable.
   """
   @type options :: [
           devpath: Path.t(),
@@ -74,6 +76,20 @@ defmodule SSHSubsystemFwup do
 
   @doc """
   Helper for creating the SSH subsystem spec
+
+  Clients can override the task by using SSH's `SendEnv` option to pass the
+  `FWUP_TASK` environment variable. For example:
+
+  ```shell
+  FWUP_TASK=complete cat firmware.fw | ssh -o SendEnv=FWUP_TASK -s device fwup
+  ```
+
+  Or using the provided upload tools:
+
+  ```shell
+  mix upload nerves.local --task complete
+  ./upload.sh --task complete nerves.local
+  ```
   """
   @spec subsystem_spec(options()) :: :ssh.subsystem_spec()
   def subsystem_spec(options \\ []) do
@@ -88,7 +104,7 @@ defmodule SSHSubsystemFwup do
       |> Keyword.merge(Application.get_all_env(:ssh_subsystem_fwup))
       |> Keyword.merge(options)
 
-    {:ok, %{state: :running_fwup, id: nil, cm: nil, fwup: nil, options: combined_options}}
+    {:ok, %{state: :running_fwup, id: nil, cm: nil, fwup: nil, options: combined_options, env: %{}}}
   end
 
   defp default_options() do
@@ -105,11 +121,19 @@ defmodule SSHSubsystemFwup do
 
   @impl :ssh_client_channel
   def handle_msg({:ssh_channel_up, channel_id, cm}, state) do
-    with {:ok, options} <- precheck(state.options[:precheck_callback], state.options),
+    # Check if FWUP_TASK was set via SendEnv and override the task option
+    options =
+      case Map.get(state.env, "FWUP_TASK") do
+        nil -> state.options
+        task when is_binary(task) -> Keyword.put(state.options, :task, validate_task(task))
+        _ -> state.options
+      end
+
+    with {:ok, options} <- precheck(options[:precheck_callback], options),
          :ok <- check_devpath(options[:devpath]) do
-      Logger.debug("ssh_subsystem_fwup: starting fwup")
+      Logger.debug("ssh_subsystem_fwup: starting fwup with task #{options[:task]}")
       fwup = FwupPort.open_port(options)
-      {:ok, %{state | id: channel_id, cm: cm, fwup: fwup}}
+      {:ok, %{state | id: channel_id, cm: cm, fwup: fwup, options: options}}
     else
       {:error, reason} ->
         _ = :ssh_connection.send(cm, channel_id, "Error: #{reason}")
@@ -160,6 +184,28 @@ defmodule SSHSubsystemFwup do
     {:ok, state}
   end
 
+  def handle_ssh_msg({:ssh_cm, cm, {:env, channel_id, want_reply, var, value}}, state) do
+    # Convert charlists to strings if needed
+    var_str = if is_list(var), do: to_string(var), else: var
+    value_str = if is_list(value), do: to_string(value), else: value
+
+    # Only accept FWUP_TASK environment variable for security
+    new_env =
+      if var_str == "FWUP_TASK" do
+        Map.put(state.env, var_str, value_str)
+      else
+        Logger.debug("ssh_subsystem_fwup: ignoring environment variable #{var_str}")
+        state.env
+      end
+
+    # Reply if requested
+    if want_reply do
+      :ssh_connection.reply_request(cm, want_reply, :success, channel_id)
+    end
+
+    {:ok, %{state | env: new_env}}
+  end
+
   def handle_ssh_msg({:ssh_cm, _cm, {:eof, _channel_id}}, state) do
     {:ok, state}
   end
@@ -212,6 +258,17 @@ defmodule SSHSubsystemFwup do
     {:ok, state}
   end
 
+  # Validate task name to only contain safe characters (alphanumeric, underscore, hyphen)
+  # to prevent potential command injection
+  defp validate_task(task) when is_binary(task) do
+    if Regex.match?(~r/^[a-zA-Z0-9_-]+$/, task) do
+      task
+    else
+      Logger.warning("ssh_subsystem_fwup: invalid task name #{inspect(task)}, using default")
+      "upgrade"
+    end
+  end
+
   defp check_devpath(devpath) do
     if is_binary(devpath) and File.exists?(devpath) do
       :ok