From 00ad27bc34598558d67aa14c59b28855d4d415e0 Mon Sep 17 00:00:00 2001
From: c4710n <c4710n@users.noreply.github.com>
Date: Sat, 8 Jan 2022 22:17:51 +0800
Subject: [PATCH] Add support for targetPort

---
 data/eval-machines.nix |  2 +-
 data/options.nix       |  8 ++++++++
 healthchecks/types.go  |  1 +
 nix/nix.go             |  8 ++++++++
 ssh/ssh.go             | 12 ++++++++++++
 5 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/data/eval-machines.nix b/data/eval-machines.nix
index a4bcda6..00b75cd 100644
--- a/data/eval-machines.nix
+++ b/data/eval-machines.nix
@@ -78,7 +78,7 @@ in rec {
 
     machines =
       flip mapAttrs nodes (n: v': let v = scrubOptionValue v'; in
-        { inherit (v.config.deployment) targetHost targetUser secrets healthChecks buildOnly substituteOnDestination tags;
+        { inherit (v.config.deployment) targetHost targetPort targetUser secrets healthChecks buildOnly substituteOnDestination tags;
           name = n;
           nixosRelease = v.config.system.nixos.release or (removeSuffix v.config.system.nixos.version.suffix v.config.system.nixos.version);
           nixConfig = mapAttrs
diff --git a/data/options.nix b/data/options.nix
index 4ffdc76..743e1d3 100644
--- a/data/options.nix
+++ b/data/options.nix
@@ -176,6 +176,14 @@ in
       '';
     };
 
+    targetPort = mkOption {
+      type = nullOr int;
+      default = null;
+      description = ''
+        The port number of remote host used for deployment.
+      '';
+    };
+
     targetUser = mkOption {
       type = str;
       default = "";
diff --git a/healthchecks/types.go b/healthchecks/types.go
index 7058ed2..c37bbdf 100644
--- a/healthchecks/types.go
+++ b/healthchecks/types.go
@@ -15,6 +15,7 @@ import (
 type Host interface {
 	GetName() string
 	GetTargetHost() string
+	GetTargetPort() int
 	GetTargetUser() string
 	GetHealthChecks() HealthChecks
 }
diff --git a/nix/nix.go b/nix/nix.go
index 9c67da8..33382ce 100644
--- a/nix/nix.go
+++ b/nix/nix.go
@@ -25,6 +25,7 @@ type Host struct {
 	Name                    string
 	NixosRelease            string
 	TargetHost              string
+	TargetPort              int
 	TargetUser              string
 	Secrets                 map[string]secrets.Secret
 	BuildOnly               bool
@@ -66,6 +67,10 @@ func (host *Host) GetTargetHost() string {
 	return host.TargetHost
 }
 
+func (host *Host) GetTargetPort() int {
+	return host.TargetPort
+}
+
 func (host *Host) GetTargetUser() string {
 	return host.TargetUser
 }
@@ -385,6 +390,9 @@ func Push(ctx *ssh.SSHContext, host Host, paths ...string) (err error) {
 	if ctx.SkipHostKeyCheck {
 		sshOpts = append(sshOpts, fmt.Sprintf("%s", "-o StrictHostkeyChecking=No -o UserKnownHostsFile=/dev/null"))
 	}
+	if host.TargetPort != 0 {
+		sshOpts = append(sshOpts, fmt.Sprintf("-p %d", host.TargetPort))
+	}
 	if ctx.ConfigFile != "" {
 		sshOpts = append(sshOpts, fmt.Sprintf("-F %s", ctx.ConfigFile))
 	}
diff --git a/ssh/ssh.go b/ssh/ssh.go
index e4920fe..caeac36 100644
--- a/ssh/ssh.go
+++ b/ssh/ssh.go
@@ -34,6 +34,7 @@ type Context interface {
 type Host interface {
 	GetName() string
 	GetTargetHost() string
+	GetTargetPort() int
 	GetTargetUser() string
 }
 
@@ -94,6 +95,17 @@ func (ctx *SSHContext) sshArgs(host Host, transfer *FileTransfer) (cmd string, a
 		args = append(args, "-F", ctx.ConfigFile)
 	}
 	var hostAndDestination = host.GetTargetHost()
+	if host.GetTargetPort() != 0 {
+		var optionName string
+		if transfer != nil {
+			optionName = "-P"
+		} else {
+			optionName = "-p"
+		}
+
+		args = append(args, optionName, fmt.Sprintf("%d", host.GetTargetPort()))
+	}
+
 	if transfer != nil {
 		args = append(args, transfer.Source)
 		hostAndDestination += ":" + transfer.Destination